update main

app.py

@@ -1,7 +1,358 @@
+# app.py
+# File triển khai hoàn chỉnh cho đồ án Chatbot Luật Giao thông
+# Tác giả: (Tên của bạn)
+# Ngày: (Ngày bạn tạo)
+
+# --- PHẦN 1: IMPORT CÁC THƯ VIỆN CẦN THIẾT ---
+print("Bắt đầu import các thư viện...")
+import os
+import sys
+import json
+import re
+import time
+from collections import defaultdict
+
+# Core ML/DL và Unsloth
+import torch
+from unsloth import FastLanguageModel
+from transformers import TextStreamer
+
+# RAG - Retrieval
+import faiss
+from sentence_transformers import SentenceTransformer
+from rank_bm25 import BM25Okapi
+import numpy as np
+
+# Deployment
 import gradio as gr
 
-def greet(name):
-    return "Hello " + name + "!!"
+print("✅ Import thư viện thành công.")
+
+# --- PHẦN 2: CẤU HÌNH VÀ TẢI TÀI NGUYÊN (MODELS & DATA) ---
+# Phần này sẽ chỉ chạy một lần khi ứng dụng khởi động.
+
+# Cấu hình mô hình
+MAX_SEQ_LENGTH = 2048
+DTYPE = None
+LOAD_IN_4BIT = True
+EMBEDDING_MODEL_NAME = "bkai-foundation-models/vietnamese-bi-encoder"
+LLM_MODEL_NAME = "unsloth/Llama-3.2-3B-Instruct-bnb-4bit"
+LAW_DATA_FILE = "luat_chi_tiet_output_openai_sdk_final_cleaned.json"
+
+# Biến toàn cục để lưu các tài nguyên đã tải
+# Điều này giúp tránh việc phải tải lại mô hình mỗi khi người dùng gửi yêu cầu.
+MODELS_AND_DATA = {
+    "llm_model": None,
+    "tokenizer": None,
+    "embedding_model": None,
+    "faiss_index": None,
+    "bm25_model": None,
+    "chunks_data": None,
+    "tokenized_corpus_bm25": None,
+}
+
+# --- Các hàm xử lý dữ liệu (từ các notebook của bạn) ---
+
+def process_law_data_to_chunks(structured_data_input):
+    """
+    Hàm làm phẳng dữ liệu luật có cấu trúc chi tiết thành danh sách các chunks.
+    Mỗi chunk chứa 'text' và 'metadata'.
+    """
+    flat_list = []
+    articles_list = []
+    if isinstance(structured_data_input, dict) and "article" in structured_data_input:
+        articles_list = [structured_data_input]
+    elif isinstance(structured_data_input, list):
+        articles_list = structured_data_input
+    else:
+        print("Lỗi: Dữ liệu đầu vào không hợp lệ.")
+        return flat_list
+
+    for article_data in articles_list:
+        if not isinstance(article_data, dict): continue
+        article_metadata_base = {
+            "source_document": article_data.get("source_document"),
+            "article": article_data.get("article"),
+            "article_title": article_data.get("article_title")
+        }
+        clauses = article_data.get("clauses", [])
+        if not isinstance(clauses, list): continue
+
+        for clause_data in clauses:
+            if not isinstance(clause_data, dict): continue
+            clause_metadata_base = article_metadata_base.copy()
+            clause_metadata_base.update({
+                "clause_number": clause_data.get("clause_number"),
+                "clause_metadata_summary": clause_data.get("clause_metadata_summary")
+            })
+            points_in_clause = clause_data.get("points_in_clause", [])
+            if not isinstance(points_in_clause, list): continue
+
+            if points_in_clause:
+                for point_data in points_in_clause:
+                    if not isinstance(point_data, dict): continue
+                    chunk_text = point_data.get("point_text_original") or point_data.get("violation_description_summary")
+                    if not chunk_text: continue
+                    
+                    current_point_metadata = clause_metadata_base.copy()
+                    point_specific_metadata = point_data.copy()
+                    if "point_text_original" in point_specific_metadata:
+                        del point_specific_metadata["point_text_original"]
+                    current_point_metadata.update(point_specific_metadata)
+                    final_metadata_cleaned = {k: v for k, v in current_point_metadata.items() if v is not None}
+                    flat_list.append({"text": chunk_text, "metadata": final_metadata_cleaned})
+            else:
+                chunk_text = clause_data.get("clause_text_original")
+                if chunk_text:
+                    current_clause_metadata = clause_metadata_base.copy()
+                    additional_clause_info = {k: v for k, value in clause_data.items() if k not in ["clause_text_original", "points_in_clause", "clause_number", "clause_metadata_summary"]}
+                    if additional_clause_info:
+                        current_clause_metadata.update(additional_clause_info)
+                    final_metadata_cleaned = {k: v for k, v in current_clause_metadata.items() if v is not None}
+                    flat_list.append({"text": chunk_text, "metadata": final_metadata_cleaned})
+    return flat_list
+
+def tokenize_vi_for_bm25(text):
+    """Hàm tokenize tiếng Việt đơn giản cho BM25."""
+    text = text.lower()
+    text = re.sub(r'[^\w\s]', '', text)
+    return text.split()
+
+def load_all_resources():
+    """
+    Hàm chính để tải tất cả mô hình và dữ liệu cần thiết.
+    Chỉ chạy một lần khi ứng dụng khởi động.
+    """
+    print("--- Bắt đầu quá trình tải tài nguyên ---")
+    
+    # 1. Tải mô hình LLM và Tokenizer
+    print(f"1. Đang tải LLM và Tokenizer: {LLM_MODEL_NAME}...")
+    llm_model, tokenizer = FastLanguageModel.from_pretrained(
+        model_name=LLM_MODEL_NAME,
+        max_seq_length=MAX_SEQ_LENGTH,
+        dtype=DTYPE,
+        load_in_4bit=LOAD_IN_4BIT,
+    )
+    FastLanguageModel.for_inference(llm_model)
+    MODELS_AND_DATA["llm_model"] = llm_model
+    MODELS_AND_DATA["tokenizer"] = tokenizer
+    print("✅ Tải LLM và Tokenizer thành công.")
+
+    # 2. Tải mô hình Embedding
+    print(f"2. Đang tải Embedding Model: {EMBEDDING_MODEL_NAME}...")
+    embedding_model = SentenceTransformer(EMBEDDING_MODEL_NAME, device="cuda" if torch.cuda.is_available() else "cpu")
+    MODELS_AND_DATA["embedding_model"] = embedding_model
+    print("✅ Tải Embedding Model thành công.")
+
+    # 3. Tải và xử lý dữ liệu luật
+    print(f"3. Đang tải và xử lý dữ liệu từ: {LAW_DATA_FILE}...")
+    if not os.path.exists(LAW_DATA_FILE):
+        raise FileNotFoundError(f"Không tìm thấy file dữ liệu luật: {LAW_DATA_FILE}. Vui lòng upload file này lên Space.")
+    with open(LAW_DATA_FILE, 'r', encoding='utf-8') as f:
+        raw_data_from_file = json.load(f)
+    chunks_data = process_law_data_to_chunks(raw_data_from_file)
+    MODELS_AND_DATA["chunks_data"] = chunks_data
+    print(f"✅ Đã xử lý thành {len(chunks_data)} chunks.")
+
+    # 4. Tạo BM25 Model
+    print("4. Đang tạo BM25 Model...")
+    corpus_texts = [chunk.get('text', '') for chunk in chunks_data]
+    tokenized_corpus = [tokenize_vi_for_bm25(text) for text in corpus_texts]
+    bm25_model = BM25Okapi(tokenized_corpus)
+    MODELS_AND_DATA["bm25_model"] = bm25_model
+    MODELS_AND_DATA["tokenized_corpus_bm25"] = tokenized_corpus
+    print("✅ Tạo BM25 Model thành công.")
+
+    # 5. Tạo FAISS Index
+    print("5. Đang tạo FAISS Index...")
+    texts_to_encode = [chunk.get('text', '') for chunk in chunks_data]
+    chunk_embeddings = embedding_model.encode(texts_to_encode, convert_to_tensor=True, device=embedding_model.device)
+    chunk_embeddings_np = chunk_embeddings.cpu().numpy().astype('float32')
+    faiss.normalize_L2(chunk_embeddings_np)
+    dimension = chunk_embeddings_np.shape[1]
+    index = faiss.IndexFlatIP(dimension)
+    index.add(chunk_embeddings_np)
+    MODELS_AND_DATA["faiss_index"] = index
+    print(f"✅ Tạo FAISS Index thành công với {index.ntotal} vectors.")
+
+    print("\n--- Tải tài nguyên hoàn tất! Ứng dụng đã sẵn sàng. ---")
+
+# --- PHẦN 3: CÁC HÀM LÕI CHO RAG ---
+
+def search_relevant_laws(query_text, k=5, initial_k_multiplier=10, rrf_k_constant=60):
+    """
+    Hàm thực hiện Hybrid Search để tìm các đoạn luật liên quan.
+    """
+    # Lấy các tài nguyên đã tải
+    embedding_model = MODELS_AND_DATA["embedding_model"]
+    faiss_index = MODELS_AND_DATA["faiss_index"]
+    chunks_data = MODELS_AND_DATA["chunks_data"]
+    bm25_model = MODELS_AND_DATA["bm25_model"]
+
+    if not all([embedding_model, faiss_index, chunks_data, bm25_model]):
+        return "Lỗi: Tài nguyên chưa được tải xong. Vui lòng chờ."
+
+    # 1. Semantic Search (FAISS)
+    query_embedding = embedding_model.encode([query_text], convert_to_tensor=True, device=embedding_model.device)
+    query_embedding_np = query_embedding.cpu().numpy().astype('float32')
+    faiss.normalize_L2(query_embedding_np)
+    num_candidates = min(k * initial_k_multiplier, faiss_index.ntotal)
+    semantic_scores, semantic_indices = faiss_index.search(query_embedding_np, num_candidates)
+
+    # 2. Keyword Search (BM25)
+    tokenized_query = tokenize_vi_for_bm25(query_text)
+    bm25_scores = bm25_model.get_scores(tokenized_query)
+    bm25_results = sorted(enumerate(bm25_scores), key=lambda x: x[1], reverse=True)[:num_candidates]
+
+    # 3. Reciprocal Rank Fusion (RRF)
+    rrf_scores = defaultdict(float)
+    if semantic_indices.size > 0:
+        for rank, doc_idx in enumerate(semantic_indices[0]):
+            if doc_idx != -1: rrf_scores[doc_idx] += 1.0 / (rrf_k_constant + rank)
+    for rank, (doc_idx, score) in enumerate(bm25_results):
+        if score > 0: rrf_scores[doc_idx] += 1.0 / (rrf_k_constant + rank)
+    
+    fused_results = sorted(rrf_scores.items(), key=lambda x: x[1], reverse=True)
+
+    # 4. Lấy kết quả cuối cùng
+    final_results = []
+    for doc_idx, score in fused_results[:k]:
+        result = chunks_data[doc_idx].copy()
+        result['score'] = score
+        final_results.append(result)
+        
+    return final_results
+
+def generate_llm_response(query, context):
+    """
+    Hàm sinh câu trả lời từ LLM dựa trên query và context.
+    """
+    llm_model = MODELS_AND_DATA["llm_model"]
+    tokenizer = MODELS_AND_DATA["tokenizer"]
+
+    prompt = f"""Dưới đây là một số thông tin trích dẫn từ văn bản luật giao thông đường bộ Việt Nam.
+Hãy SỬ DỤNG CÁC THÔNG TIN NÀY để trả lời câu hỏi một cách chính xác và đầy đủ.
+Nếu câu hỏi đưa ra nhiều đáp án thì chọn 1 đáp án đúng nhất.
+
+### Thông tin luật:
+{context}
+
+### Câu hỏi:
+{query}
+
+### Trả lời:"""
+
+    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
+    generation_config = dict(
+        max_new_tokens=300,
+        temperature=0.2,
+        top_p=0.7,
+        do_sample=True,
+        pad_token_id=tokenizer.eos_token_id,
+        eos_token_id=tokenizer.eos_token_id
+    )
+    output_ids = llm_model.generate(**inputs, **generation_config)
+    input_length = inputs.input_ids.shape[1]
+    generated_ids = output_ids[0][input_length:]
+    response_text = tokenizer.decode(generated_ids, skip_special_tokens=True)
+    return response_text
+
+# --- PHẦN 4: CÁC HÀM XỬ LÝ CHO GRADIO INTERFACE ---
+
+def run_retriever_only(query):
+    """
+    Chức năng 1: Chỉ tìm kiếm và trả về các điều luật liên quan.
+    """
+    print(f"Chạy chức năng Retriever cho query: '{query}'")
+    retrieved_results = search_relevant_laws(query)
+    
+    if isinstance(retrieved_results, str): # Xử lý trường hợp lỗi
+        return retrieved_results
+
+    if not retrieved_results:
+        return "Không tìm thấy điều luật nào liên quan."
+
+    # Định dạng output cho Gradio Markdown
+    formatted_output = f"### Các điều luật liên quan nhất đến truy vấn của bạn:\n\n"
+    for i, res in enumerate(retrieved_results):
+        metadata = res.get('metadata', {})
+        article = metadata.get('article', 'N/A')
+        clause = metadata.get('clause_number', 'N/A')
+        source = metadata.get('source_document', 'N/A')
+        text = res.get('text', 'N/A')
+        
+        formatted_output += f"**{i+1}. Nguồn: {source} | Điều {article} | Khoản {clause}**\n"
+        formatted_output += f"> {text}\n\n---\n\n"
+        
+    return formatted_output
+
+def run_full_rag(query, progress=gr.Progress()):
+    """
+    Chức năng 2: Thực hiện toàn bộ pipeline RAG.
+    """
+    progress(0, desc="Bắt đầu...")
+    
+    # Bước 1: Truy xuất ngữ cảnh
+    progress(0.2, desc="Đang tìm kiếm các điều luật liên quan (Hybrid Search)...")
+    print(f"Chạy chức năng RAG cho query: '{query}'")
+    retrieved_results = search_relevant_laws(query)
+    
+    if isinstance(retrieved_results, str) or not retrieved_results:
+        context_for_llm = "Không tìm thấy thông tin luật liên quan."
+        context_for_display = context_for_llm
+    else:
+        # Định dạng context cho LLM
+        context_parts = []
+        for res in retrieved_results:
+            text = res.get('text', '')
+            context_parts.append(text)
+        context_for_llm = "\n\n---\n\n".join(context_parts)
+        
+        # Định dạng context để hiển thị cho người dùng
+        context_for_display = run_retriever_only(query) # Tái sử dụng hàm retriever
+    
+    # Bước 2: Sinh câu trả lời
+    progress(0.7, desc="Đã có ngữ cảnh, đang yêu cầu LLM tạo câu trả lời...")
+    final_answer = generate_llm_response(query, context_for_llm)
+    
+    progress(1, desc="Hoàn tất!")
+    
+    return final_answer, context_for_display
+
+
+# --- PHẦN 5: KHỞI CHẠY ỨNG DỤNG GRADIO ---
+
+# Tải tài nguyên ngay khi script được chạy
+load_all_resources()
+
+with gr.Blocks(theme=gr.themes.Soft(), title="Chatbot Luật Giao thông") as demo:
+    gr.Markdown(
+        """
+        # ⚖️ Chatbot Luật Giao thông Việt Nam
+        Ứng dụng này sử dụng mô hình RAG (Retrieval-Augmented Generation) để trả lời các câu hỏi về luật giao thông.
+        """
+    )
+    
+    with gr.Tabs():
+        # Tab 1: Chỉ tìm kiếm
+        with gr.TabItem("Tìm kiếm Điều luật (Retriever)"):
+            with gr.Row():
+                retriever_query = gr.Textbox(label="Nhập nội dung cần tìm kiếm", placeholder="Ví dụ: Vượt đèn đỏ bị phạt bao nhiêu tiền?", scale=4)
+                retriever_button = gr.Button("Tìm kiếm", variant="secondary", scale=1)
+            retriever_output = gr.Markdown(label="Các điều luật liên quan")
+        
+        # Tab 2: Hỏi-đáp RAG đầy đủ
+        with gr.TabItem("Hỏi-Đáp (RAG)"):
+            with gr.Row():
+                rag_query = gr.Textbox(label="Nhập câu hỏi của bạn", placeholder="Ví dụ: Phương tiện giao thông đường bộ gồm những loại nào?", scale=4)
+                rag_button = gr.Button("Gửi câu hỏi", variant="primary", scale=1)
+            rag_answer = gr.Textbox(label="Câu trả lời của Chatbot", lines=5)
+            with gr.Accordion("Xem ngữ cảnh đã sử dụng để tạo câu trả lời", open=False):
+                rag_context = gr.Markdown(label="Ngữ cảnh")
+
+    # Xử lý sự kiện click
+    retriever_button.click(fn=run_retriever_only, inputs=retriever_query, outputs=retriever_output)
+    rag_button.click(fn=run_full_rag, inputs=rag_query, outputs=[rag_answer, rag_context])
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch()
+if __name__ == "__main__":
+    demo.launch(share=True) # share=True để tạo link public nếu chạy trên Colab/local