source/embedding.py

from sentence_transformers import SentenceTransformer, util

# Initialize the model
model = SentenceTransformer("sentence-transformers/all-mpnet-base-v2")

def text_to_embedding(text):
    """Convert text to embedding using the all-MiniLM-L6-v2 model.
    Assumes the input is a single string."""
    return model.encode([text])[0]

def compare_embeddings(embedding1, embedding2):
    """Compare two embeddings using cosine similarity.
    
    This function computes the cosine similarity between two embeddings and returns the score.
    """
    # Ensure the embeddings are in the correct shape for comparison
    # util.cos_sim expects 2D arrays
    embedding1 = embedding1.reshape(1, -1)
    embedding2 = embedding2.reshape(1, -1)
    
    # Compute and return cosine similarity
    return util.cos_sim(embedding1, embedding2).item()