Skip to content

fix: per-sequence token counts in batch embedding averaging #3792

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
alvinttang wants to merge 2 commits into
base: main
Choose a base branch
from
+8 −4
Open

fix: per-sequence token counts in batch embedding averaging #3792

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
5d168c5
fix: use per-sequence token counts in batch embedding averaging
Mar 11, 2026
67b8d42
style: apply black formatting to model_worker.py
Mar 16, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 8 additions & 4 deletions fastchat/serve/model_worker.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -171,7 +171,7 @@ def __process_embed_chunk(self, input_ids, attention_mask, **model_type_dict):
mask = attention_mask.unsqueeze(-1).expand(data.size()).float()
masked_embeddings = data * mask
sum_embeddings = torch.sum(masked_embeddings, dim=1)
token_num = torch.sum(attention_mask).item()
token_num = attention_mask.sum(dim=1, keepdim=True)

return sum_embeddings, token_num

Expand Down Expand Up @@ -224,10 +224,10 @@ def get_embeddings(self, params):
):
embedding = embedding / token_num
normalized_embeddings = F.normalize(embedding, p=2, dim=1)
ret["token_num"] = token_num
ret["token_num"] = token_num.sum().item()
else:
all_embeddings = []
all_token_num = 0
all_token_num = 0 # per-sequence tensor, accumulated across chunks
for i in range(0, input_ids.size(1), self.context_len):
chunk_input_ids = input_ids[:, i : i + self.context_len]
chunk_attention_mask = attention_mask[:, i : i + self.context_len]
Expand Down Expand Up @@ -273,7 +273,11 @@ def get_embeddings(self, params):
embedding = torch.sum(all_embeddings_tensor, dim=0) / all_token_num
normalized_embeddings = F.normalize(embedding, p=2, dim=1)

ret["token_num"] = all_token_num
ret["token_num"] = (
all_token_num.sum().item()
if isinstance(all_token_num, torch.Tensor)
else all_token_num
)

if base64_encode == "base64":
out_embeddings = self.__encode_base64(normalized_embeddings)
Expand Down
Loading