test_vision_openai_server_common.py 15.8 KB
Newer Older
1
2
import base64
import io
Ying Sheng's avatar
Ying Sheng committed
3
import json
4
import os
Ying Sheng's avatar
Ying Sheng committed
5
import unittest
6
from concurrent.futures import ThreadPoolExecutor
Ying Sheng's avatar
Ying Sheng committed
7

8
import numpy as np
Ying Sheng's avatar
Ying Sheng committed
9
import openai
10
11
import requests
from PIL import Image
Ying Sheng's avatar
Ying Sheng committed
12

13
from sglang.srt.utils import kill_process_tree
14
15
16
from sglang.test.test_utils import (
    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
    DEFAULT_URL_FOR_TEST,
17
    CustomTestCase,
18
19
    popen_launch_server,
)
Ying Sheng's avatar
Ying Sheng committed
20

21
22
23
24
25
26
27
28
29
30
31
# image
IMAGE_MAN_IRONING_URL = "https://raw.githubusercontent.com/sgl-project/sgl-test-files/refs/heads/main/images/man_ironing_on_back_of_suv.png"
IMAGE_SGL_LOGO_URL = "https://raw.githubusercontent.com/sgl-project/sgl-test-files/refs/heads/main/images/sgl_logo.png"

# video
VIDEO_JOBS_URL = "https://raw.githubusercontent.com/sgl-project/sgl-test-files/refs/heads/main/videos/jobs_presenting_ipod.mp4"

# audio
AUDIO_TRUMP_SPEECH_URL = "https://raw.githubusercontent.com/sgl-project/sgl-test-files/refs/heads/main/audios/Trump_WEF_2018_10s.mp3"
AUDIO_BIRD_SONG_URL = "https://raw.githubusercontent.com/sgl-project/sgl-test-files/refs/heads/main/audios/bird_song.mp3"

Ying Sheng's avatar
Ying Sheng committed
32

33
class TestOpenAIVisionServer(CustomTestCase):
Ying Sheng's avatar
Ying Sheng committed
34
35
    @classmethod
    def setUpClass(cls):
36
        cls.model = "lmms-lab/llava-onevision-qwen2-0.5b-ov"
37
        cls.base_url = DEFAULT_URL_FOR_TEST
Ying Sheng's avatar
Ying Sheng committed
38
39
40
41
        cls.api_key = "sk-123456"
        cls.process = popen_launch_server(
            cls.model,
            cls.base_url,
42
            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
Ying Sheng's avatar
Ying Sheng committed
43
44
45
46
47
48
            api_key=cls.api_key,
        )
        cls.base_url += "/v1"

    @classmethod
    def tearDownClass(cls):
49
        kill_process_tree(cls.process.pid)
Ying Sheng's avatar
Ying Sheng committed
50

51
    def test_single_image_chat_completion(self):
Ying Sheng's avatar
Ying Sheng committed
52
53
54
55
56
57
58
59
60
61
        client = openai.Client(api_key=self.api_key, base_url=self.base_url)

        response = client.chat.completions.create(
            model="default",
            messages=[
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "image_url",
62
                            "image_url": {"url": IMAGE_MAN_IRONING_URL},
Ying Sheng's avatar
Ying Sheng committed
63
                        },
Ying Sheng's avatar
Ying Sheng committed
64
65
66
67
                        {
                            "type": "text",
                            "text": "Describe this image in a very short sentence.",
                        },
Ying Sheng's avatar
Ying Sheng committed
68
69
70
71
72
73
74
                    ],
                },
            ],
            temperature=0,
        )

        assert response.choices[0].message.role == "assistant"
Ying Sheng's avatar
Ying Sheng committed
75
76
        text = response.choices[0].message.content
        assert isinstance(text, str)
77
        # `driver` is for gemma-3-it
78
79
80
81
82
83
84
85
86
87
        assert (
            "man" in text or "person" or "driver" in text
        ), f"text: {text}, should contain man, person or driver"
        assert (
            "cab" in text
            or "taxi" in text
            or "SUV" in text
            or "vehicle" in text
            or "car" in text
        ), f"text: {text}, should contain cab, taxi, SUV, vehicle or car"
Mick's avatar
Mick committed
88
        # MiniCPMO fails to recognize `iron`, but `hanging`
89
90
91
        assert (
            "iron" in text or "hang" in text or "cloth" in text or "holding" in text
        ), f"text: {text}, should contain iron, hang, cloth or holding"
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
        assert response.id
        assert response.created
        assert response.usage.prompt_tokens > 0
        assert response.usage.completion_tokens > 0
        assert response.usage.total_tokens > 0

    def test_multi_turn_chat_completion(self):
        client = openai.Client(api_key=self.api_key, base_url=self.base_url)

        response = client.chat.completions.create(
            model="default",
            messages=[
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "image_url",
109
                            "image_url": {"url": IMAGE_MAN_IRONING_URL},
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
                        },
                        {
                            "type": "text",
                            "text": "Describe this image in a very short sentence.",
                        },
                    ],
                },
                {
                    "role": "assistant",
                    "content": [
                        {
                            "type": "text",
                            "text": "There is a man at the back of a yellow cab ironing his clothes.",
                        }
                    ],
                },
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": "Repeat your previous answer."}
                    ],
                },
            ],
            temperature=0,
        )

        assert response.choices[0].message.role == "assistant"
        text = response.choices[0].message.content
        assert isinstance(text, str)
139
140
141
        assert (
            "man" in text or "cab" in text
        ), f"text: {text}, should contain man or cab"
Ying Sheng's avatar
Ying Sheng committed
142
143
144
145
        assert response.id
        assert response.created
        assert response.usage.prompt_tokens > 0
        assert response.usage.completion_tokens > 0
146
147
        assert response.usage.total_tokens > 0

148
    def test_multi_images_chat_completion(self):
149
150
151
152
153
154
155
156
157
158
        client = openai.Client(api_key=self.api_key, base_url=self.base_url)

        response = client.chat.completions.create(
            model="default",
            messages=[
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "image_url",
Mick's avatar
Mick committed
159
                            "image_url": {"url": IMAGE_MAN_IRONING_URL},
160
                            "modalities": "multi-images",
161
162
163
                        },
                        {
                            "type": "image_url",
164
                            "image_url": {"url": IMAGE_SGL_LOGO_URL},
165
                            "modalities": "multi-images",
166
167
168
                        },
                        {
                            "type": "text",
169
170
                            "text": "I have two very different images. They are not related at all. "
                            "Please describe the first image in one sentence, and then describe the second image in another sentence.",
171
172
173
174
175
176
177
178
179
180
                        },
                    ],
                },
            ],
            temperature=0,
        )

        assert response.choices[0].message.role == "assistant"
        text = response.choices[0].message.content
        assert isinstance(text, str)
Mick's avatar
Mick committed
181
182
183
        print("-" * 30)
        print(f"Multi images response:\n{text}")
        print("-" * 30)
184
185
186
187
188
189
        assert (
            "man" in text or "cab" in text or "SUV" in text or "taxi" in text
        ), f"text: {text}, should contain man, cab, SUV or taxi"
        assert (
            "logo" in text or '"S"' in text or "SG" in text
        ), f"text: {text}, should contain logo, S or SG"
190
191
192
193
        assert response.id
        assert response.created
        assert response.usage.prompt_tokens > 0
        assert response.usage.completion_tokens > 0
Ying Sheng's avatar
Ying Sheng committed
194
195
        assert response.usage.total_tokens > 0

196
    def prepare_video_messages(self, video_path):
197
198
        # the memory consumed by the Vision Attention varies a lot, e.g. blocked qkv vs full-sequence sdpa
        # the size of the video embeds differs from the `modality` argument when preprocessed
199
200
201
202
203
204
205

        # We import decord here to avoid a strange Segmentation fault (core dumped) issue.
        # The following import order will cause Segmentation fault.
        # import decord
        # from transformers import AutoTokenizer
        from decord import VideoReader, cpu

206
        max_frames_num = 20
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
        vr = VideoReader(video_path, ctx=cpu(0))
        total_frame_num = len(vr)
        uniform_sampled_frames = np.linspace(
            0, total_frame_num - 1, max_frames_num, dtype=int
        )
        frame_idx = uniform_sampled_frames.tolist()
        frames = vr.get_batch(frame_idx).asnumpy()

        base64_frames = []
        for frame in frames:
            pil_img = Image.fromarray(frame)
            buff = io.BytesIO()
            pil_img.save(buff, format="JPEG")
            base64_str = base64.b64encode(buff.getvalue()).decode("utf-8")
            base64_frames.append(base64_str)

        messages = [{"role": "user", "content": []}]
        frame_format = {
            "type": "image_url",
            "image_url": {"url": "data:image/jpeg;base64,{}"},
227
            "modalities": "video",
228
229
230
231
232
233
234
235
236
237
238
239
240
        }

        for base64_frame in base64_frames:
            frame_format["image_url"]["url"] = "data:image/jpeg;base64,{}".format(
                base64_frame
            )
            messages[0]["content"].append(frame_format.copy())

        prompt = {"type": "text", "text": "Please describe the video in detail."}
        messages[0]["content"].append(prompt)

        return messages

241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
    def prepare_video_messages_video_direct(self, video_path):
        messages = [
            {
                "role": "user",
                "content": [
                    {
                        "type": "image_url",
                        "image_url": {"url": f"video:{video_path}"},
                        "modalities": "video",
                    },
                    {"type": "text", "text": "Please describe the video in detail."},
                ],
            },
        ]
        return messages

257
    def get_or_download_file(self, url: str) -> str:
258
        cache_dir = os.path.expanduser("~/.cache")
259
260
261
262
        if url is None:
            raise ValueError()
        file_name = url.split("/")[-1]
        file_path = os.path.join(cache_dir, file_name)
263
264
265
266
267
268
269
270
        os.makedirs(cache_dir, exist_ok=True)

        if not os.path.exists(file_path):
            response = requests.get(url)
            response.raise_for_status()

            with open(file_path, "wb") as f:
                f.write(response.content)
271
272
273
274
275
        return file_path

    def test_video_chat_completion(self):
        url = VIDEO_JOBS_URL
        file_path = self.get_or_download_file(url)
276
277
278

        client = openai.Client(api_key=self.api_key, base_url=self.base_url)

279
        # messages = self.prepare_video_messages_video_direct(file_path)
280
281
        messages = self.prepare_video_messages(file_path)

Mick's avatar
Mick committed
282
        response = client.chat.completions.create(
283
284
285
286
            model="default",
            messages=messages,
            temperature=0,
            max_tokens=1024,
Mick's avatar
Mick committed
287
            stream=False,
288
        )
289

Mick's avatar
Mick committed
290
291
        video_response = response.choices[0].message.content

292
        print("-" * 30)
Mick's avatar
Mick committed
293
        print(f"Video response:\n{video_response}")
294
295
296
        print("-" * 30)

        # Add assertions to validate the video response
Mick's avatar
Mick committed
297
298
299
300
301
        assert "iPod" in video_response or "device" in video_response, video_response
        assert (
            "man" in video_response
            or "person" in video_response
            or "individual" in video_response
302
            or "speaker" in video_response
Mick's avatar
Mick committed
303
304
305
306
307
        ), video_response
        assert (
            "present" in video_response
            or "examine" in video_response
            or "display" in video_response
308
            or "hold" in video_response
Mick's avatar
Mick committed
309
310
        )
        assert "black" in video_response or "dark" in video_response
311
312
313
        self.assertIsNotNone(video_response)
        self.assertGreater(len(video_response), 0)

Ying Sheng's avatar
Ying Sheng committed
314
315
316
317
    def test_regex(self):
        client = openai.Client(api_key=self.api_key, base_url=self.base_url)

        regex = (
318
319
320
            r"""\{"""
            + r""""color":"[\w]+","""
            + r""""number_of_cars":[\d]+"""
Ying Sheng's avatar
Ying Sheng committed
321
322
323
324
325
326
327
328
329
330
331
            + r"""\}"""
        )

        response = client.chat.completions.create(
            model="default",
            messages=[
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "image_url",
332
                            "image_url": {"url": IMAGE_MAN_IRONING_URL},
Ying Sheng's avatar
Ying Sheng committed
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
                        },
                        {
                            "type": "text",
                            "text": "Describe this image in the JSON format.",
                        },
                    ],
                },
            ],
            temperature=0,
            extra_body={"regex": regex},
        )
        text = response.choices[0].message.content

        try:
            js_obj = json.loads(text)
        except (TypeError, json.decoder.JSONDecodeError):
            print("JSONDecodeError", text)
            raise
        assert isinstance(js_obj["color"], str)
        assert isinstance(js_obj["number_of_cars"], int)

354
355
356
357
358
359
360
361
    def run_decode_with_image(self, image_id):
        client = openai.Client(api_key=self.api_key, base_url=self.base_url)

        content = []
        if image_id == 0:
            content.append(
                {
                    "type": "image_url",
362
                    "image_url": {"url": IMAGE_MAN_IRONING_URL},
363
364
365
366
367
368
                }
            )
        elif image_id == 1:
            content.append(
                {
                    "type": "image_url",
369
                    "image_url": {"url": IMAGE_SGL_LOGO_URL},
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
                }
            )
        else:
            pass

        content.append(
            {
                "type": "text",
                "text": "Describe this image in a very short sentence.",
            }
        )

        response = client.chat.completions.create(
            model="default",
            messages=[
                {"role": "user", "content": content},
            ],
            temperature=0,
        )

        assert response.choices[0].message.role == "assistant"
        text = response.choices[0].message.content
        assert isinstance(text, str)

    def test_mixed_batch(self):
        image_ids = [0, 1, 2] * 4
        with ThreadPoolExecutor(4) as executor:
            list(executor.map(self.run_decode_with_image, image_ids))

Mick's avatar
Mick committed
399
400
401
402
403
404
405
406
407
    def prepare_audio_messages(self, prompt, audio_file_name):
        messages = [
            {
                "role": "user",
                "content": [
                    {
                        "type": "audio_url",
                        "audio_url": {"url": f"{audio_file_name}"},
                    },
Mick's avatar
Mick committed
408
409
410
411
                    {
                        "type": "text",
                        "text": prompt,
                    },
Mick's avatar
Mick committed
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
                ],
            }
        ]

        return messages

    def get_audio_response(self, url: str, prompt, category):
        audio_file_path = self.get_or_download_file(url)
        client = openai.Client(api_key="sk-123456", base_url=self.base_url)

        messages = self.prepare_audio_messages(prompt, audio_file_path)

        response = client.chat.completions.create(
            model="default",
            messages=messages,
            temperature=0,
            max_tokens=128,
            stream=False,
        )

        audio_response = response.choices[0].message.content

        print("-" * 30)
        print(f"audio {category} response:\n{audio_response}")
        print("-" * 30)

        audio_response = audio_response.lower()

        self.assertIsNotNone(audio_response)
        self.assertGreater(len(audio_response), 0)

        return audio_response

    def _test_audio_speech_completion(self):
        # a fragment of Trump's speech
        audio_response = self.get_audio_response(
            AUDIO_TRUMP_SPEECH_URL,
            "I have an audio sample. Please repeat the person's words",
            category="speech",
        )
        assert "thank you" in audio_response
        assert "it's a privilege to be here" in audio_response
        assert "leader" in audio_response
        assert "science" in audio_response
        assert "art" in audio_response

    def _test_audio_ambient_completion(self):
        # bird song
        audio_response = self.get_audio_response(
            AUDIO_BIRD_SONG_URL,
            "Please listen to the audio snippet carefully and transcribe the content.",
            "ambient",
        )
        assert "bird" in audio_response

    def test_audio_chat_completion(self):
        pass