Spaces:
Runtime error
Runtime error
File size: 596 Bytes
06696b5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 |
from retriever.reranker import rerank_documents as rerank_fn
from type.state_schema import RAGState # Optional: for type hint
def rerank_documents(state: RAGState) -> RAGState:
query = state.query
law = state.retrieved_docs["law"]
exam = state.retrieved_docs["exam"]
law_reranked = rerank_fn(query, [doc for doc, _, _ in law], top_k=5)
exam_reranked = rerank_fn(query, [doc for doc, _, _ in exam], top_k=5)
return state.model_copy(update={
"reranked_docs": {
"law": law_reranked,
"exam": exam_reranked
}
})
|