|
""" |
|
Computes embeddings |
|
""" |
|
|
|
|
|
import unittest |
|
from sentence_transformers import SentenceTransformer |
|
import numpy as np |
|
|
|
class ComputeEmbeddingsTest(unittest.TestCase): |
|
def setUp(self): |
|
self.model = SentenceTransformer('paraphrase-distilroberta-base-v1') |
|
|
|
|
|
def test_encode_token_embeddings(self): |
|
""" |
|
Test that encode(output_value='token_embeddings') works |
|
:return: |
|
""" |
|
sent = ["Hello Word, a test sentence", "Here comes another sentence", "My final sentence", "Sentences", "Sentence five five five five five five five"] |
|
emb = self.model.encode(sent, output_value='token_embeddings', batch_size=2) |
|
assert len(emb) == len(sent) |
|
for s, e in zip(sent, emb): |
|
assert len(self.model.tokenize([s])['input_ids'][0]) == e.shape[0] |
|
|
|
|
|
def test_encode_single_sentences(self): |
|
|
|
emb = self.model.encode("Hello Word, a test sentence") |
|
assert emb.shape == (768,) |
|
assert abs(np.sum(emb) - 7.9811716) < 0.001 |
|
|
|
|
|
emb = self.model.encode(["Hello Word, a test sentence"]) |
|
assert emb.shape == (1, 768) |
|
assert abs(np.sum(emb) - 7.9811716) < 0.001 |
|
|
|
|
|
emb = self.model.encode(["Hello Word, a test sentence", "Here comes another sentence", "My final sentence"]) |
|
assert emb.shape == (3, 768) |
|
assert abs(np.sum(emb) - 22.968266) < 0.001 |
|
|
|
def test_encode_normalize(self): |
|
emb = self.model.encode(["Hello Word, a test sentence", "Here comes another sentence", "My final sentence"], normalize_embeddings=True) |
|
assert emb.shape == (3, 768) |
|
for norm in np.linalg.norm(emb, axis=1): |
|
assert abs(norm - 1) < 0.001 |
|
|
|
def test_encode_tuple_sentences(self): |
|
|
|
emb = self.model.encode([("Hello Word, a test sentence", "Second input for model")]) |
|
assert emb.shape == (1, 768) |
|
assert abs(np.sum(emb) - 9.503508) < 0.001 |
|
|
|
|
|
emb = self.model.encode([("Hello Word, a test sentence", "Second input for model"), ("My second tuple", "With two inputs"), ("Final tuple", "final test")]) |
|
assert emb.shape == (3, 768) |
|
assert abs(np.sum(emb) - 32.14627) < 0.001 |
|
|
|
|
|
|
|
|
|
|
|
|
