@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 135 lines 3.75 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Tests for the Cohere /v2/embed API with a multimodal model (SigLIP).
4
5 Validates image embedding, batching, normalisation, and embedding type
6 conversions through the /v2/embed endpoint.
7 """
8
9 import struct
10 import zlib
11
12 import numpy as np
13 import pybase64 as base64
14 import pytest
15 import requests
16
17 from tests.utils import RemoteOpenAIServer
18
19 MODEL_NAME = "google/siglip-so400m-patch14-384"
20 DTYPE = "bfloat16"
21
22
23 @pytest.fixture(scope="module")
24 def server():
25 args = [
26 "--runner",
27 "pooling",
28 "--dtype",
29 DTYPE,
30 "--enforce-eager",
31 "--max-model-len",
32 "64",
33 "--gpu-memory-utilization",
34 "0.3",
35 ]
36 with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
37 yield remote_server
38
39
40 def _make_tiny_png(r: int, g: int, b: int, w: int = 2, h: int = 2) -> str:
41 raw = b""
42 for _ in range(h):
43 raw += b"\x00" + bytes([r, g, b]) * w
44 compressed = zlib.compress(raw)
45
46 def chunk(ctype: bytes, cdata: bytes) -> bytes:
47 c = ctype + cdata
48 return (
49 struct.pack(">I", len(cdata))
50 + c
51 + struct.pack(">I", zlib.crc32(c) & 0xFFFFFFFF)
52 )
53
54 ihdr = struct.pack(">IIBBBBB", w, h, 8, 2, 0, 0, 0)
55 png = (
56 b"\x89PNG\r\n\x1a\n"
57 + chunk(b"IHDR", ihdr)
58 + chunk(b"IDAT", compressed)
59 + chunk(b"IEND", b"")
60 )
61 return "data:image/png;base64," + base64.b64encode(png).decode()
62
63
64 def _cohere_embed(
65 server: RemoteOpenAIServer,
66 texts: list[str] | None = None,
67 images: list[str] | None = None,
68 embedding_types: list[str] | None = None,
69 ) -> dict:
70 body: dict = {"model": MODEL_NAME}
71 if texts is not None:
72 body["texts"] = texts
73 if images is not None:
74 body["images"] = images
75 if embedding_types is not None:
76 body["embedding_types"] = embedding_types
77 resp = requests.post(server.url_for("/v2/embed"), json=body)
78 resp.raise_for_status()
79 return resp.json()
80
81
82 def test_image_embed(server: RemoteOpenAIServer):
83 img_uri = _make_tiny_png(255, 0, 0)
84 r = _cohere_embed(
85 server,
86 images=[img_uri],
87 embedding_types=["float"],
88 )
89 assert "embeddings" in r
90 assert len(r["embeddings"]["float"]) == 1
91 assert len(r["embeddings"]["float"][0]) > 0
92 assert r["meta"]["billed_units"]["image_tokens"] > 0
93 assert r["meta"]["billed_units"]["input_tokens"] == 0
94
95
96 def test_image_batch(server: RemoteOpenAIServer):
97 red = _make_tiny_png(255, 0, 0)
98 blue = _make_tiny_png(0, 0, 255)
99 r = _cohere_embed(
100 server,
101 images=[red, blue],
102 embedding_types=["float"],
103 )
104 assert len(r["embeddings"]["float"]) == 2
105
106
107 def test_image_l2_normalized(server: RemoteOpenAIServer):
108 img_uri = _make_tiny_png(0, 255, 0)
109 r = _cohere_embed(
110 server,
111 images=[img_uri],
112 embedding_types=["float"],
113 )
114 emb = np.array(r["embeddings"]["float"][0])
115 assert abs(float(np.linalg.norm(emb)) - 1.0) < 0.01
116
117
118 def test_image_embedding_types(server: RemoteOpenAIServer):
119 img_uri = _make_tiny_png(128, 128, 128)
120 r = _cohere_embed(
121 server,
122 images=[img_uri],
123 embedding_types=["float", "binary", "ubinary"],
124 )
125 dim = len(r["embeddings"]["float"][0])
126 assert len(r["embeddings"]["binary"][0]) == dim // 8
127 assert len(r["embeddings"]["ubinary"][0]) == dim // 8
128
129
130 def test_text_embed_on_multimodal(server: RemoteOpenAIServer):
131 """SigLIP also supports text-only embedding via /v2/embed."""
132 r = _cohere_embed(server, texts=["hello world"], embedding_types=["float"])
133 assert "embeddings" in r
134 assert len(r["embeddings"]["float"]) == 1
135 assert len(r["embeddings"]["float"][0]) > 0