pls-rag / service /search.py
m97j's picture
Initial codes commit
4fdc679
raw
history blame
382 Bytes
from modules.embedder import get_embedding
from modules.retriever import retrieve_ids
from modules.corpus import fetch_contexts_by_ids
from modules.reranker import rerank
def search(query: str) -> list[dict]:
embedding = get_embedding(query)
ids = retrieve_ids(embedding)
contexts = fetch_contexts_by_ids(ids)
reranked = rerank(query, contexts)
return reranked