A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 8, 20:15 UTC
[Bugfix] Fix corner case of sparse embedding (#33886)
Signed-off-by: wang.yuqi <yuqi.wang@daocloud.io>
wang.yuqi committed
Feb 5, 2026 at 18:51 UTC
1c3a221d3b0f7a82cd9a6d56e10ea360e2435a1c
2 files changed
+11
-1
tests/models/language/pooling/test_bge_m3.py
+10
@@ -136,6 +136,16 @@ async def test_bge_m3_api_server_sparse_embedding(client: openai.AsyncOpenAI):
136
)
137
138
139
+@pytest.mark.asyncio
140
+async def test_bge_m3_api_server_sparse_embedding_corner_case(
141
+ client: openai.AsyncOpenAI,
142
+):
143
+ embeddings = await sparse_embeddings(client, ["Hi"])
144
+ assert len(embeddings) == 1
145
+ assert 2673 in embeddings[0]
146
+ assert embeddings[0][2673] == pytest.approx(0.26710861921310425, rel=0.01)
147
+
148
+
149
# https://github.com/FlagOpen/FlagEmbedding/blob/6fd176266f2382878bcc69cd656cff425d52f49b/FlagEmbedding/inference/embedder/encoder_only/m3.py#L163
150
def colbert_score(q_reps: torch.Tensor, p_reps: torch.Tensor) -> torch.Tensor:
151
token_scores = torch.einsum("in,jn->ij", q_reps, p_reps)
vllm/model_executor/layers/pooler/special.py
+1
-1
@@ -165,7 +165,7 @@ class BOSEOSFilter(Pooler):
165
pooled_data = pooled_data[1:]
166
if token_ids[-1] == self.eos_token_id:
167
pooled_data = pooled_data[:-1]
168
- pooled_outputs[i] = pooled_data.squeeze()
168
+ pooled_outputs[i] = pooled_data.squeeze(-1)
169
170
return pooled_outputs
171