sentence_transformers_embeddings_with_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 to encode
7sentences = [
8    "The cat sits outside",
9    "A man is playing guitar",
10    "The movies are awesome"
11]
12
13# 2. Calculate embeddings by calling model.encode()
14embeddings = model.encode(sentences)
15
16# 3. Calculate the embedding similarities
17# (Optional: If you want to compute similarities between sentences)
18sentences2 = [
19    "The new movie is awesome",
20    "The cat plays in the garden",
21    "A woman is watching TV"
22]
23embeddings2 = model.encode(sentences2)
24
25# Compute cosine similarity between all pairs
26cos_sim = util.cos_sim(embeddings, embeddings2)
27
28# Print the results
29for i in range(len(sentences)):
30    for j in range(len(sentences2)):
31        print(f"{sentences[i]} \t\t {sentences2[j]} \t\t Score: {cos_sim[i][j]:.4f}")