Spaces:
Sleeping
Sleeping
File size: 2,893 Bytes
1a3b3aa |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 |
import torch
import torch.nn.functional as F
from sentence_transformers.util import cos_sim
from transformers import AutoTokenizer, AutoModel
class OutcomeSimilarity:
""" similarity detector between outcomes statements"""
ID2LABEL = ["different", "similar"]
def __init__(self, model_path: str):
self.tokenizer = AutoTokenizer.from_pretrained(model_path)
self.model = AutoModel.from_pretrained(model_path)
def _mean_pooling(self, model_output, attention_mask: torch.Tensor):
""" Mean Pooling - Take attention mask into account for correct averaging"""
# First element of model_output contains all token embeddings
token_embeddings = model_output[0]
input_mask_expanded = attention_mask.unsqueeze(
-1).expand(token_embeddings.size()).float()
return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)
def _encode(self, outcomes_lot: list[tuple[str,str]]):
# Parse sentences
sentences = []
if len(outcomes_lot) > 0:
_, sentences = zip(*outcomes_lot)
# Tokenize sentences
encoded_input = self.tokenizer(
sentences, padding=True, truncation=True, return_tensors='pt')
# Compute token embeddings
with torch.no_grad():
model_output = self.model(**encoded_input)
# Perform pooling
sentence_embeddings = self._mean_pooling(
model_output, encoded_input['attention_mask'])
# Normalize embeddings
sentence_embeddings = F.normalize(sentence_embeddings, p=2, dim=1)
return sentence_embeddings
def get_similarity(
self,
registry_outcomes:list[tuple[str,str]],
article_outcomes:list[tuple[str,str]]
) -> list[tuple[int,int,float]]:
"""For each outcome in true_dict, find the most similar outcome in compared_dict and return a mapping
of all matchs , for each tuple : registry is the first index (at i=0); article is the second index (at i=1)
and the third index (i=3) is the cosine similarity score"""
connections = set()
rembs = self._encode(registry_outcomes)
aembs = self._encode(article_outcomes)
cosines_scores = cos_sim(rembs, aembs)
lines_max = torch.argmax(cosines_scores, dim=1)
col_max = torch.argmax(cosines_scores, dim=0)
remaining_cols = set(range(len(col_max)))
for i in range(len(lines_max)):
connection = (i, lines_max[i].item(), cosines_scores[i, lines_max[i]].item())
remaining_cols.discard(lines_max[i].item())
connections.add(connection)
for j in remaining_cols:
connection = (col_max[j].item(), j, cosines_scores[col_max[j], j].item())
connections.add(connection)
return connections
|