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-1.5B-Instruct', torch_dtype=torch.float16)
fc_tokenizer = AutoTokenizer.from_pretrained('Qwen/Qwen2.5-1.5B-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 = 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]
    
    # 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ý 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}"
    )
    
    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")
    
    # 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()