File size: 574 Bytes
bb0b469
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
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]]