GodsDevProject's picture
Create semantic.py
ff39e36 verified
raw
history blame contribute delete
574 Bytes
import faiss
import numpy as np
from sentence_transformers import SentenceTransformer
from typing import List, Dict
model = SentenceTransformer("all-MiniLM-L6-v2")
def build_faiss_index(docs: List[Dict]):
texts = [d["content"] for d in docs]
embeddings = model.encode(texts)
index = faiss.IndexFlatL2(embeddings.shape[1])
index.add(np.array(embeddings))
return index, embeddings
def semantic_search(query: str, docs: List[Dict], index):
q_emb = model.encode([query])
D, I = index.search(np.array(q_emb), k=5)
return [docs[i] for i in I[0]]