Alverciito
/

wikipedia_segmentation

Sentence Similarity

sentence_cosenet

feature-extraction

sentence-embeddings

information-retrieval

semantic-search

Model card Files Files and versions

alverciito commited on Jan 17

Commit

00f1b20

·

1 Parent(s): 6faa82b

fix docstrings

Files changed (1) hide show

model.py +2 -2

model.py CHANGED Viewed

@@ -232,14 +232,14 @@ class SentenceCoseNet(PreTrainedModel):
         Returns:
             torch.Tensor:
-                Similarity scores of shape (B, S, S)
         """
         # Concatenate embeddings (B, S, 2, D)
         embeddings = torch.stack([embeddings_1, embeddings_2], dim=-2)
         # Compute distances (B, S, 2, 2):
         embeddings = self.model.distance_layer(embeddings)
         # Return cosine similarities (B, S):
-        return embeddings[..., 0, 1]
     def forward(
         self,

         Returns:
             torch.Tensor:
+                Similarity scores of shape (B, S)
         """
         # Concatenate embeddings (B, S, 2, D)
         embeddings = torch.stack([embeddings_1, embeddings_2], dim=-2)
         # Compute distances (B, S, 2, 2):
         embeddings = self.model.distance_layer(embeddings)
         # Return cosine similarities (B, S):
+        return (embeddings[..., 0, 1] + embeddings[..., 1, 0]) / 2
     def forward(
         self,