| | """ |
| | 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 |
| |
|
| |
|
| |
|
| |
|
| |
|
| |
|
