Fix rerank top k

pyproject_gpu.toml

@@ -103,6 +103,7 @@ aspose-slides = "^24.10.0"
 ultralytics = "8.3.43"
 datasketch = "^1.6.5"
 tablestore = "^6.1.0"
+anyio = "4.6.2.post1"
 
 [tool.poetry.scripts]
 pai_rag = "pai_rag.main:run"

src/pai_rag/app/web/tabs/chat_tab.py

@@ -141,14 +141,36 @@ def create_chat_tab() -> Dict[str, Any]:
                         visible=(retrieval_mode == "Hybrid"),
                     )
 
+                    similarity_top_k = gr.Slider(
+                        minimum=0,
+                        maximum=100,
+                        step=1,
+                        elem_id="similarity_top_k",
+                        label="Text Top K (choose between 0 and 100)",
+                    )
+                    image_similarity_top_k = gr.Slider(
+                        minimum=0,
+                        maximum=10,
+                        step=1,
+                        elem_id="image_similarity_top_k",
+                        label="Image Top K (choose between 0 and 10)",
+                    )
+                    similarity_threshold = gr.Slider(
+                        minimum=0,
+                        maximum=1,
+                        step=0.01,
+                        elem_id="similarity_threshold",
+                        label="Similarity Score Threshold (The more similar the items, the bigger the value.)",
+                    )
+
                     reranker_type = gr.Radio(
                         ["no-reranker", "model-based-reranker"],
                         label="Reranker Type",
                         elem_id="reranker_type",
                     )
-
                     with gr.Column(
-                        visible=(reranker_type == "model-based-reranker")
+                        visible=(reranker_type == "model-based-reranker"),
+                        elem_id="model_reranker_col",
                     ) as model_reranker_col:
                         reranker_model = gr.Radio(
                             [
@@ -165,28 +187,12 @@ def create_chat_tab() -> Dict[str, Any]:
                             elem_id="reranker_similarity_threshold",
                             label="Reranker Similarity Score Threshold (The more similar the items, the bigger the value.)",
                         )
-
-                    with gr.Column():
-                        similarity_top_k = gr.Slider(
+                        reranker_similarity_top_k = gr.Slider(
                             minimum=0,
-                            maximum=100,
+                            maximum=50,
                             step=1,
-                            elem_id="similarity_top_k",
-                            label="Text Top K (choose between 0 and 100)",
-                        )
-                        image_similarity_top_k = gr.Slider(
-                            minimum=0,
-                            maximum=10,
-                            step=1,
-                            elem_id="image_similarity_top_k",
-                            label="Image Top K (choose between 0 and 10)",
-                        )
-                        similarity_threshold = gr.Slider(
-                            minimum=0,
-                            maximum=1,
-                            step=0.01,
-                            elem_id="similarity_threshold",
-                            label="Similarity Score Threshold (The more similar the items, the bigger the value.)",
+                            elem_id="reranker_similarity_top_k",
+                            label="Reranker Text Top K (choose between 0 and 50)",
                         )
 
                     def change_weight(change_weight):
@@ -246,6 +252,7 @@ def change_retrieval_mode(retrieval_mode):
                     similarity_threshold,
                     reranker_similarity_threshold,
                     reranker_model,
+                    reranker_similarity_top_k,
                 }
 
             with gr.Column(visible=True) as llm_col:
@@ -454,11 +461,13 @@ def change_query_radio(query_type):
             keyword_weight.elem_id: keyword_weight,
             similarity_threshold.elem_id: similarity_threshold,
             reranker_similarity_threshold.elem_id: reranker_similarity_threshold,
+            reranker_similarity_top_k.elem_id: reranker_similarity_top_k,
             multimodal_qa_template.elem_id: multimodal_qa_template,
             citation_multimodal_qa_template.elem_id: citation_multimodal_qa_template,
             citation_text_qa_template.elem_id: citation_text_qa_template,
             text_qa_template.elem_id: text_qa_template,
             search_lang.elem_id: search_lang,
             search_api_key.elem_id: search_api_key,
             search_count.elem_id: search_count,
+            model_reranker_col.elem_id: model_reranker_col,
         }

src/pai_rag/app/web/view_model.py

@@ -316,8 +316,6 @@ def to_app_config(self):
         config["retriever"]["vector_weight"] = self.vector_weight
         config["retriever"]["keyword_weight"] = self.keyword_weight
 
-        config["retriever"]["image_similarity_top_k"] = self.image_similarity_top_k
-
         config["retriever"]["search_image"] = self.need_image
         if self.retrieval_mode == "Hybrid":
             config["retriever"]["vector_store_query_mode"] = VectorStoreQueryMode.HYBRID
@@ -386,7 +384,7 @@ def to_app_config(self):
             config["postprocessor"][
                 "similarity_threshold"
             ] = self.reranker_similarity_threshold
-        config["postprocessor"]["top_n"] = 5
+            config["postprocessor"]["top_n"] = self.reranker_similarity_top_k
 
         config["synthesizer"]["use_multimodal_llm"] = self.use_mllm
         config["synthesizer"]["text_qa_template"] = self.text_qa_template
@@ -573,6 +571,12 @@ def to_component_settings(self) -> Dict[str, Dict[str, Any]]:
         settings["reranker_similarity_threshold"] = {
             "value": self.reranker_similarity_threshold
         }
+        settings["reranker_similarity_top_k"] = {
+            "value": self.reranker_similarity_top_k
+        }
+        settings["model_reranker_col"] = {
+            "visible": self.reranker_type == "model-based-reranker"
+        }
 
         settings["text_qa_template"] = {"value": self.text_qa_template}
         settings["multimodal_qa_template"] = {"value": self.multimodal_qa_template}

src/pai_rag/config/settings.toml

@@ -69,8 +69,6 @@ endpoint = "oss-cn-hangzhou.aliyuncs.com"
 [rag.postprocessor]
 reranker_type = "no-reranker" # [simple-weighted-reranker, model-based-reranker]
 reranker_model = "bge-reranker-base" # [bge-reranker-base, bge-reranker-large]
-keyword_weight = 0.3
-vector_weight = 0.7
 similarity_threshold = 0.5
 top_n = 2
 

src/pai_rag/core/rag_config.py

@@ -47,8 +47,6 @@ def validate_case_insensitive(value: Dict) -> Dict:
             # fix old config
             if value[key] == "simple-weighted-reranker":
                 value[key] = "no-reranker"
-            if value[key] == "pai-llm-trace":
-                value[key] = "pai_trace"
     return value