File size: 5,684 Bytes
0b360c1
 
 
 
 
 
 
 
 
 
 
 
 
1167cb4
0b360c1
 
 
1167cb4
 
59083bb
0b360c1
 
1167cb4
 
 
 
 
 
 
 
 
 
 
0b360c1
 
 
 
1167cb4
0b360c1
 
 
 
1167cb4
0b360c1
 
 
 
 
 
1167cb4
0b360c1
 
 
 
1167cb4
 
0b360c1
1167cb4
0b360c1
 
 
 
 
 
 
 
 
 
 
 
 
 
 
59083bb
 
 
 
 
 
 
 
 
 
 
 
 
 
0b360c1
59083bb
0b360c1
59083bb
 
0b360c1
 
 
 
 
 
 
59083bb
0b360c1
59083bb
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
0b360c1
59083bb
1167cb4
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
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
# Copyright 2020 The HuggingFace Datasets Authors and the current dataset script contributor.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""SBERT consime similarity metric."""

import evaluate
import datasets
import torch
import torch.nn as nn
from transformers import AutoTokenizer, BertModel

_CITATION = """\
@article{Reimers2019,
archivePrefix = {arXiv},
arxivId = {1908.10084},
author = {Reimers, Nils and Gurevych, Iryna},
doi = {10.18653/v1/d19-1410},
eprint = {1908.10084},
isbn = {9781950737901},
journal = {EMNLP-IJCNLP 2019 - 2019 Conference on Empirical Methods in Natural Language Processing and 9th International Joint Conference on Natural Language Processing, Proceedings of the Conference},
pages = {3982--3992},
title = {{Sentence-BERT: Sentence embeddings using siamese BERT-networks}},
year = {2019}
}
"""

_DESCRIPTION = """\
Use SBERT to produce embedding and score the similarity by cosine similarity
"""


_KWARGS_DESCRIPTION = """
Calculates how semantic similarity are predictions given some references, using certain scores
Args:
    predictions: list of predictions to score. Each predictions
        should be a string with tokens separated by spaces.
    references: list of reference for each prediction. Each
        reference should be a string with tokens separated by spaces.
Returns:
    score: description of the first score,
Examples:
    Examples should be written in doctest format, and should illustrate how
    to use the function.

    >>> sbert_cosine = evaluate.load("transZ/sbert_cosine")
    >>> results = my_new_module.compute(references=["Nice to meet you"], predictions=["It is my pleasure to meet you"])
    >>> print(results)
    {'score': 0.85}
"""

@evaluate.utils.file_utils.add_start_docstrings(_DESCRIPTION, _KWARGS_DESCRIPTION)
class sbert_cosine(evaluate.Metric):
    """TODO: Short description of my evaluation module."""

    def _info(self):
        # TODO: Specifies the evaluate.EvaluationModuleInfo object
        return evaluate.MetricInfo(
            # This is the description that will appear on the modules page.
            module_type="metric",
            description=_DESCRIPTION,
            citation=_CITATION,
            inputs_description=_KWARGS_DESCRIPTION,
            # This defines the format of each prediction and reference
            features=[
                datasets.Features(
                    {
                        "predictions": datasets.Value("string", id="sequence"),
                        "references": datasets.Sequence(datasets.Value("string", id="sequence"), id="references"),
                    }
                ),
                datasets.Features(
                    {
                        "predictions": datasets.Value("string", id="sequence"),
                        "references": datasets.Value("string", id="sequence"),
                    }
                ),
            ],
            # Homepage of the module for documentation
            homepage="http://sbert.net",
            # Additional links to the codebase or references
            codebase_urls=["https://github.com/UKPLab/sentence-transformers"],
            reference_urls=["https://github.com/UKPLab/sentence-transformers"]
        )

    def _download_and_prepare(self, dl_manager):
        """Optional: download external resources useful to compute the scores"""
        # TODO: Download external resources if needed
        pass

    def _compute(self, predictions, references, model_type='sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2'):
        """Returns the scores"""
        def mean_pooling(model_output, attention_mask):
            token_embeddings = model_output[0] #First element of model_output contains all token embeddings
            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 batch_to_device(batch, target_device: device):
            """
            send a pytorch batch to a device (CPU/GPU)
            """
            for key in batch:
                if isinstance(batch[key], torch.Tensor):
                    batch[key] = batch[key].to(target_device)
            return batch

        device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

        tokenizer = AutoTokenizer.from_pretrained(model_type)
        model = BertModel.from_pretrained(model_type)
        model = model.to(device)
        cosine = nn.CosineSimilarity()

        def calculate(x: str, y: str):
            encoded_input = tokenizer([x, y], padding=True, truncation=True, return_tensors='pt')
            encoded_input = batch_to_device(encode_input, device)
            model_output = model(**encoded_input)
            embeds = mean_pooling(model_output, encoded_input['attention_mask'])
            res = cosine(embeds[0, :], embeds[1, :]).item()
            return res

        with torch.no_grad():
            score = torch.mean([calculate(pred, ref) for pred, ref in zip(predictions, references)]).item()

        return {
            "score": score,
        }