Spaces:
Sleeping
Sleeping
| 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() | |