fix bug

rag_pipeline.py

@@ -83,14 +83,16 @@ def initialize_components(data_path):
 def generate_response(query: str, components: dict) -> str:
     """
     Tạo câu trả lời (single-turn) bằng cách sử dụng các thành phần đã được khởi tạo.
-    Phiên bản này được sửa đổi để tương thích với các mô hình Vision (đa phương thức)
-    bằng cách sử dụng chat template.
+    Phiên bản cuối cùng:
+    - Tương thích với mô hình Vision bằng cách sử dụng chat template.
+    - Nhận và sử dụng thông tin `matched_vehicle` từ retriever.
+    - Định dạng context với tóm tắt thông minh từ metadata.
     """
     print("--- Bắt đầu quy trình RAG cho query mới ---")
     
-    # --- Bước 1: Truy xuất Ngữ cảnh (Không thay đổi) ---
-    # (Hàm này giả định bạn đã có phiên bản retriever.py đã sửa lỗi logic boosting)
-    retrieved_results = search_relevant_laws(
+    # === THAY ĐỔI 1: Nhận cả 2 giá trị trả về từ retriever ===
+    # 1. Truy xuất ngữ cảnh bằng retriever đã được nâng cấp
+    retrieved_results, matched_vehicle = search_relevant_laws(
         query_text=query,
         embedding_model=components["embedding_model"],
         faiss_index=components["faiss_index"],
@@ -100,7 +102,7 @@ def generate_response(query: str, components: dict) -> str:
         initial_k_multiplier=15
     )
 
-    # --- Bước 2: Định dạng Ngữ cảnh (Không thay đổi, nhưng nên tách ra hàm riêng) ---
+    # 2. Định dạng Context một cách thông minh
     if not retrieved_results:
         context = "Không tìm thấy thông tin luật liên quan trong cơ sở dữ liệu."
     else:
@@ -109,19 +111,30 @@ def generate_response(query: str, components: dict) -> str:
             metadata = res.get('metadata', {})
             header = f"Trích dẫn {i+1}: Điều {metadata.get('article', 'N/A')}, Khoản {metadata.get('clause_number', 'N/A')} (Nguồn: {metadata.get('source_document', 'N/A')})"
             text = res.get('text', '*Nội dung không có*')
+
+            # === THAY ĐỔI 2: Thêm gợi ý về loại xe vào header ===
+            if matched_vehicle:
+                vehicle_keywords = {
+                    "ô tô": ["ô tô", "xe con"], "xe máy": ["xe máy", "xe mô tô"],
+                    "xe đạp": ["xe đạp", "xe thô sơ"], "máy kéo": ["máy kéo", "xe chuyên dùng"]
+                }
+                article_title_lower = metadata.get("article_title", "").lower()
+                if any(keyword in article_title_lower for keyword in vehicle_keywords.get(matched_vehicle, [])):
+                    header += f" [GỢI Ý: Thông tin này áp dụng cho {matched_vehicle.upper()}]"
+            
             context_parts.append(f"{header}\n{text}")
         context = "\n\n---\n\n".join(context_parts)
 
-    # --- Bước 3: Xây dựng Prompt bằng Chat Template (Đây là phần thay đổi chính) ---
+    # 3. Xây dựng Prompt bằng Chat Template
     print("--- Xây dựng prompt bằng chat template ---")
     llm_model = components["llm_model"]
     tokenizer = components["tokenizer"]
     
-    # Tạo cấu trúc tin nhắn theo chuẩn của các mô hình hội thoại
+    # Tạo cấu trúc tin nhắn theo chuẩn
     messages = [
         {
             "role": "system",
-            "content": "Bạn là một trợ lý pháp luật chuyên trả lời các câu hỏi liên quan đến luật giao thông đường bộ Việt Nam. Hãy dựa vào các thông tin được cung cấp để trả lời một cách chính xác và dễ hiểu."
+            "content": "Bạn là một trợ lý pháp luật chuyên trả lời các câu hỏi về luật giao thông Việt Nam. Hãy dựa vào các thông tin được cung cấp để trả lời một cách chính xác và dễ hiểu."
         },
         {
             "role": "user",
@@ -136,14 +149,19 @@ def generate_response(query: str, components: dict) -> str:
         }
     ]
     
-    # Sử dụng apply_chat_template để tạo prompt hoàn chỉnh.
-    # Thêm `add_generation_prompt=True` để nó tự động thêm vai trò "assistant" ở cuối,
-    # báo hiệu cho model bắt đầu sinh câu trả lời.
-    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    # === THAY ĐỔI 3 (Cốt lõi): Sử dụng apply_chat_template ===
+    # Phương thức này sẽ tạo ra chuỗi prompt hoàn chỉnh với các token đặc biệt,
+    # tương thích với cả mô hình text và vision (khi không có ảnh).
+    prompt = tokenizer.apply_chat_template(
+        messages, 
+        tokenize=False, 
+        add_generation_prompt=True
+    )
 
-    # --- Bước 4: Tạo câu trả lời từ LLM (Không thay đổi) ---
+    # 4. Tạo câu trả lời từ LLM
     print("--- Bắt đầu tạo câu trả lời từ LLM ---")
     
+    # Tokenize chuỗi prompt đã được định dạng đúng
     inputs = tokenizer([prompt], return_tensors="pt").to(llm_model.device)
 
     generation_config = dict(
@@ -158,4 +176,4 @@ def generate_response(query: str, components: dict) -> str:
     response_text = tokenizer.decode(output_ids[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
 
     print("--- Tạo câu trả lời hoàn tất ---")
-    return response_text
+    return response_text