app.py

import duckdb
import pandas as pd
import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
from sentence_transformers import SentenceTransformer, util
import threading

# --------------------------
# Setup: Load dữ liệu và mô hình
# --------------------------

# Đọc dữ liệu từ file Excel vào DataFrame
df = pd.read_excel("mau_bao_cao.xlsx")

# 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');
""")

# 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 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')

# --------------------------
# Hàm tạo phản hồi streaming theo thời gian thực
# --------------------------

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.
    """
    # Tính embedding cho câu truy vấn
    question_embedding = embedding_model.encode(user_query, convert_to_tensor=True)
    
    # 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]
    
    # 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]
    
    # 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ọ._**

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)
    
    # Dùng TextIteratorStreamer để stream phản hồi
    streamer = TextIteratorStreamer(fc_tokenizer, skip_prompt=True, skip_special_tokens=True)
    
    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
# --------------------------

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ử.
    """
    history.append([user_message, ""])
    yield "", history
    for partial_response in generate_response(user_message):
        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()