from __future__ import annotations import csv import logging import os from contextlib import nullcontext from typing import TYPE_CHECKING, Literal import numpy as np from scipy.stats import pearsonr, spearmanr from sklearn.metrics.pairwise import paired_cosine_distances, paired_euclidean_distances, paired_manhattan_distances from sentence_transformers.evaluation.SentenceEvaluator import SentenceEvaluator from sentence_transformers.readers import InputExample from sentence_transformers.similarity_functions import SimilarityFunction if TYPE_CHECKING: from sentence_transformers.SentenceTransformer import SentenceTransformer logger = logging.getLogger(__name__) class EmbeddingSimilarityEvaluator(SentenceEvaluator): """ Evaluate a model based on the similarity of the embeddings by calculating the Spearman and Pearson rank correlation in comparison to the gold standard labels. The metrics are the cosine similarity as well as euclidean and Manhattan distance The returned score is the Spearman correlation with a specified metric. Example: :: from datasets import load_dataset from sentence_transformers import SentenceTransformer from sentence_transformers.evaluation import EmbeddingSimilarityEvaluator, SimilarityFunction # Load a model model = SentenceTransformer('all-mpnet-base-v2') # Load the STSB dataset (https://huggingface.co/datasets/sentence-transformers/stsb) eval_dataset = load_dataset("sentence-transformers/stsb", split="validation") # Initialize the evaluator dev_evaluator = EmbeddingSimilarityEvaluator( sentences1=eval_dataset["sentence1"], sentences2=eval_dataset["sentence2"], scores=eval_dataset["score"], name="sts_dev", ) results = dev_evaluator(model) ''' EmbeddingSimilarityEvaluator: Evaluating the model on the sts-dev dataset: Cosine-Similarity : Pearson: 0.8806 Spearman: 0.8810 ''' print(dev_evaluator.primary_metric) # => "sts_dev_pearson_cosine" print(results[dev_evaluator.primary_metric]) # => 0.881019449484294 """ def __init__( self, sentences1: list[str], sentences2: list[str], scores: list[float], batch_size: int = 16, main_similarity: str | SimilarityFunction | None = None, similarity_fn_names: list[Literal["cosine", "euclidean", "manhattan", "dot"]] | None = None, name: str = "", show_progress_bar: bool = False, write_csv: bool = True, precision: Literal["float32", "int8", "uint8", "binary", "ubinary"] | None = None, truncate_dim: int | None = None, ): """ Constructs an evaluator based for the dataset. Args: sentences1 (List[str]): List with the first sentence in a pair. sentences2 (List[str]): List with the second sentence in a pair. scores (List[float]): Similarity score between sentences1[i] and sentences2[i]. batch_size (int, optional): The batch size for processing the sentences. Defaults to 16. main_similarity (Optional[Union[str, SimilarityFunction]], optional): The main similarity function to use. Can be a string (e.g. "cosine", "dot") or a SimilarityFunction object. Defaults to None. similarity_fn_names (List[str], optional): List of similarity function names to use. If None, the ``similarity_fn_name`` attribute of the model is used. Defaults to None. name (str, optional): The name of the evaluator. Defaults to "". show_progress_bar (bool, optional): Whether to show a progress bar during evaluation. Defaults to False. write_csv (bool, optional): Whether to write the evaluation results to a CSV file. Defaults to True. precision (Optional[Literal["float32", "int8", "uint8", "binary", "ubinary"]], optional): The precision to use for the embeddings. Can be "float32", "int8", "uint8", "binary", or "ubinary". Defaults to None. truncate_dim (Optional[int], optional): The dimension to truncate sentence embeddings to. `None` uses the model's current truncation dimension. Defaults to None. """ super().__init__() self.sentences1 = sentences1 self.sentences2 = sentences2 self.scores = scores self.write_csv = write_csv self.precision = precision self.truncate_dim = truncate_dim assert len(self.sentences1) == len(self.sentences2) assert len(self.sentences1) == len(self.scores) self.main_similarity = SimilarityFunction(main_similarity) if main_similarity else None self.similarity_fn_names = similarity_fn_names or [] self.name = name self.batch_size = batch_size if show_progress_bar is None: show_progress_bar = ( logger.getEffectiveLevel() == logging.INFO or logger.getEffectiveLevel() == logging.DEBUG ) self.show_progress_bar = show_progress_bar self.csv_file = ( "similarity_evaluation" + ("_" + name if name else "") + ("_" + precision if precision else "") + "_results.csv" ) self.csv_headers = [ "epoch", "steps", ] self._append_csv_headers(self.similarity_fn_names) def _append_csv_headers(self, similarity_fn_names: list[str]) -> None: metrics = ["pearson", "spearman"] for v in similarity_fn_names: for m in metrics: self.csv_headers.append(f"{v}_{m}") @classmethod def from_input_examples(cls, examples: list[InputExample], **kwargs): sentences1 = [] sentences2 = [] scores = [] for example in examples: sentences1.append(example.texts[0]) sentences2.append(example.texts[1]) scores.append(example.label) return cls(sentences1, sentences2, scores, **kwargs) def __call__( self, model: SentenceTransformer, output_path: str = None, epoch: int = -1, steps: int = -1 ) -> dict[str, float]: if epoch != -1: if steps == -1: out_txt = f" after epoch {epoch}" else: out_txt = f" in epoch {epoch} after {steps} steps" else: out_txt = "" if self.truncate_dim is not None: out_txt += f" (truncated to {self.truncate_dim})" logger.info(f"EmbeddingSimilarityEvaluator: Evaluating the model on the {self.name} dataset{out_txt}:") with nullcontext() if self.truncate_dim is None else model.truncate_sentence_embeddings(self.truncate_dim): embeddings1 = model.encode( self.sentences1, batch_size=self.batch_size, show_progress_bar=self.show_progress_bar, convert_to_numpy=True, precision=self.precision, normalize_embeddings=bool(self.precision), ) embeddings2 = model.encode( self.sentences2, batch_size=self.batch_size, show_progress_bar=self.show_progress_bar, convert_to_numpy=True, precision=self.precision, normalize_embeddings=bool(self.precision), ) # Binary and ubinary embeddings are packed, so we need to unpack them for the distance metrics if self.precision == "binary": embeddings1 = (embeddings1 + 128).astype(np.uint8) embeddings2 = (embeddings2 + 128).astype(np.uint8) if self.precision in ("ubinary", "binary"): embeddings1 = np.unpackbits(embeddings1, axis=1) embeddings2 = np.unpackbits(embeddings2, axis=1) labels = self.scores if not self.similarity_fn_names: self.similarity_fn_names = [model.similarity_fn_name] self._append_csv_headers(self.similarity_fn_names) similarity_functions = { "cosine": lambda x, y: 1 - paired_cosine_distances(x, y), "manhattan": lambda x, y: -paired_manhattan_distances(x, y), "euclidean": lambda x, y: -paired_euclidean_distances(x, y), "dot": lambda x, y: [np.dot(emb1, emb2) for emb1, emb2 in zip(x, y)], } metrics = {} for fn_name in self.similarity_fn_names: if fn_name in similarity_functions: scores = similarity_functions[fn_name](embeddings1, embeddings2) eval_pearson, _ = pearsonr(labels, scores) eval_spearman, _ = spearmanr(labels, scores) metrics[f"pearson_{fn_name}"] = eval_pearson metrics[f"spearman_{fn_name}"] = eval_spearman logger.info( f"{fn_name.capitalize()}-Similarity :\tPearson: {eval_pearson:.4f}\tSpearman: {eval_spearman:.4f}" ) if output_path is not None and self.write_csv: csv_path = os.path.join(output_path, self.csv_file) output_file_exists = os.path.isfile(csv_path) with open(csv_path, newline="", mode="a" if output_file_exists else "w", encoding="utf-8") as f: writer = csv.writer(f) if not output_file_exists: writer.writerow(self.csv_headers) writer.writerow( [ epoch, steps, ] + [ metrics[f"{fn_name}_{m}"] for fn_name in self.similarity_fn_names for m in ["pearson", "spearman"] ] ) if len(self.similarity_fn_names) > 1: metrics["pearson_max"] = max(metrics[f"pearson_{fn_name}"] for fn_name in self.similarity_fn_names) metrics["spearman_max"] = max(metrics[f"spearman_{fn_name}"] for fn_name in self.similarity_fn_names) if self.main_similarity: self.primary_metric = { SimilarityFunction.COSINE: "spearman_cosine", SimilarityFunction.EUCLIDEAN: "spearman_euclidean", SimilarityFunction.MANHATTAN: "spearman_manhattan", SimilarityFunction.DOT_PRODUCT: "spearman_dot", }.get(self.main_similarity) else: if len(self.similarity_fn_names) > 1: self.primary_metric = "spearman_max" else: self.primary_metric = f"spearman_{self.similarity_fn_names[0]}" metrics = self.prefix_name_to_metrics(metrics, self.name) self.store_metrics_in_model_card_data(model, metrics) return metrics @property def description(self) -> str: return "Semantic Similarity"