Update app.py

app.py

@@ -13,7 +13,16 @@ def get_embedding(text):
     inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
     with torch.no_grad():
         output = model(**inputs)
-    return output.last_hidden_state[:, 0, :].squeeze()  # Use CLS token embedding
+    
+    # Mean pooling over token embeddings
+    embeddings = output.last_hidden_state  # Shape: (batch_size, seq_len, hidden_dim)
+    attention_mask = inputs["attention_mask"].unsqueeze(-1)  # Shape: (batch_size, seq_len, 1)
+    
+    # Apply mean pooling: Sum(token_embeddings * mask) / Sum(mask)
+    pooled_embedding = (embeddings * attention_mask).sum(dim=1) / attention_mask.sum(dim=1)
+
+    # Normalize embedding
+    return F.normalize(pooled_embedding, p=2, dim=1).squeeze()
 
 def get_similarity_and_excerpt(query, paragraph1, paragraph2, paragraph3, threshold_weight):
     paragraphs = [p for p in [paragraph1, paragraph2, paragraph3] if p.strip()]