test_vision_openai_server_common.py 18.5 KB
Newer Older
1
2
3
import base64
import io
import os
Ying Sheng's avatar
Ying Sheng committed
4

5
import numpy as np
Ying Sheng's avatar
Ying Sheng committed
6
import openai
7
8
import requests
from PIL import Image
Ying Sheng's avatar
Ying Sheng committed
9

10
from sglang.srt.utils import kill_process_tree
11
from sglang.test.test_utils import DEFAULT_URL_FOR_TEST, CustomTestCase
Ying Sheng's avatar
Ying Sheng committed
12

13
14
15
16
17
18
19
20
21
22
23
# 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
24

25
class TestOpenAIOmniServerBase(CustomTestCase):
Ying Sheng's avatar
Ying Sheng committed
26
27
    @classmethod
    def setUpClass(cls):
28
        cls.model = ""
29
        cls.base_url = DEFAULT_URL_FOR_TEST
Ying Sheng's avatar
Ying Sheng committed
30
        cls.api_key = "sk-123456"
31
        cls.process = None
Ying Sheng's avatar
Ying Sheng committed
32
33
34
35
        cls.base_url += "/v1"

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

38
39
40
    def get_vision_request_kwargs(self):
        return self.get_request_kwargs()

41
42
43
    def get_request_kwargs(self):
        return {}

44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
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
139
140
141
    def get_or_download_file(self, url: str) -> str:
        cache_dir = os.path.expanduser("~/.cache")
        if url is None:
            raise ValueError()
        file_name = url.split("/")[-1]
        file_path = os.path.join(cache_dir, file_name)
        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)
        return file_path


class AudioOpenAITestMixin(TestOpenAIOmniServerBase):
    def prepare_audio_messages(self, prompt, audio_file_name):
        messages = [
            {
                "role": "user",
                "content": [
                    {
                        "type": "audio_url",
                        "audio_url": {"url": f"{audio_file_name}"},
                    },
                    {
                        "type": "text",
                        "text": prompt,
                    },
                ],
            }
        ]

        return messages

    def get_audio_request_kwargs(self):
        return self.get_request_kwargs()

    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,
            **(self.get_audio_request_kwargs()),
        )

        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.lower()

    def test_audio_speech_completion(self):
        # a fragment of Trump's speech
        audio_response = self.get_audio_response(
            AUDIO_TRUMP_SPEECH_URL,
            "Listen to this audio and write down the audio transcription in English.",
            category="speech",
        )
        check_list = [
            "thank you",
            "it's a privilege to be here",
            "leader",
            "science",
            "art",
        ]
        for check_word in check_list:
            assert (
                check_word in audio_response
            ), f"audio_response: |{audio_response}| should contain |{check_word}|"

    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 in English.",
            "ambient",
        )
        assert "bird" in audio_response


class ImageOpenAITestMixin(TestOpenAIOmniServerBase):
142
    def test_single_image_chat_completion(self):
Ying Sheng's avatar
Ying Sheng committed
143
144
145
146
147
148
149
150
151
152
        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",
153
                            "image_url": {"url": IMAGE_MAN_IRONING_URL},
Ying Sheng's avatar
Ying Sheng committed
154
                        },
Ying Sheng's avatar
Ying Sheng committed
155
156
                        {
                            "type": "text",
157
                            "text": "Describe this image in a sentence.",
Ying Sheng's avatar
Ying Sheng committed
158
                        },
Ying Sheng's avatar
Ying Sheng committed
159
160
161
162
                    ],
                },
            ],
            temperature=0,
163
            **(self.get_vision_request_kwargs()),
Ying Sheng's avatar
Ying Sheng committed
164
165
166
        )

        assert response.choices[0].message.role == "assistant"
Ying Sheng's avatar
Ying Sheng committed
167
168
        text = response.choices[0].message.content
        assert isinstance(text, str)
169
        # `driver` is for gemma-3-it
170
171
172
173
174
175
176
177
178
179
        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
180
        # MiniCPMO fails to recognize `iron`, but `hanging`
181
        assert (
182
183
184
185
186
187
188
            "iron" in text
            or "hang" in text
            or "cloth" in text
            or "coat" in text
            or "holding" in text
            or "outfit" in text
        ), f"text: {text}, should contain iron, hang, cloth, coat or holding or outfit"
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
        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",
206
                            "image_url": {"url": IMAGE_MAN_IRONING_URL},
207
208
209
                        },
                        {
                            "type": "text",
210
                            "text": "Describe this image in a sentence.",
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
                        },
                    ],
                },
                {
                    "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,
231
            **(self.get_vision_request_kwargs()),
232
233
234
235
236
        )

        assert response.choices[0].message.role == "assistant"
        text = response.choices[0].message.content
        assert isinstance(text, str)
237
238
239
        assert (
            "man" in text or "cab" in text
        ), f"text: {text}, should contain man or cab"
Ying Sheng's avatar
Ying Sheng committed
240
241
242
243
        assert response.id
        assert response.created
        assert response.usage.prompt_tokens > 0
        assert response.usage.completion_tokens > 0
244
245
        assert response.usage.total_tokens > 0

246
    def test_multi_images_chat_completion(self):
247
248
249
250
251
252
253
254
255
256
        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
257
                            "image_url": {"url": IMAGE_MAN_IRONING_URL},
258
                            "modalities": "multi-images",
259
260
261
                        },
                        {
                            "type": "image_url",
262
                            "image_url": {"url": IMAGE_SGL_LOGO_URL},
263
                            "modalities": "multi-images",
264
265
266
                        },
                        {
                            "type": "text",
267
268
                            "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.",
269
270
271
272
273
                        },
                    ],
                },
            ],
            temperature=0,
274
            **(self.get_vision_request_kwargs()),
275
276
277
278
279
        )

        assert response.choices[0].message.role == "assistant"
        text = response.choices[0].message.content
        assert isinstance(text, str)
Mick's avatar
Mick committed
280
281
282
        print("-" * 30)
        print(f"Multi images response:\n{text}")
        print("-" * 30)
283
        assert (
284
285
286
287
288
289
            "man" in text
            or "cab" in text
            or "SUV" in text
            or "taxi" in text
            or "car" in text
        ), f"text: {text}, should contain man, cab, SUV, taxi or car"
290
        assert (
291
292
            "logo" in text or '"S"' in text or "SG" in text or "graphic" in text
        ), f"text: {text}, should contain logo, S or SG or graphic"
293
294
295
296
        assert response.id
        assert response.created
        assert response.usage.prompt_tokens > 0
        assert response.usage.completion_tokens > 0
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
        assert response.usage.total_tokens > 0

    def _test_mixed_image_audio_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",
                            "image_url": {"url": IMAGE_MAN_IRONING_URL},
                        },
                        {
                            "type": "audio_url",
                            "audio_url": {"url": AUDIO_TRUMP_SPEECH_URL},
                        },
                        {
                            "type": "text",
                            "text": "Please describe the image in one sentence, and then write down the audio transcription in English.",
                        },
                    ],
                },
            ],
            temperature=0,
            **(self.get_vision_request_kwargs()),
        )

        assert response.choices[0].message.role == "assistant"
        text = response.choices[0].message.content
        assert isinstance(text, str)
        print("-" * 30)
        print(f"Mixed image & audio response:\n{text}")
        print("-" * 30)
        assert (
            "man" in text
            or "cab" in text
            or "SUV" in text
            or "taxi" in text
            or "car" in text
        ), f"text: {text}, should contain man, cab, SUV, taxi or car"
        check_list = [
            "thank you",
            "it's a privilege to be here",
            "leader",
            "science",
            "art",
        ]
        for check_word in check_list:
            assert (
                check_word in text
            ), f"text: |{text}| should contain |{check_word}|"
        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
355
356
        assert response.usage.total_tokens > 0

357
    def prepare_video_images_messages(self, video_path):
358
359
        # 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
360
361
362
363
364
365
366

        # 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

367
        max_frames_num = 10
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
        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,{}"},
388
            "modalities": "image",
389
390
391
392
393
394
395
396
397
398
399
400
401
        }

        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

402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
    def test_video_images_chat_completion(self):
        url = VIDEO_JOBS_URL
        file_path = self.get_or_download_file(url)

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

        messages = self.prepare_video_images_messages(file_path)

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

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

        print("-" * 30)
        print(f"Video images response:\n{video_response}")
        print("-" * 30)

        # Add assertions to validate the video response
        assert (
            "iPod" in video_response
            or "device" in video_response
            or "microphone" in video_response
429
430
431
432
433
434
        ), f"""
        ====================== video_response =====================
        {video_response}
        ===========================================================
        should contain 'iPod' or 'device' or 'microphone'
        """
435
436
437
438
439
        assert (
            "man" in video_response
            or "person" in video_response
            or "individual" in video_response
            or "speaker" in video_response
440
            or "presenter" in video_response
441
            or "Steve" in video_response
442
            or "hand" in video_response
443
444
445
446
        ), f"""
        ====================== video_response =====================
        {video_response}
        ===========================================================
447
        should contain 'man' or 'person' or 'individual' or 'speaker' or 'presenter' or 'Steve' or 'hand'
448
        """
449
450
451
452
453
        assert (
            "present" in video_response
            or "examine" in video_response
            or "display" in video_response
            or "hold" in video_response
454
455
456
457
458
459
        ), f"""
        ====================== video_response =====================
        {video_response}
        ===========================================================
        should contain 'present' or 'examine' or 'display' or 'hold'
        """
460
461
462
        self.assertIsNotNone(video_response)
        self.assertGreater(len(video_response), 0)

463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480

class VideoOpenAITestMixin(TestOpenAIOmniServerBase):
    def prepare_video_messages(self, video_path):
        messages = [
            {
                "role": "user",
                "content": [
                    {
                        "type": "video_url",
                        "video_url": {"url": f"{video_path}"},
                    },
                    {"type": "text", "text": "Please describe the video in detail."},
                ],
            },
        ]
        return messages

    def test_video_chat_completion(self):
481
482
        url = VIDEO_JOBS_URL
        file_path = self.get_or_download_file(url)
483
484
485
486
487

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

        messages = self.prepare_video_messages(file_path)

Mick's avatar
Mick committed
488
        response = client.chat.completions.create(
489
490
491
492
            model="default",
            messages=messages,
            temperature=0,
            max_tokens=1024,
Mick's avatar
Mick committed
493
            stream=False,
494
            **(self.get_vision_request_kwargs()),
495
        )
496

Mick's avatar
Mick committed
497
498
        video_response = response.choices[0].message.content

499
        print("-" * 30)
Mick's avatar
Mick committed
500
        print(f"Video response:\n{video_response}")
501
502
503
        print("-" * 30)

        # Add assertions to validate the video response
504
        assert (
505
506
507
            "iPod" in video_response
            or "device" in video_response
            or "microphone" in video_response
508
        ), f"video_response: {video_response}, should contain 'iPod' or 'device'"
Mick's avatar
Mick committed
509
510
511
512
        assert (
            "man" in video_response
            or "person" in video_response
            or "individual" in video_response
513
            or "speaker" in video_response
514
            or "presenter" in video_response
515
            or "hand" in video_response
516
        ), f"video_response: {video_response}, should either have 'man' in video_response, or 'person' in video_response, or 'individual' in video_response or 'speaker' in video_response or 'presenter' or 'hand' in video_response"
Mick's avatar
Mick committed
517
518
519
520
        assert (
            "present" in video_response
            or "examine" in video_response
            or "display" in video_response
521
            or "hold" in video_response
522
523
524
525
        ), f"video_response: {video_response}, should contain 'present', 'examine', 'display', or 'hold'"
        assert (
            "black" in video_response or "dark" in video_response
        ), f"video_response: {video_response}, should contain 'black' or 'dark'"
526
527
        self.assertIsNotNone(video_response)
        self.assertGreater(len(video_response), 0)