| """ |
| Tests that the pretrained models produce the correct scores on the STSbenchmark dataset |
| """ |
| from sentence_transformers import SentenceTransformer, InputExample, util |
| from sentence_transformers.evaluation import EmbeddingSimilarityEvaluator |
| import unittest |
| import os |
| import gzip |
| import csv |
|
|
| class PretrainedSTSbTest(unittest.TestCase): |
|
|
| def pretrained_model_score(self, model_name, expected_score): |
| model = SentenceTransformer(model_name) |
| sts_dataset_path = 'datasets/stsbenchmark.tsv.gz' |
|
|
| if not os.path.exists(sts_dataset_path): |
| util.http_get('https://sbert.net/datasets/stsbenchmark.tsv.gz', sts_dataset_path) |
|
|
| train_samples = [] |
| dev_samples = [] |
| test_samples = [] |
| with gzip.open(sts_dataset_path, 'rt', encoding='utf8') as fIn: |
| reader = csv.DictReader(fIn, delimiter='\t', quoting=csv.QUOTE_NONE) |
| for row in reader: |
| score = float(row['score']) / 5.0 |
| inp_example = InputExample(texts=[row['sentence1'], row['sentence2']], label=score) |
|
|
| if row['split'] == 'dev': |
| dev_samples.append(inp_example) |
| elif row['split'] == 'test': |
| test_samples.append(inp_example) |
| else: |
| train_samples.append(inp_example) |
|
|
| evaluator = EmbeddingSimilarityEvaluator.from_input_examples(test_samples, name='sts-test') |
|
|
| score = model.evaluate(evaluator)*100 |
| print(model_name, "{:.2f} vs. exp: {:.2f}".format(score, expected_score)) |
| assert score > expected_score or abs(score-expected_score) < 0.1 |
|
|
| def test_bert_base(self): |
| self.pretrained_model_score('bert-base-nli-mean-tokens', 77.12) |
| self.pretrained_model_score('bert-base-nli-max-tokens', 77.21) |
| self.pretrained_model_score('bert-base-nli-cls-token', 76.30) |
| self.pretrained_model_score('bert-base-nli-stsb-mean-tokens', 85.14) |
|
|
|
|
| def test_bert_large(self): |
| self.pretrained_model_score('bert-large-nli-mean-tokens', 79.19) |
| self.pretrained_model_score('bert-large-nli-max-tokens', 78.41) |
| self.pretrained_model_score('bert-large-nli-cls-token', 78.29) |
| self.pretrained_model_score('bert-large-nli-stsb-mean-tokens', 85.29) |
|
|
| def test_roberta(self): |
| self.pretrained_model_score('roberta-base-nli-mean-tokens', 77.49) |
| self.pretrained_model_score('roberta-large-nli-mean-tokens', 78.69) |
| self.pretrained_model_score('roberta-base-nli-stsb-mean-tokens', 85.30) |
| self.pretrained_model_score('roberta-large-nli-stsb-mean-tokens', 86.39) |
|
|
| def test_distilbert(self): |
| self.pretrained_model_score('distilbert-base-nli-mean-tokens', 78.69) |
| self.pretrained_model_score('distilbert-base-nli-stsb-mean-tokens', 85.16) |
| self.pretrained_model_score('paraphrase-distilroberta-base-v1', 81.81) |
|
|
| def test_multiling(self): |
| self.pretrained_model_score('distiluse-base-multilingual-cased', 80.75) |
| self.pretrained_model_score('paraphrase-xlm-r-multilingual-v1', 83.50) |
| self.pretrained_model_score('paraphrase-multilingual-MiniLM-L12-v2', 84.42) |
|
|
| def test_mpnet(self): |
| self.pretrained_model_score('paraphrase-mpnet-base-v2', 86.99) |
|
|
| def test_other_models(self): |
| self.pretrained_model_score('average_word_embeddings_komninos', 61.56) |
|
|
| def test_msmarco(self): |
| self.pretrained_model_score('msmarco-roberta-base-ance-firstp', 77.0) |
| self.pretrained_model_score('msmarco-distilbert-base-v3', 78.85) |
|
|
| def test_sentence_t5(self): |
| self.pretrained_model_score('sentence-t5-base', 85.52) |
|
|
| if "__main__" == __name__: |
| unittest.main() |
