Update README.md

README.md

@@ -155,15 +155,17 @@ from torch import Tensor
 from transformers import AutoTokenizer, AutoModel
 
 
-def last_token_pool(last_hidden_states: Tensor,
+def mean_pool(last_hidden_states: Tensor,
                  attention_mask: Tensor) -> Tensor:
-    left_padding = (attention_mask[:, -1].sum() == attention_mask.shape[0])
-    if left_padding:
-        return last_hidden_states[:, -1]
-    else:
-        sequence_lengths = attention_mask.sum(dim=1) - 1
-        batch_size = last_hidden_states.shape[0]
-        return last_hidden_states[torch.arange(batch_size, device=last_hidden_states.device), sequence_lengths]
+
+    seq_lengths = attention_mask.sum(dim=-1)
+    return torch.stack(
+                [
+                    last_hidden_states[i, -length:, :].sum(dim=0) / length
+                    for i, length in enumerate(seq_lengths)
+                ],
+                dim=0,
+            )
 
 
 def get_detailed_instruct(task_description: str, query: str) -> str:
@@ -195,7 +197,7 @@ batch_dict = tokenizer(
 )
 batch_dict.to(model.device)
 outputs = model(**batch_dict)
-embeddings = last_token_pool(outputs.last_hidden_state, batch_dict['attention_mask'])
+embeddings = mean_pool(outputs.last_hidden_state, batch_dict['attention_mask'])
 
 embeddings = F.normalize(embeddings, p=2, dim=1)
 scores = (embeddings[:2] @ embeddings[2:].T)