File size: 377 Bytes
71d239c | 1 2 3 4 5 6 7 8 9 10 11 12 | """Find the document chunks most relevant to a question."""
import config
from embed import embed_text
from ingest import collection
def retrieve(question):
"""Return the top-K most relevant documents for the question."""
q_vec = embed_text(question)
results = collection.query(query_embeddings=[q_vec], n_results=config.TOP_K)
return results["documents"][0]
|