sửa rerank

core/embeddings/retrival.py

@@ -34,12 +34,13 @@ class RetrievalConfig:
     rerank_api_base_url: str = "https://api.siliconflow.com/v1"
     rerank_model: str = "Qwen/Qwen3-Reranker-4B"
     rerank_top_n: int = 10
-    initial_k: int = 100
+    initial_k: int = 25  # Reduced to minimize reranker time
     top_k: int = 5
     vector_weight: float = 0.5
     bm25_weight: float = 0.5
 
 
+
 _retrieval_config: RetrievalConfig | None = None
 
 
@@ -82,7 +83,7 @@ class SiliconFlowReranker(BaseDocumentCompressor):
                         "documents": [doc.page_content for doc in documents],
                         "top_n": self.top_n or len(documents),
                     },
-                    timeout=30,
+                    timeout=120,
                 )
                 response.raise_for_status()
                 data = response.json()
@@ -109,11 +110,10 @@ class SiliconFlowReranker(BaseDocumentCompressor):
         return list(documents)
 
 
+
+
 class Retriever:
     def __init__(self, vector_db: "ChromaVectorDB", use_reranker: bool = True):
-        import time
-        start = time.time()
-        
         self._vector_db = vector_db
         self._config = get_retrieval_config()
         self._reranker: Optional[SiliconFlowReranker] = None
@@ -138,7 +138,8 @@ class Retriever:
         if use_reranker:
             self._reranker = self._init_reranker()
         
-        logger.info(f"Retriever initialized in {time.time() - start:.2f}s (BM25 lazy-loaded)")
+        logger.info("Retriever initialized")
+
     
     def _save_bm25_cache(self, bm25: BM25Retriever) -> None:
         """Save BM25 retriever to disk for fast loading."""
@@ -153,9 +154,9 @@ class Retriever:
             logger.warning(f"Failed to save BM25 cache: {e}")
     
     def _load_bm25_cache(self) -> Optional[BM25Retriever]:
-        """Load BM25 retriever from disk cache."""
         if not self._bm25_cache_path or not self._bm25_cache_path.exists():
             return None
+            
         try:
             import pickle
             import time
@@ -168,9 +169,10 @@ class Retriever:
         except Exception as e:
             logger.warning(f"Failed to load BM25 cache: {e}")
             return None
+
+
     
     def _init_bm25(self) -> Optional[BM25Retriever]:
-        """Initialize BM25 retriever (lazy-loaded, with disk cache)."""
         if self._bm25_initialized:
             return self._bm25_retriever
         
@@ -330,6 +332,8 @@ class Retriever:
         where: Optional[Dict[str, Any]] = None,
         initial_k: int | None = None,
     ) -> List[Dict[str, Any]]:
+        import time
+        
         if not text.strip():
             return []
         
@@ -353,12 +357,20 @@ class Retriever:
             if bm25:
                 bm25.k = initial_k
         
-        final_retriever = self._build_final()
-        results = final_retriever.invoke(text)
+        ensemble = self._get_ensemble_retriever()
+        ensemble_results = ensemble.invoke(text)
+        
+        if self._reranker:
+            results = self._reranker.compress_documents(ensemble_results, text)
+        else:
+            results = ensemble_results
+        
         return [
             self._to_result(doc, i + 1, rerank_score=doc.metadata.get("rerank_score"))
             for i, doc in enumerate(results[:k])
         ]
+
+
     
     def flexible_search(
         self,

core/gradio/gradio_rag_qwen.py

@@ -30,7 +30,7 @@ from core.embeddings.generator import RAGContextBuilder, build_context, build_pr
 
 _load_env()
 
-RETRIEVAL_MODE = RetrievalMode.VECTOR_ONLY  # Fastest mode - no BM25/reranker
+RETRIEVAL_MODE = RetrievalMode.HYBRID_RERANK  # Test with debug logs
 
 # LLM Config (hardcoded sau khi xóa LLMConfig từ generator)
 LLM_MODEL = os.getenv("LLM_MODEL", "qwen/qwen3-32b")
@@ -83,12 +83,7 @@ def _init_resources() -> None:
 
 
 def rag_chat(message: str, history: List[Dict[str, str]] | None = None):
-    import time
-    total_start = time.time()
-    
-    init_start = time.time()
     _init_resources()
-    init_time = time.time() - init_start
 
     assert STATE.db is not None
     assert STATE.client is not None
@@ -96,14 +91,12 @@ def rag_chat(message: str, history: List[Dict[str, str]] | None = None):
     assert STATE.rag_builder is not None
 
     # Bước 1: Retrieve và prepare context
-    retrieval_start = time.time()
     prepared = STATE.rag_builder.retrieve_and_prepare(
         message,
         k=RETRIEVAL_CFG.top_k, 
         initial_k=RETRIEVAL_CFG.initial_k,
         mode=RETRIEVAL_MODE.value,
     )
-    retrieval_time = time.time() - retrieval_start
     results = prepared["results"]
 
     if not results:
@@ -111,7 +104,6 @@ def rag_chat(message: str, history: List[Dict[str, str]] | None = None):
         return
 
     # Bước 2: Gọi LLM streaming để generate answer
-    llm_start = time.time()
     completion = STATE.client.chat.completions.create(
         model=LLM_MODEL,
         messages=[{"role": "user", "content": prepared["prompt"]}],
@@ -121,26 +113,11 @@ def rag_chat(message: str, history: List[Dict[str, str]] | None = None):
     )
     
     acc = ""
-    first_token_time = None
     for chunk in completion:
-        if first_token_time is None:
-            first_token_time = time.time() - llm_start
         delta = getattr(chunk.choices[0].delta, "content", "") or ""
         if delta:
             acc += delta
             yield acc
-    
-    llm_time = time.time() - llm_start
-    total_time = time.time() - total_start
-    
-    # Timing info
-    timing_info = f"\n\n---\n**⏱️ Timing:**\n"
-    timing_info += f"- Init: {init_time:.2f}s\n"
-    timing_info += f"- Retrieval: {retrieval_time:.2f}s\n"
-    timing_info += f"- LLM (first token): {first_token_time:.2f}s\n" if first_token_time else ""
-    timing_info += f"- LLM (total): {llm_time:.2f}s\n"
-    timing_info += f"- **Total: {total_time:.2f}s**\n"
-
 
     # Debug info with mode indicator
     debug_info = f"\n\n---\n\n**Retrieved (Top {len(results)} | Mode: {RETRIEVAL_MODE.value})**\n\n"
@@ -182,7 +159,8 @@ def rag_chat(message: str, history: List[Dict[str, str]] | None = None):
             debug_info += f"   - **Mục:** {header[:80]}{'...' if len(header) > 80 else ''}\n"
         debug_info += f"   - **Content:** {content[:200]}{'...' if len(content) > 200 else ''}\n\n"
 
-    yield acc + timing_info + debug_info
+    yield acc + debug_info
+