|
""" |
|
Source: https://github.com/ZurichNLP/recognizing-semantic-differences |
|
MIT License |
|
Copyright (c) 2023 University of Zurich |
|
""" |
|
|
|
from dataclasses import dataclass |
|
from typing import Tuple, Optional |
|
|
|
import torch |
|
from tokenizers.pre_tokenizers import Whitespace |
|
from torch import Tensor |
|
|
|
|
|
@dataclass |
|
class DifferenceSample: |
|
tokens_a: Tuple[str, ...] |
|
tokens_b: Tuple[str, ...] |
|
labels_a: Tuple[float, ...] |
|
labels_b: Optional[Tuple[float, ...]] |
|
|
|
def add_whitespace(self, encoding_a, encoding_b): |
|
self.tokens_a = self._add_whitespace(self.tokens_a, encoding_a) |
|
self.tokens_b = self._add_whitespace(self.tokens_b, encoding_b) |
|
|
|
def _add_whitespace(self, tokens, encoding) -> Tuple[str, ...]: |
|
assert len(tokens) == len(encoding) |
|
new_tokens = [] |
|
for i in range(len(encoding)): |
|
token = tokens[i] |
|
if i < len(encoding) - 1: |
|
cur_end = encoding[i][1][1] |
|
next_start = encoding[i + 1][1][0] |
|
token += " " * (next_start - cur_end) |
|
new_tokens.append(token) |
|
return tuple(new_tokens) |
|
|
|
|
|
@property |
|
def token_labels_a(self) -> Tuple[Tuple[str, float], ...]: |
|
return tuple(zip(self.tokens_a, self.labels_a)) |
|
|
|
@property |
|
def token_labels_b(self) -> Tuple[Tuple[str, float], ...]: |
|
return tuple(zip(self.tokens_b, self.labels_b)) |
|
|
|
@property |
|
def min(self) -> float: |
|
return min(self.labels_a + self.labels_b) |
|
|
|
@property |
|
def max(self) -> float: |
|
return max(self.labels_a + self.labels_b) |
|
|
|
|
|
def tokenize(text: str) -> Tuple[str]: |
|
""" |
|
Apply Moses-like tokenization to a string. |
|
""" |
|
whitespace_tokenizer = Whitespace() |
|
output = whitespace_tokenizer.pre_tokenize_str(text) |
|
|
|
tokens = [str(token[0]) for token in output] |
|
return tuple(tokens) |
|
|
|
|
|
def cos_sim(a: Tensor, b: Tensor): |
|
""" |
|
Copied from https://github.com/UKPLab/sentence-transformers/blob/d928410803bb90f555926d145ee7ad3bd1373a83/sentence_transformers/util.py#L31 |
|
|
|
Computes the cosine similarity cos_sim(a[i], b[j]) for all i and j. |
|
:return: Matrix with res[i][j] = cos_sim(a[i], b[j]) |
|
""" |
|
if not isinstance(a, torch.Tensor): |
|
a = torch.tensor(a) |
|
|
|
if not isinstance(b, torch.Tensor): |
|
b = torch.tensor(b) |
|
|
|
if len(a.shape) == 1: |
|
a = a.unsqueeze(0) |
|
|
|
if len(b.shape) == 1: |
|
b = b.unsqueeze(0) |
|
|
|
a_norm = torch.nn.functional.normalize(a, p=2, dim=1) |
|
b_norm = torch.nn.functional.normalize(b, p=2, dim=1) |
|
return torch.mm(a_norm, b_norm.transpose(0, 1)) |
|
|
|
|
|
def pairwise_dot_score(a: Tensor, b: Tensor): |
|
""" |
|
Copied from https://github.com/UKPLab/sentence-transformers/blob/d928410803bb90f555926d145ee7ad3bd1373a83/sentence_transformers/util.py#L73 |
|
|
|
Computes the pairwise dot-product dot_prod(a[i], b[i]) |
|
:return: Vector with res[i] = dot_prod(a[i], b[i]) |
|
""" |
|
if not isinstance(a, torch.Tensor): |
|
a = torch.tensor(a) |
|
|
|
if not isinstance(b, torch.Tensor): |
|
b = torch.tensor(b) |
|
|
|
return (a * b).sum(dim=-1) |
|
|
|
|
|
def normalize_embeddings(embeddings: Tensor): |
|
""" |
|
Copied from https://github.com/UKPLab/sentence-transformers/blob/d928410803bb90f555926d145ee7ad3bd1373a83/sentence_transformers/util.py#L101 |
|
|
|
Normalizes the embeddings matrix, so that each sentence embedding has unit length |
|
""" |
|
return torch.nn.functional.normalize(embeddings, p=2, dim=1) |
|
|
|
|
|
def pairwise_cos_sim(a: Tensor, b: Tensor): |
|
""" |
|
Copied from https://github.com/UKPLab/sentence-transformers/blob/d928410803bb90f555926d145ee7ad3bd1373a83/sentence_transformers/util.py#L87 |
|
|
|
Computes the pairwise cossim cos_sim(a[i], b[i]) |
|
:return: Vector with res[i] = cos_sim(a[i], b[i]) |
|
""" |
|
if not isinstance(a, torch.Tensor): |
|
a = torch.tensor(a) |
|
|
|
if not isinstance(b, torch.Tensor): |
|
b = torch.tensor(b) |
|
|
|
return pairwise_dot_score(normalize_embeddings(a), normalize_embeddings(b)) |
|
|
