Update app.py

app.py

@@ -1,5 +1,6 @@
 import duckdb
 import pandas as pd
+from tabulate import tabulate
 import gradio as gr
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
@@ -7,83 +8,144 @@ from sentence_transformers import SentenceTransformer, util
 import threading
 
 # --------------------------
-# Setup: Load dữ liệu và mô hình
+# Setup: Load data and models
 # --------------------------
 
-# Đọc 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ạo bảng production_data 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');
 """)
 
-# Lấy mẫu bảng production_data để hiển thị (ở đây dùng 10 dòng đầu)
-production_data_df = df.head(10)
-
-# Load mô hình embedding để tính 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-1.5B-Instruct', torch_dtype=torch.float16)
-fc_tokenizer = AutoTokenizer.from_pretrained('Qwen/Qwen2.5-1.5B-Instruct')
+# Load Qwen model and tokenizer for conversational generation.
+fc_model = AutoModelForCausalLM.from_pretrained(
+    "Qwen/Qwen2.5-3B-Instruct",
+    torch_dtype=torch.float16,
+)
+
+fc_tokenizer = AutoTokenizer.from_pretrained('Qwen/Qwen2.5-3B-Instruct')
+
+def extract_table(user_query: str):
+    """
+    Dựa trên câu truy vấn của người dùng:
+      - Tính embedding cho câu truy vấn.
+      - Lấy top k cột và top m dòng phù hợp.
+      - Trả về DataFrame đã lọc, danh sách tên cột và bảng dạng text.
+    """
+    # Embed câu truy vấn
+    question_embedding = embedding_model.encode(user_query, convert_to_tensor=True)
+
+    # Lấy top 3 cột phù hợp
+    k = 3
+    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]
+
+    # Lấy top 10 dòng phù hợp
+    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]
+
+    # Tạo bảng text (dùng cho prompt cho mô hình)
+    table_text = tabulate(filtered_df, headers=best_column_names, tablefmt="grid")
+
+    return filtered_df, best_column_names, table_text
 
 # --------------------------
-# Hàm tạo phản hồi streaming theo thời gian thực
+# Define the streaming generation function
 # --------------------------
 
 def generate_response(user_query: str):
     """
-    Hàm này sẽ:
-      - Tính embedding cho câu truy vấn của người dùng.
-      - Chọn ra top 3 cột và top 10 dòng phù hợp từ dữ liệu.
-      - Tạo system prompt bao gồm bảng dữ liệu đã được format bằng tabulate.
-      - Sử dụng TextIteratorStreamer để stream phản hồi từ mô hình theo thời gian thực.
+    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.
     """
-    # Tính embedding cho câu truy vấn
+    # 1. Embed the user query.
     question_embedding = embedding_model.encode(user_query, convert_to_tensor=True)
-    
-    # Chọn top 3 cột phù hợp
-    k = 3
+
+    # 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]
-    
-    # Chọn top 10 dòng phù hợp
+
+    # 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]
-    
-    # Format bảng dữ liệu sử dụng tabulate
-    from tabulate import tabulate
-    table_text = tabulate(filtered_df, headers=best_column_names, tablefmt="grid")
-    
-    # Tạo system prompt chứa thông tin bảng dữ liệu
-    system_prompt = (
-        f"Bạn là một trợ lý AI thông minh, luôn trả lời bằng tiếng Việt.\n"
-        f"Dưới đây là dữ liệu liên quan đến câu hỏi của bạn:\n\n"
-        f"🔹 **Các cột phù hợp**: {', '.join(best_column_names)}\n"
-        f"🔹 **Dữ liệu:**\n{table_text}"
-    )
-    
+
+    # 4. Format the filtered data as a table.
+    table_output = tabulate(filtered_df, headers=best_column_names, tablefmt="grid")
+
+    # 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ọ._**
+
+Dưới đây là dữ liệu bạn cần phân tích:
+🔹 Các cột dữ liệu liên quan: {', '.join(best_column_names)}
+🔹 Bảng dữ liệu:
+{table_output}
+
+📌 Nhiệm vụ của bạn:
+
+Tóm tắt số liệu quan trọng, tránh liệt kê máy móc.
+
+Nhận xét về xu hướng và điểm bất thường.
+
+Nếu có thể, đề xuất giải pháp hoặc hành động tiếp theo.
+
+📊 Cách trả lời:
+✔️ Tự nhiên, dễ hiểu, không quá cứng nhắc.
+✔️ Không cần nhắc lại bảng dữ liệu, hãy diễn giải nó.
+✔️ Trả lời đúng trọng tâm, không dư thừa.
+✔️ Nếu người dùng không hỏi về bảng dữ liệu, hãy chỉ giao tiếp bình thường.
+✔️ Mô hình hóa dữ câu trả lời nếu cần thiết, giúp người dùng dễ hiểu hơn về câu trả lời.
+
+Ví dụ:
+
+🔹 "Hôm nay, sản lượng đạt 95%, cao hơn 5% so với tuần trước."
+
+⚠️ "Dây chuyền A đang giảm hiệu suất, cần theo dõi thêm."
+
+🚀 "Nếu duy trì tốc độ này, sản lượng tháng có thể vượt kế hoạch 10%."
+
+🚀 "Không có gì nếu bạn cần thêm thông tin chi tiết hãy nói cho tôi biết nhé ;))"
+
+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")
-    
-    # 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)
-    
+
+    # Start generation in a separate thread so we can yield tokens as they arrive.
     thread = threading.Thread(
         target=lambda: fc_model.generate(
             **response_inputs,
@@ -94,44 +156,49 @@ def generate_response(user_query: str):
         )
     )
     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
+# Build the Gradio conversation interface
 # --------------------------
 
 def chat_interface(user_message, history):
     """
-    Hàm này sẽ:
-      - Thêm tin nhắn của người dùng vào lịch sử chat (dưới dạng cặp [tin nhắn người dùng, phản hồi AI]).
-      - Stream phản hồi từ mô hình theo thời gian thực và cập nhật lịch sử.
+    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].
     """
+    # Create a new conversation entry with user message and an empty bot response.
     history.append([user_message, ""])
+    # Yield the initial state.
     yield "", history
+
+    # 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
 
-# --------------------------
-# Xây dựng giao diện Gradio với 2 tab: Chat và Production Data Sample
-# --------------------------
-
 with gr.Blocks() as demo:
-    gr.Markdown("## Giao diện Chat và Hiển thị Bảng production_data Mẫu")
-    with gr.Tabs():
-        with gr.TabItem("Chat"):
-            chatbot = gr.Chatbot()
-            state = gr.State([])
-            with gr.Row():
-                txt = gr.Textbox(show_label=False, placeholder="Nhập câu hỏi của bạn...", container=False)
-                send_btn = gr.Button("Gửi")
-            txt.submit(chat_interface, inputs=[txt, state], outputs=[txt, chatbot], queue=True)
-            send_btn.click(chat_interface, inputs=[txt, state], outputs=[txt, chatbot], queue=True)
-        with gr.TabItem("Production Data Sample"):
-            gr.Markdown("Dưới đây là bảng **production_data** mẫu:")
-            production_table = gr.Dataframe(value=production_data_df, label="Production Data Sample")
-    demo.launch()
+    gr.Markdown("## Giao diện Chat với Streaming và Hiển thị Bảng Dữ liệu")
+    chatbot = gr.Chatbot()
+    state = gr.State([])  # duy trì lịch sử chat dưới dạng danh sách các cặp
+    table_display = gr.Dataframe(label="Bảng dữ liệu liên quan")  # hiển thị bảng dữ liệu cho người dùng
+
+    with gr.Row():
+        txt = gr.Textbox(show_label=False, placeholder="Nhập câu hỏi của bạn...", container=False)
+        send_btn = gr.Button("Gửi")
+
+    # Cả submit của Textbox và click của nút gửi đều kích hoạt hàm chat_interface,
+    # trả về 3 outputs: textbox, chatbot và table_display.
+    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()