Update README.md

README.md

@@ -50,7 +50,7 @@ def generate_sentence_embedding(sentence, language):
         model.set_default_language("it_CH")
     if "rm" in language:
         model.set_default_language("rm_CH")
-    
+
     # Tokenize input sentence
     inputs = tokenizer(sentence, padding=True, truncation=True, return_tensors="pt", max_length=512)
 
@@ -58,8 +58,12 @@ def generate_sentence_embedding(sentence, language):
     with torch.no_grad():
         outputs = model(**inputs)
 
-    # Extract average sentence embeddings from the last hidden layer
-    embedding = outputs.last_hidden_state.mean(dim=1)
+    # Extract sentence embeddings via mean pooling
+    token_embeddings = outputs.last_hidden_state
+    attention_mask = inputs['attention_mask'].unsqueeze(-1).expand(token_embeddings.size()).float()
+    sum_embeddings = torch.sum(token_embeddings * attention_mask, 1)
+    sum_mask = torch.clamp(attention_mask.sum(1), min=1e-9)
+    embedding = sum_embeddings / sum_mask
 
     return embedding