Politometro / search_engine.py
MaNmAxImO's picture
Create search_engine.py
f1dd3d0 verified
raw
history blame contribute delete
339 Bytes
import numpy as np
from sentence_transformers import SentenceTransformer
def search_relevant_chunks(query, index, documents, model, top_k=5):
query_embedding = model.encode(query, convert_to_tensor=True)
distances, indices = index.search(np.array([query_embedding]), top_k)
return [documents[i]["content"] for i in indices[0]]