From f97a51ce24b09004ad79fbb3790e5bc4380b588b Mon Sep 17 00:00:00 2001 From: Joe <79627742+ZhouhaoJiang@users.noreply.github.com> Date: Fri, 2 Aug 2024 16:34:50 +0800 Subject: [PATCH] fix: reranking disable timer error (#6910) --- api/core/rag/retrieval/dataset_retrieval.py | 23 ++++++++++++--------- 1 file changed, 13 insertions(+), 10 deletions(-) diff --git a/api/core/rag/retrieval/dataset_retrieval.py b/api/core/rag/retrieval/dataset_retrieval.py index a69fcffbb4..2c5d920a9a 100644 --- a/api/core/rag/retrieval/dataset_retrieval.py +++ b/api/core/rag/retrieval/dataset_retrieval.py @@ -322,23 +322,26 @@ class DatasetRetrieval: for thread in threads: thread.join() - if reranking_enable: - # do rerank for searched documents - data_post_processor = DataPostProcessor(tenant_id, reranking_mode, - reranking_model, weights, False) + with measure_time() as timer: + if reranking_enable: + # do rerank for searched documents + data_post_processor = DataPostProcessor( + tenant_id, reranking_mode, + reranking_model, weights, False + ) - with measure_time() as timer: all_documents = data_post_processor.invoke( query=query, documents=all_documents, score_threshold=score_threshold, top_n=top_k ) - else: - if index_type == "economy": - all_documents = self.calculate_keyword_score(query, all_documents, top_k) - elif index_type == "high_quality": - all_documents = self.calculate_vector_score(all_documents, top_k, score_threshold) + else: + if index_type == "economy": + all_documents = self.calculate_keyword_score(query, all_documents, top_k) + elif index_type == "high_quality": + all_documents = self.calculate_vector_score(all_documents, top_k, score_threshold) + self._on_query(query, dataset_ids, app_id, user_from, user_id) if all_documents: