@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 121 lines 3.76 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Parity test between Cohere /v2/embed and OpenAI /v1/embeddings.
4
5 Verifies that both endpoints produce identical float embeddings when
6 no prompt prefix is applied (input_type omitted for Cohere /v2/embed).
7 """
8
9 import numpy as np
10 import pytest
11 import requests
12
13 from tests.utils import ROCM_EXTRA_ARGS, RemoteOpenAIServer
14
15 MODEL_NAME = "BAAI/bge-base-en-v1.5"
16 DTYPE = "bfloat16"
17
18
19 @pytest.fixture(scope="module")
20 def server():
21 args = [
22 "--runner",
23 "pooling",
24 "--dtype",
25 DTYPE,
26 "--enforce-eager",
27 "--max-model-len",
28 "512",
29 "--gpu-memory-utilization",
30 "0.02",
31 ] + ROCM_EXTRA_ARGS
32 with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
33 yield remote_server
34
35
36 def _cohere_embed(
37 server: RemoteOpenAIServer,
38 texts: list[str],
39 ) -> list[list[float]]:
40 body = {
41 "model": MODEL_NAME,
42 "texts": texts,
43 "embedding_types": ["float"],
44 }
45 resp = requests.post(server.url_for("/v2/embed"), json=body)
46 resp.raise_for_status()
47 return resp.json()["embeddings"]["float"]
48
49
50 def _openai_embed(
51 server: RemoteOpenAIServer,
52 texts: list[str],
53 ) -> list[list[float]]:
54 body = {"model": MODEL_NAME, "input": texts, "encoding_format": "float"}
55 resp = requests.post(server.url_for("/v1/embeddings"), json=body)
56 resp.raise_for_status()
57 return [item["embedding"] for item in resp.json()["data"]]
58
59
60 def _cosine_sim(a: list[float], b: list[float]) -> float:
61 va, vb = np.array(a), np.array(b)
62 return float(np.dot(va, vb) / (np.linalg.norm(va) * np.linalg.norm(vb)))
63
64
65 def test_single_text_parity(server: RemoteOpenAIServer):
66 """A single text should produce equivalent embeddings via both APIs."""
67 texts = ["the quick brown fox jumps over the lazy dog"]
68 v2 = _cohere_embed(server, texts)
69 v1 = _openai_embed(server, texts)
70 # Full-suite BF16 runs can introduce tiny numerical drift even when both
71 # endpoints are functionally equivalent, so compare semantic equivalence
72 # instead of exact elementwise equality.
73 cos = _cosine_sim(v2[0], v1[0])
74 assert cos > 0.9999, f"single-text parity failed, cosine={cos}"
75
76
77 def test_batch_parity(server: RemoteOpenAIServer):
78 """A batch of texts should produce equivalent embeddings via both APIs,
79 in the same order."""
80 texts = [
81 "machine learning",
82 "deep learning",
83 "natural language processing",
84 ]
85 v2 = _cohere_embed(server, texts)
86 v1 = _openai_embed(server, texts)
87 assert len(v2) == len(v1) == 3
88
89 similarities = np.array(
90 [[_cosine_sim(v2_emb, v1_emb) for v1_emb in v1] for v2_emb in v2]
91 )
92 for i in range(3):
93 assert int(np.argmax(similarities[i])) == i, (
94 f"batch parity order mismatch at index {i}: "
95 f"similarities={similarities[i].tolist()}"
96 )
97 assert similarities[i, i] > 0.9999, (
98 f"batch parity failed at index {i}, cosine={similarities[i, i]}"
99 )
100
101
102 def test_token_count_parity(server: RemoteOpenAIServer):
103 """Both APIs should report the same prompt token count."""
104 texts = ["hello world"]
105 v2_resp = requests.post(
106 server.url_for("/v2/embed"),
107 json={
108 "model": MODEL_NAME,
109 "texts": texts,
110 "embedding_types": ["float"],
111 },
112 )
113 v1_resp = requests.post(
114 server.url_for("/v1/embeddings"),
115 json={"model": MODEL_NAME, "input": texts, "encoding_format": "float"},
116 )
117 v2_resp.raise_for_status()
118 v1_resp.raise_for_status()
119 v2_tokens = v2_resp.json()["meta"]["billed_units"]["input_tokens"]
120 v1_tokens = v1_resp.json()["usage"]["prompt_tokens"]
121 assert v2_tokens == v1_tokens