Update app.py

app.py

@@ -1,6 +1,5 @@
 import duckdb
 import pandas as pd
-from tabulate import tabulate
 import gradio as gr
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
@@ -8,93 +7,65 @@ from sentence_transformers import SentenceTransformer, util
 import threading
 
 # --------------------------
-# Setup: Load data and models
+# Setup: Load dữ liệu và mô hình
 # --------------------------
 
-# Load Excel data into a pandas DataFrame.
+# Đọc dữ liệu từ file Excel vào DataFrame
 df = pd.read_excel("mau_bao_cao.xlsx")
 
-# (Optional) Create a DuckDB table if needed.
+# Tạo bảng production_data trong DuckDB (nếu cần)
 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 embedding model for computing embeddings.
+# 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
 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 Qwen model and tokenizer for conversational generation.
-fc_model = AutoModelForCausalLM.from_pretrained(
-    "Qwen/Qwen2.5-3B-Instruct",
-    torch_dtype=torch.float16,
-)
-
+# 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')
 
-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
-
 # --------------------------
-# Define the streaming generation function
+# Hàm tạo phản hồi streaming theo thời gian thực
 # --------------------------
 
 def generate_response(user_query: str):
     """
-    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.
+    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.
     """
-    # 1. Embed the user query.
+    # Tính embedding cho câu truy vấn
     question_embedding = embedding_model.encode(user_query, convert_to_tensor=True)
-
-    # 2. Find best matching columns (top 10).
-    k = 10
+    
+    # Chọn top 3 cột phù hợp
+    k = 7
     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).
+    
+    # Chọn 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]
-
-    # 4. Format the filtered data as a table.
-    table_output = tabulate(filtered_df, headers=best_column_names, tablefmt="grid")
-
-    # 5. Build the system prompt.
+    
+    # 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ý 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ọ._**
@@ -131,21 +102,17 @@ 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")
-
-    # 8. Use TextIteratorStreamer to yield tokens as they are generated.
+    response_inputs = fc_tokenizer(response_template, return_tensors="pt").to(fc_model.device)
+    
+    # Dùng TextIteratorStreamer để stream phản hồi
     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,
@@ -156,49 +123,44 @@ 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
 
 # --------------------------
-# Build the Gradio conversation interface
+# Hàm giao diện chat
 # --------------------------
 
 def chat_interface(user_message, history):
     """
-    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].
+    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ử.
     """
-    # 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
 
-with gr.Blocks() as demo:
-    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)
-
+# --------------------------
+# Xây dựng giao diện Gradio với 2 tab: Chat và Production Data Sample
+# --------------------------
 
-demo.launch()
+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()