Update app.py

app.py

@@ -11,30 +11,32 @@ import threading
 # Setup: Load data and models
 # --------------------------
 
-# Load dữ liệu từ file Excel vào DataFrame
+# Load Excel data into a pandas DataFrame.
 df = pd.read_excel("mau_bao_cao.xlsx")
 
-# (Tùy chọn) Tạo bảng trong DuckDB nếu cần
+# (Optional) Create a DuckDB table if needed.
 conn = duckdb.connect('mau_bao_cao.db')
 conn.execute("""\
 CREATE TABLE IF NOT EXISTS production_data AS
     SELECT * FROM read_xlsx('mau_bao_cao.xlsx');
 """)
 
-# Load mô hình embedding để tính toán embedding cho cột và dòng dữ liệu
+# Load embedding model for computing embeddings.
 embedding_model = SentenceTransformer("intfloat/multilingual-e5-large-instruct")
 column_names = df.columns.tolist()
 column_embeddings = embedding_model.encode(column_names, convert_to_tensor=True)
 row_texts = df.apply(lambda row: " | ".join(row.astype(str)), axis=1)
 row_embeddings = embedding_model.encode(row_texts.tolist(), convert_to_tensor=True)
 
-# Load mô hình Qwen và tokenizer cho việc tạo phản hồi
-fc_model = AutoModelForCausalLM.from_pretrained('Qwen/Qwen2.5-3B-Instruct', torch_dtype=torch.float16, device_map="auto")
-fc_tokenizer = AutoTokenizer.from_pretrained('Qwen/Qwen2.5-3B-Instruct')
+# Load Qwen model and tokenizer for conversational generation.
+fc_model = AutoModelForCausalLM.from_pretrained(
+    "Qwen/Qwen2.5-3B-Instruct",
+    torch_dtype=torch.float16,   # Sử dụng float16 để tiết kiệm VRAM
+    device_map="cuda",           # Ép mô hình chạy trên GPU
+    trust_remote_code=True        # Cho phép tải mã nguồn từ repo
+).cuda()  # Chắc chắn mô hình được chuyển lên GPU
 
-# --------------------------
-# Helper function: Trích xuất bảng dữ liệu
-# --------------------------
+fc_tokenizer = AutoTokenizer.from_pretrained('Qwen/Qwen2.5-3B-Instruct')
 
 def extract_table(user_query: str):
     """
@@ -64,21 +66,37 @@ def extract_table(user_query: str):
     return filtered_df, best_column_names, table_text
 
 # --------------------------
-# Hàm streaming tạo phản hồi từ mô hình
+# Define the streaming generation function
 # --------------------------
 
 def generate_response(user_query: str):
     """
-    Hàm generator để:
-      - Trích xuất bảng dữ liệu dựa trên câu truy vấn.
-      - Tạo system prompt dựa trên bảng dữ liệu đã trích xuất.
-      - Dùng TextIteratorStreamer để tạo phản hồi theo thời gian thực.
-    Yields (trả về) phản hồi được cập nhật theo từng token.
+    A generator function that:
+      1. Embeds the query.
+      2. Selects top matching columns and rows from the data.
+      3. Prepares a system prompt with the extracted table.
+      4. Uses TextIteratorStreamer to stream the generated response.
+    Yields the partial generated text as it is updated.
     """
-    # Lấy bảng dữ liệu liên quan
-    filtered_df, best_column_names, table_text = extract_table(user_query)
+    # 1. Embed the user query.
+    question_embedding = embedding_model.encode(user_query, convert_to_tensor=True)
+    
+    # 2. Find best matching columns (top 10).
+    k = 10
+    column_similarities = util.cos_sim(question_embedding, column_embeddings)[0]
+    best_column_indices = torch.topk(column_similarities, k).indices.tolist()
+    best_column_names = [column_names[i] for i in best_column_indices]
+    
+    # 3. Select top matching rows (top 10).
+    row_similarities = util.cos_sim(question_embedding, row_embeddings).squeeze(0)
+    m = 10
+    best_row_indices = torch.topk(row_similarities, m).indices.tolist()
+    filtered_df = df.iloc[best_row_indices][best_column_names]
+    
+    # 4. Format the filtered data as a table.
+    table_output = tabulate(filtered_df, headers=best_column_names, tablefmt="grid")
     
-    # Tạo system prompt có chứa thông tin bảng dữ liệu
+    # 5. Build the system prompt.
     system_prompt = f"""\
 Bạn là một trợ lý báo cáo sản xuất thông minh, chuyên phân tích và tổng hợp dữ liệu một cách rõ ràng, dễ hiểu.
 **_Chỉ báo cáo nếu người dùng yêu cầu mà nếu không thì cứ giao tiếp bình thường với họ._**
@@ -116,18 +134,20 @@ Ví dụ:
 Bạn đã sẵn sàng phân tích và đưa ra báo cáo!    
 """
     
+    # 6. Create the conversation messages.
     messages = [
         {'role': 'system', 'content': system_prompt},
         {'role': 'user', 'content': user_query}
     ]
     
+    # 7. Prepare the prompt for the model.
     response_template = fc_tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
     response_inputs = fc_tokenizer(response_template, return_tensors="pt").to(fc_model.device)
     
-    # Sử dụng TextIteratorStreamer để stream phản hồi
+    # 8. Use TextIteratorStreamer to yield tokens as they are generated.
     streamer = TextIteratorStreamer(fc_tokenizer, skip_prompt=True, skip_special_tokens=True)
     
-    # Khởi chạy generation trong một thread riêng
+    # Start generation in a separate thread so we can yield tokens as they arrive.
     thread = threading.Thread(
         target=lambda: fc_model.generate(
             **response_inputs,
@@ -139,40 +159,33 @@ Bạn đã sẵn sàng phân tích và đưa ra báo cáo!
     )
     thread.start()
     
+    # 9. Yield tokens incrementally.
     collected_text = ""
     for new_text in streamer:
         collected_text += new_text
         yield collected_text
 
 # --------------------------
-# Hàm giao diện chat của Gradio
+# Build the Gradio conversation interface
 # --------------------------
 
 def chat_interface(user_message, history):
     """
-    Generator cho giao diện chat:
-      - Cập nhật lịch sử cuộc trò chuyện với tin nhắn của người dùng.
-      - Tính toán bảng dữ liệu dựa trên truy vấn và cập nhật component hiển thị bảng.
-      - Stream phản hồi của mô hình theo thời gian thực.
-    Lịch sử cuộc trò chuyện được duy trì dưới dạng danh sách các cặp [tin nhắn người dùng, phản hồi AI].
-    Hàm trả về 3 giá trị: giá trị cho Textbox (reset), lịch sử chat, và bảng dữ liệu (dạng DataFrame).
+    A generator function for Gradio that:
+      - Updates the conversation history with the user message.
+      - Streams the model's response token-by-token in real time.
+    The history is maintained as a list of pairs [user_message, bot_response].
     """
-    # Trích xuất bảng để hiển thị cho người dùng
-    filtered_df, _, _ = extract_table(user_message)
-    
-    # Thêm một cặp tin nhắn mới với phản hồi AI ban đầu là chuỗi rỗng.
+    # Create a new conversation entry with user message and an empty bot response.
     history.append([user_message, ""])
-    # Yield trạng thái ban đầu: clear textbox, lịch sử chat cập nhật, và bảng dữ liệu đã trích xuất.
-    yield "", history, filtered_df
+    # Yield the initial state.
+    yield "", history
 
-    # Stream phản hồi từ mô hình theo thời gian thực
+    # Stream tokens from the generate_response generator.
     for partial_response in generate_response(user_message):
+        # Update the latest conversation entry with the partial bot response.
         history[-1][1] = partial_response
-        yield "", history, filtered_df
-
-# --------------------------
-# Xây dựng giao diện Gradio
-# --------------------------
+        yield "", history
 
 with gr.Blocks() as demo:
     gr.Markdown("## Giao diện Chat với Streaming và Hiển thị Bảng Dữ liệu")
@@ -189,4 +202,5 @@ with gr.Blocks() as demo:
     txt.submit(chat_interface, inputs=[txt, state], outputs=[txt, chatbot, table_display], queue=True)
     send_btn.click(chat_interface, inputs=[txt, state], outputs=[txt, chatbot, table_display], queue=True)
 
+
 demo.launch()