xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2026-07-09 22:47:14 +08:00
Code Issues Packages Projects Releases Wiki Activity

[Bugfix] Embedding model pooling_type equals ALL and multi input's bug (#10494)

Browse Source
This commit is contained in:
Xiaoyu Zhang 2024-11-21 22:40:02 +08:00 committed by GitHub
parent d5ec121f95
commit 4d676f0852
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 17 additions and 12 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

29
vllm/model_executor/layers/pooler.py
View File

@ -94,14 +94,10 @@ class Pooler(nn.Module):
pooled_data = hidden_states[last_token_flat_indices] pooled_data = hidden_states[last_token_flat_indices]
elif self.pooling_type == PoolingType.ALL: elif self.pooling_type == PoolingType.ALL:
offset = 0 offset = 0
pooled_data_lst = [] pooled_data = []
for prompt_len in prompt_lens: for prompt_len in prompt_lens:
pooled_data_i = hidden_states[offset:offset + prompt_len] pooled_data.append(hidden_states[offset:offset + prompt_len])
pooled_data_lst.append(pooled_data_i)
offset += prompt_len offset += prompt_len
pooled_data = torch.stack(pooled_data_lst)
elif self.pooling_type == PoolingType.MEAN: elif self.pooling_type == PoolingType.MEAN:
# Calculate mean pooling # Calculate mean pooling
cumsum = torch.cumsum(hidden_states, dim=0) cumsum = torch.cumsum(hidden_states, dim=0)
@ -121,7 +117,7 @@ class Pooler(nn.Module):
step_tag_id = self.step_tag_id step_tag_id = self.step_tag_id
offset = 0 offset = 0
pooled_data_lst = [] pooled_data = []
for prompt_len, seq_data_i in zip( for prompt_len, seq_data_i in zip(
prompt_lens, pooling_metadata.seq_data.values()): prompt_lens, pooling_metadata.seq_data.values()):
pooled_data_i = hidden_states[offset:offset + prompt_len] pooled_data_i = hidden_states[offset:offset + prompt_len]
@ -130,17 +126,26 @@ class Pooler(nn.Module):
pooled_data_i = pooled_data_i[token_ids == step_tag_id] pooled_data_i = pooled_data_i[token_ids == step_tag_id]
offset += prompt_len offset += prompt_len
pooled_data_lst.append(pooled_data_i) pooled_data.append(pooled_data_i)
pooled_data = torch.stack(pooled_data_lst)
else: else:
raise ValueError(f"Invalid pooling type: {self.pooling_type}") raise ValueError(f"Invalid pooling type: {self.pooling_type}")
if self.normalize: if self.normalize:
pooled_data = nn.functional.normalize(pooled_data, p=2, dim=1) if isinstance(pooled_data, list):
pooled_data = [
nn.functional.normalize(data, p=2, dim=1)
for data in pooled_data
]
else:
pooled_data = nn.functional.normalize(pooled_data, p=2, dim=1)
if self.softmax: if self.softmax:
pooled_data = nn.functional.softmax(pooled_data, dim=-1) if isinstance(pooled_data, list):
pooled_data = [
nn.functional.softmax(data, dim=-1) for data in pooled_data
]
else:
pooled_data = nn.functional.softmax(pooled_data, dim=-1)
pooled_outputs = [ pooled_outputs = [
EmbeddingSequenceGroupOutput(data.tolist()) for data in pooled_data EmbeddingSequenceGroupOutput(data.tolist()) for data in pooled_data