File size: 665 Bytes
027da24 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 |
from sentence_transformers import SentenceTransformer
import faiss
import numpy as np
# Load MiniLM embedder
embedder = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
def embed_texts(texts):
return embedder.encode(texts, convert_to_tensor=False)
def build_faiss_index(texts):
embeddings = embed_texts(texts)
index = faiss.IndexFlatL2(embeddings[0].shape[0])
index.add(np.array(embeddings))
return index, embeddings
def retrieve(query, index, docs, k=3):
query_embedding = embed_texts([query])
distances, indices = index.search(np.array(query_embedding), k)
return [docs[i] for i in indices[0]]
|