cossim.py

import datasets
import evaluate
from typing import List, Union
import torch

import torch.nn.functional as F
import torch.nn as nn
_DESCRIPTION = """
Cosine similarity between two pairs of embeddings where each embedding represents the semantics of object .

"""

_KWARGS_DESCRIPTION = """
Args:
    predictions (`list` of a list of `int`): a group of embeddings
    references (`list` of `int`): the other group of embeddings paired with the predictions
    
Returns:
    cos_similarity ("float") : average cosine similarity between two pairs of embeddings
    
Examples:

    Example 1-A simple example
        >>> cos_similarity_metrics = evaluate.load("ahnyeonchan/cosine_sim_btw_embeddings_of_same_semantics")
        >>> results = accuracy_metric.compute(references=[[1.0, 1.0], [0.0, 1.0]], predictions=[[1.0, 1.0], [0.0, 1.0]])
        >>> print(results)
        {'cos_similarity': 1.0}
"""

_CITATION = """"""

@evaluate.utils.file_utils.add_start_docstrings(_DESCRIPTION, _KWARGS_DESCRIPTION)
class CosSim(evaluate.Metric):
    def __init__(self, *args, **kwargs):
        super(CosSim, self).__init__(*args, **kwargs)
        self.cossim = nn.CosineSimilarity()
        
    def _info(self):
        return evaluate.MetricInfo(
            description=_DESCRIPTION,
            citation=_CITATION,
            inputs_description=_KWARGS_DESCRIPTION,
            features=datasets.Features(
                {
                    "predictions": datasets.Sequence(datasets.Value("float32")),
                    "references": datasets.Sequence(datasets.Value("float32")),
                }
            ),
            reference_urls=[],
        )

    def _compute(self, predictions: List[List], references: List[List]):
        
        if isinstance(predictions, torch.Tensor):
            predictions = torch.Tensor(predictions)
        elif isinstance(predictions, list):
            predictions = torch.Tensor(predictions)
        else:
            raise NotImplementedError()
        
        if isinstance(references, torch.Tensor):
            references = torch.Tensor(references)
        elif isinstance(references, list):
            references = torch.Tensor(references)
        else:
            raise NotImplementedError()
        
        cosims = self.cossim(predictions, references)
        val = torch.mean(cossim).item()
        
        return {
            "cos_similarity": float(val)
        }