|
from sentence_transformers import util, SentenceTransformer |
|
import unittest |
|
import numpy as np |
|
import sklearn |
|
import torch |
|
|
|
class UtilTest(unittest.TestCase): |
|
|
|
def test_normalize_embeddings(self): |
|
"""Tests the correct computation of util.normalize_embeddings""" |
|
embedding_size = 100 |
|
a = torch.tensor(np.random.randn(50, embedding_size)) |
|
a_norm = util.normalize_embeddings(a) |
|
|
|
for embedding in a_norm: |
|
assert len(embedding) == embedding_size |
|
emb_norm = torch.norm(embedding) |
|
assert abs(emb_norm.item() - 1) < 0.0001 |
|
|
|
|
|
def test_pytorch_cos_sim(self): |
|
"""Tests the correct computation of util.pytorch_cos_scores""" |
|
a = np.random.randn(50, 100) |
|
b = np.random.randn(50, 100) |
|
|
|
sklearn_pairwise = sklearn.metrics.pairwise.cosine_similarity(a, b) |
|
pytorch_cos_scores = util.pytorch_cos_sim(a, b).numpy() |
|
for i in range(len(sklearn_pairwise)): |
|
for j in range(len(sklearn_pairwise[i])): |
|
assert abs(sklearn_pairwise[i][j] - pytorch_cos_scores[i][j]) < 0.001 |
|
|
|
|
|
def test_semantic_search(self): |
|
"""Tests util.semantic_search function""" |
|
num_queries = 20 |
|
num_k = 10 |
|
|
|
doc_emb = torch.tensor(np.random.randn(1000, 100)) |
|
q_emb = torch.tensor(np.random.randn(num_queries, 100)) |
|
hits = util.semantic_search(q_emb, doc_emb, top_k=num_k, query_chunk_size=5, corpus_chunk_size=17) |
|
assert len(hits) == num_queries |
|
assert len(hits[0]) == num_k |
|
|
|
|
|
cos_scores = util.pytorch_cos_sim(q_emb, doc_emb) |
|
cos_scores_values, cos_scores_idx = cos_scores.topk(num_k) |
|
cos_scores_values = cos_scores_values.cpu().tolist() |
|
cos_scores_idx = cos_scores_idx.cpu().tolist() |
|
|
|
for qid in range(num_queries): |
|
for hit_num in range(num_k): |
|
assert hits[qid][hit_num]['corpus_id'] == cos_scores_idx[qid][hit_num] |
|
assert np.abs(hits[qid][hit_num]['score'] - cos_scores_values[qid][hit_num]) < 0.001 |
|
|
|
def test_paraphrase_mining(self): |
|
model = SentenceTransformer('paraphrase-distilroberta-base-v1') |
|
sentences = [ |
|
"This is a test", "This is a test!", |
|
"The cat sits on mat", "The cat sits on the mat", "On the mat a cat sits", |
|
"A man eats pasta", "A woman eats pasta", "A man eats spaghetti" |
|
] |
|
duplicates = util.paraphrase_mining(model, sentences) |
|
|
|
for score, a, b in duplicates: |
|
if score > 0.5: |
|
assert (a,b) in [(0,1), (2,3), (2,4), (3,4), (5,6), (5,7), (6,7)] |
|
|
|
|
|
def test_pairwise_scores(self): |
|
a = np.random.randn(50, 100) |
|
b = np.random.randn(50, 100) |
|
|
|
|
|
sklearn_pairwise = 1-sklearn.metrics.pairwise.paired_cosine_distances(a, b) |
|
pytorch_cos_scores = util.pairwise_cos_sim(a, b).numpy() |
|
assert np.allclose(sklearn_pairwise, pytorch_cos_scores) |
|
|
|
if "__main__" == __name__: |
|
unittest.main() |
