main
py 171 lines 5.18 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 # Copyright 2025 The vLLM team.
5 # Copyright 2025 NVIDIA CORPORATION and the HuggingFace Inc. team. All rights
6 # reserved.
7 #
8 # Licensed under the Apache License, Version 2.0 (the "License");
9 # you may not use this file except in compliance with the License.
10 # You may obtain a copy of the License at
11 #
12 # http://www.apache.org/licenses/LICENSE-2.0
13 #
14 # Unless required by applicable law or agreed to in writing, software
15 # distributed under the License is distributed on an "AS IS" BASIS,
16 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
17 # See the License for the specific language governing permissions and
18 # limitations under the License.
19
20 import json
21 import os
22
23 import pytest
24
25 from tests.models.registry import HF_EXAMPLE_MODELS
26 from vllm import LLM, SamplingParams
27
28 MODEL_NAME = "nvidia/audio-flamingo-3-hf"
29 SINGLE_CONVERSATION = [
30 {
31 "role": "user",
32 "content": [
33 {
34 "type": "text",
35 "text": "What is surprising about the relationship between "
36 "the barking and the music?",
37 },
38 {
39 "type": "audio_url",
40 "audio_url": {
41 "url": "https://huggingface.co/datasets/nvidia/AudioSkills/"
42 "resolve/main/assets/"
43 "dogs_barking_in_sync_with_the_music.wav",
44 },
45 },
46 ],
47 }
48 ]
49 BATCHED_CONVERSATIONS = [
50 SINGLE_CONVERSATION,
51 [
52 {
53 "role": "user",
54 "content": [
55 {
56 "type": "text",
57 "text": "Why is the philosopher's name mentioned in the "
58 "lyrics? (A) To express a sense of nostalgia "
59 "(B) To indicate that language cannot express clearly, "
60 "satirizing the inversion of black and white in the world "
61 "(C) To add depth and complexity to the lyrics "
62 "(D) To showcase the wisdom and influence of the "
63 "philosopher",
64 },
65 {
66 "type": "audio_url",
67 "audio_url": {
68 "url": "https://huggingface.co/datasets/nvidia/"
69 "AudioSkills/resolve/main/assets/"
70 "Ch6Ae9DT6Ko_00-04-03_00-04-31.wav",
71 },
72 },
73 ],
74 }
75 ],
76 ]
77
78
79 def get_fixture_path(filename):
80 return os.path.join(
81 os.path.dirname(__file__), "../../fixtures/audioflamingo3", filename
82 )
83
84
85 def assert_output_matches(output, expected_text, expected_token_ids):
86 generated = output.outputs[0]
87 assert generated.text.strip() == expected_text
88 actual_token_ids = list(generated.token_ids)
89 assert (
90 actual_token_ids == expected_token_ids
91 or actual_token_ids == expected_token_ids[:-1]
92 or actual_token_ids[:-1] == expected_token_ids
93 )
94
95
96 @pytest.fixture(scope="module")
97 def llm():
98 model_info = HF_EXAMPLE_MODELS.get_hf_info("AudioFlamingo3ForConditionalGeneration")
99 model_info.check_transformers_version(on_fail="skip")
100
101 try:
102 return LLM(
103 model=MODEL_NAME,
104 dtype="bfloat16",
105 enforce_eager=True,
106 limit_mm_per_prompt={"audio": 1},
107 )
108 except Exception as e:
109 pytest.skip(f"Failed to load model {MODEL_NAME}: {e}")
110
111
112 def test_single_generation(llm):
113 fixture_path = get_fixture_path("expected_results_single.json")
114 if not os.path.exists(fixture_path):
115 pytest.skip(f"Fixture not found: {fixture_path}")
116
117 with open(fixture_path) as f:
118 expected = json.load(f)
119
120 sampling_params = SamplingParams(temperature=0.0, max_tokens=128)
121
122 outputs = llm.chat(
123 messages=SINGLE_CONVERSATION,
124 sampling_params=sampling_params,
125 )
126 assert_output_matches(
127 outputs[0],
128 expected["transcriptions"][0],
129 expected["token_ids"][0],
130 )
131
132
133 def test_batched_generation(llm):
134 fixture_path = get_fixture_path("expected_results_batched.json")
135 if not os.path.exists(fixture_path):
136 pytest.skip(f"Fixture not found: {fixture_path}")
137
138 with open(fixture_path) as f:
139 expected = json.load(f)
140
141 sampling_params = SamplingParams(temperature=0.0, max_tokens=128)
142
143 outputs = llm.chat(
144 messages=BATCHED_CONVERSATIONS,
145 sampling_params=sampling_params,
146 )
147
148 for i, output in enumerate(outputs):
149 assert_output_matches(
150 output,
151 expected["transcriptions"][i],
152 expected["token_ids"][i],
153 )
154
155
156 def test_single_and_batched_generation_match(llm):
157 sampling_params = SamplingParams(temperature=0.0, max_tokens=128)
158
159 single_output = llm.chat(
160 messages=SINGLE_CONVERSATION,
161 sampling_params=sampling_params,
162 )[0]
163 batched_output = llm.chat(
164 messages=BATCHED_CONVERSATIONS,
165 sampling_params=sampling_params,
166 )[0]
167
168 assert single_output.outputs[0].text == batched_output.outputs[0].text
169 assert list(single_output.outputs[0].token_ids) == list(
170 batched_output.outputs[0].token_ids
171 )