| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import httpx |
| 4 | import openai |
| 5 | import pytest |
| 6 | import torch |
| 7 | |
| 8 | from ....utils import RemoteOpenAIServer |
| 9 | from .embed_utils import run_client_embeddings |
| 10 | |
| 11 | MODEL_NAME = "BAAI/bge-m3" |
| 12 | MAX_MODEL_LEN = 512 |
| 13 | |
| 14 | |
| 15 | # Example from https://huggingface.co/BAAI/bge-m3 |
| 16 | sentences_1 = ["What is BGE M3?", "Definition of BM25"] |
| 17 | sentences_2 = [ |
| 18 | "BGE M3 is an embedding model supporting dense retrieval, " |
| 19 | "lexical matching and multi-vector interaction.", |
| 20 | "BM25 is a bag-of-words retrieval function that ranks a set " |
| 21 | "of documents based on the query terms appearing in each document", |
| 22 | ] |
| 23 | |
| 24 | similarity_reference = [[0.6259, 0.3474], [0.3309, 0.6734]] |
| 25 | lexical_score_reference = [0.19554901123046875, 0.0] |
| 26 | colbert_score_reference = [0.7797, 0.4620] |
| 27 | SUPPORTED_TASKS = ["embed", "token_embed", "token_classify"] |
| 28 | |
| 29 | |
| 30 | @pytest.fixture(scope="module", params=SUPPORTED_TASKS) |
| 31 | def pooling_task(request): |
| 32 | yield request.param |
| 33 | |
| 34 | |
| 35 | @pytest.fixture(scope="module") |
| 36 | def server(pooling_task): |
| 37 | args = [ |
| 38 | "--max-model-len", |
| 39 | str(MAX_MODEL_LEN), |
| 40 | "--hf-overrides", |
| 41 | '{"architectures": ["BgeM3EmbeddingModel"]}', |
| 42 | "--pooler-config.task", |
| 43 | pooling_task, |
| 44 | ] |
| 45 | |
| 46 | with RemoteOpenAIServer(MODEL_NAME, args) as remote_server: |
| 47 | yield remote_server |
| 48 | |
| 49 | |
| 50 | @pytest.mark.asyncio |
| 51 | async def test_bge_m3_api_server_embedding(server, pooling_task): |
| 52 | client = server.get_async_client() |
| 53 | |
| 54 | if pooling_task != "embed": |
| 55 | with pytest.raises(openai.InternalServerError): |
| 56 | await run_client_embeddings( |
| 57 | client, |
| 58 | MODEL_NAME, |
| 59 | sentences_1, |
| 60 | ) |
| 61 | return |
| 62 | |
| 63 | embeddings_list_1 = await run_client_embeddings( |
| 64 | client, |
| 65 | MODEL_NAME, |
| 66 | sentences_1, |
| 67 | ) |
| 68 | embeddings_list_2 = await run_client_embeddings( |
| 69 | client, |
| 70 | MODEL_NAME, |
| 71 | sentences_2, |
| 72 | ) |
| 73 | |
| 74 | embeddings_1 = torch.tensor(embeddings_list_1) |
| 75 | embeddings_2 = torch.tensor(embeddings_list_2) |
| 76 | similarity = embeddings_1 @ embeddings_2.T |
| 77 | |
| 78 | # reference values from BAAI/bge-m3 documentation |
| 79 | reference = torch.tensor(similarity_reference) |
| 80 | |
| 81 | assert torch.allclose(similarity, reference, rtol=0.01) |
| 82 | |
| 83 | |
| 84 | async def tokenize(client: openai.AsyncOpenAI, sentences: list[str]) -> list[list[int]]: |
| 85 | futures = [] |
| 86 | for sentence in sentences: |
| 87 | futures.append( |
| 88 | client.post( |
| 89 | "../tokenize", |
| 90 | body={"model": MODEL_NAME, "prompt": sentence}, |
| 91 | cast_to=httpx.Response, |
| 92 | ) |
| 93 | ) |
| 94 | return [(await future).json()["tokens"] for future in futures] |
| 95 | |
| 96 | |
| 97 | async def sparse_embeddings( |
| 98 | client: openai.AsyncOpenAI, sentences: list[str] |
| 99 | ) -> list[dict[int, float]]: |
| 100 | all_tokens = await tokenize(client, sentences) |
| 101 | result = await client.post( |
| 102 | "../pooling", |
| 103 | body={"model": MODEL_NAME, "input": sentences, "task": "token_classify"}, |
| 104 | cast_to=httpx.Response, |
| 105 | ) |
| 106 | all_embeddings = [data["data"] for data in result.json()["data"]] |
| 107 | |
| 108 | ret = [] |
| 109 | |
| 110 | for sent_tokens, sent_emb in zip(all_tokens, all_embeddings): |
| 111 | token_embs = dict[int, float]() |
| 112 | if sent_tokens[0] == 0: |
| 113 | sent_tokens = sent_tokens[1:] |
| 114 | for token, val in zip(sent_tokens, sent_emb): |
| 115 | token_embs[token] = max(val, token_embs.get(token, 0.0)) |
| 116 | ret.append(token_embs) |
| 117 | return ret |
| 118 | |
| 119 | |
| 120 | # Based on https://github.com/FlagOpen/FlagEmbedding/blob/6fd176266f2382878bcc69cd656cff425d52f49b/FlagEmbedding/inference/embedder/encoder_only/m3.py#L129 |
| 121 | def compute_lexical_matching_score( |
| 122 | lw1: dict[int, float], lw2: dict[int, float] |
| 123 | ) -> float: |
| 124 | scores = 0.0 |
| 125 | for token, weight in lw1.items(): |
| 126 | if token in lw2: |
| 127 | scores += weight * lw2[token] |
| 128 | return scores |
| 129 | |
| 130 | |
| 131 | @pytest.mark.asyncio |
| 132 | async def test_bge_m3_api_server_sparse_embedding(server, pooling_task): |
| 133 | client = server.get_async_client() |
| 134 | |
| 135 | if pooling_task != "token_classify": |
| 136 | with pytest.raises(openai.BadRequestError): |
| 137 | await sparse_embeddings(client, sentences_1) |
| 138 | return |
| 139 | |
| 140 | embeddings_1 = await sparse_embeddings(client, sentences_1) |
| 141 | embeddings_2 = await sparse_embeddings(client, sentences_2) |
| 142 | |
| 143 | lexical_scores_1_0_x_2_0 = compute_lexical_matching_score( |
| 144 | embeddings_1[0], embeddings_2[0] |
| 145 | ) |
| 146 | assert lexical_scores_1_0_x_2_0 == pytest.approx( |
| 147 | lexical_score_reference[0], rel=0.01 |
| 148 | ) |
| 149 | |
| 150 | lexical_scores_1_0_x_1_1 = compute_lexical_matching_score( |
| 151 | embeddings_1[0], embeddings_1[1] |
| 152 | ) |
| 153 | assert lexical_scores_1_0_x_1_1 == pytest.approx( |
| 154 | lexical_score_reference[1], rel=0.01 |
| 155 | ) |
| 156 | |
| 157 | |
| 158 | @pytest.mark.asyncio |
| 159 | async def test_bge_m3_api_server_sparse_embedding_corner_case(server, pooling_task): |
| 160 | if pooling_task != "token_classify": |
| 161 | return |
| 162 | |
| 163 | client = server.get_async_client() |
| 164 | embeddings = await sparse_embeddings(client, ["Hi"]) |
| 165 | assert len(embeddings) == 1 |
| 166 | assert 2673 in embeddings[0] |
| 167 | assert embeddings[0][2673] == pytest.approx(0.26710861921310425, rel=0.01) |
| 168 | |
| 169 | |
| 170 | # https://github.com/FlagOpen/FlagEmbedding/blob/6fd176266f2382878bcc69cd656cff425d52f49b/FlagEmbedding/inference/embedder/encoder_only/m3.py#L163 |
| 171 | def colbert_score(q_reps: torch.Tensor, p_reps: torch.Tensor) -> torch.Tensor: |
| 172 | token_scores = torch.einsum("in,jn->ij", q_reps, p_reps) |
| 173 | scores, _ = token_scores.max(-1) |
| 174 | scores = torch.sum(scores) / q_reps.size(0) |
| 175 | return scores |
| 176 | |
| 177 | |
| 178 | @pytest.mark.asyncio |
| 179 | async def test_bge_m3_api_server_multi_vector(server, pooling_task): |
| 180 | client = server.get_async_client() |
| 181 | |
| 182 | if pooling_task != "token_embed": |
| 183 | with pytest.raises(openai.BadRequestError): |
| 184 | await client.post( |
| 185 | "../pooling", |
| 186 | body={"model": MODEL_NAME, "input": sentences_1, "task": "token_embed"}, |
| 187 | cast_to=httpx.Response, |
| 188 | ) |
| 189 | return |
| 190 | |
| 191 | result_1 = await client.post( |
| 192 | "../pooling", |
| 193 | body={"model": MODEL_NAME, "input": sentences_1, "task": "token_embed"}, |
| 194 | cast_to=httpx.Response, |
| 195 | ) |
| 196 | embeddings_1 = [torch.tensor(data["data"]) for data in result_1.json()["data"]] |
| 197 | |
| 198 | result_2 = await client.post( |
| 199 | "../pooling", |
| 200 | body={"model": MODEL_NAME, "input": sentences_2, "task": "token_embed"}, |
| 201 | cast_to=httpx.Response, |
| 202 | ) |
| 203 | embeddings_2 = [torch.tensor(data["data"]) for data in result_2.json()["data"]] |
| 204 | |
| 205 | colbert_score_1_0_x_2_0 = colbert_score(embeddings_1[0], embeddings_2[0]) |
| 206 | assert colbert_score_1_0_x_2_0 == pytest.approx( |
| 207 | colbert_score_reference[0], rel=0.01 |
| 208 | ) |
| 209 | colbert_score_1_0_x_2_1 = colbert_score(embeddings_1[0], embeddings_2[1]) |
| 210 | assert colbert_score_1_0_x_2_1 == pytest.approx( |
| 211 | colbert_score_reference[1], rel=0.01 |
| 212 | ) |