vector-search.py with requirements.txt

requirements.txt

sentence-transformers
torch
huggingface_hub

vector-search.py

from sentence_transformers import SentenceTransformer

print("\nLoading embedding model...")
model = SentenceTransformer("google/embeddinggemma-300m", device="mps")

print("\nGenerating vectors...")
documents = ["cymbals", "guitar"]
vectors = model.encode(documents, prompt_name="document", normalize_embeddings=True)


def search():
    query = input("\nEnter a query: ")
    query_vector = model.encode([query], prompt_name="query", normalize_embeddings=True)[0]
    similarity_scores = (vectors @ query_vector.T).reshape(-1)
    for doc, score in zip(documents, similarity_scores):
        print(f'\n"{query}" vs "{doc}": {score:.4f}')

while True:
    search()