diff --git a/buster/retriever/service.py b/buster/retriever/service.py index 5e1ce5d..362571f 100644 --- a/buster/retriever/service.py +++ b/buster/retriever/service.py @@ -92,4 +92,7 @@ def get_topk_documents(self, query: str, source: str, top_k: int) -> pd.DataFram matched_documents["similarity"] = matched_documents["_id"].apply(lambda x: matching_scores[str(x)]) matched_documents["embedding"] = matched_documents["_id"].apply(lambda x: matching_embeddings[str(x)]) + # sort by similarity + matched_documents = matched_documents.sort_values(by="similarity", ascending=False, ignore_index=True) + return matched_documents