BrandSight / src /utils.py
akashjayampu's picture
Update src/utils.py
9e62dbe verified
import faiss
import numpy as np
from sentence_transformers import SentenceTransformer
model = SentenceTransformer("all-MiniLM-L6-v2")
def get_embeddings(texts):
return model.encode(texts, convert_to_numpy=True)
def build_index(embeddings):
dim = embeddings.shape[1]
index = faiss.IndexFlatIP(dim)
faiss.normalize_L2(embeddings)
index.add(embeddings)
return index
def search_similar(index, query_embedding, k=5):
faiss.normalize_L2(query_embedding)
distances, indices = index.search(query_embedding, k)
return distances, indices