import duckdb import pandas as pd from tabulate import tabulate import gradio as gr import torch from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer from sentence_transformers import SentenceTransformer, util import threading # -------------------------- # Setup: Load data and models # -------------------------- # Load dữ liệu từ file Excel vào DataFrame df = pd.read_excel("mau_bao_cao.xlsx") # (Tùy chọn) Tạo bảng 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 mô hình embedding để tính toán 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 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') # -------------------------- # Helper function: Trích xuất bảng dữ liệu # -------------------------- 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 streaming tạo phản hồi từ mô hình # -------------------------- 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. """ # Lấy bảng dữ liệu liên quan filtered_df, best_column_names, table_text = extract_table(user_query) # Tạo system prompt có 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ọ._** 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! """ messages = [ {'role': 'system', 'content': system_prompt}, {'role': 'user', 'content': user_query} ] 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 streamer = TextIteratorStreamer(fc_tokenizer, skip_prompt=True, skip_special_tokens=True) # Khởi chạy generation trong một thread riêng thread = threading.Thread( target=lambda: fc_model.generate( **response_inputs, max_new_tokens=512, temperature=1, top_p=0.95, streamer=streamer ) ) thread.start() collected_text = "" for new_text in streamer: collected_text += new_text yield collected_text # -------------------------- # Hàm giao diện chat của Gradio # -------------------------- 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). """ # 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. 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 # Stream phản hồi từ mô hình theo thời gian thực for partial_response in generate_response(user_message): history[-1][1] = partial_response yield "", history, filtered_df # -------------------------- # Xây dựng giao diện Gradio # -------------------------- 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) demo.launch()