sentence_transformers_encode_sentences_cosine_similarity.py

1from sentence_transformers import SentenceTransformer, util
2
3# 1. Load a pretrained Sentence Transformer model
4model = SentenceTransformer("all-MiniLM-L6-v2")
5
6# The sentences we wish to encode
7sentences = [
8    "This framework generates embeddings for each input sentence",
9    "Sentences are passed as a list of strings.",
10    "The quick brown fox jumps over the lazy dog.",
11]
12
13# 2. Calculate embeddings by calling model.encode()
14embeddings = model.encode(sentences)
15
16# Print the embeddings
17for sentence, embedding in zip(sentences, embeddings):
18    print("Sentence:", sentence)
19    print("Embedding:", embedding)
20    print("")
21
22# 3. Calculate the cosine similarity between all pairs of sentences
23cosine_scores = util.cos_sim(embeddings, embeddings)
24
25# Output the pairs with their score
26for i in range(len(sentences)):
27    for j in range(i + 1, len(sentences)):
28        print(f"{sentences[i]} \n{sentences[j]} \nScore: {cosine_scores[i][j]:.4f}\n")