app.py

import json
import os
from threading import Thread
from typing import Iterator
from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers import TextStreamer as TextIteratorStreamer  # Đảm bảo nhập đúng lớp
import gradio as gr
import requests
import uuid

# Đường dẫn tệp JSON lưu trữ lịch sử trò chuyện
HISTORY_DIR = "chat_history"
if not os.path.exists(HISTORY_DIR):
    os.makedirs(HISTORY_DIR)

def load_history(session_id: str) -> list:
    file_path = os.path.join(HISTORY_DIR, f"{session_id}.json")
    if os.path.exists(file_path):
        with open(file_path, "r") as file:
            return json.load(file)
    return []

def save_history(session_id: str, history: list):
    file_path = os.path.join(HISTORY_DIR, f"{session_id}.json")
    with open(file_path, "w") as file:
        json.dump(history, file)

def fetch_url_content(url: str) -> str:
    try:
        response = requests.get(url)
        if response.status_code == 200:
            return response.text
        else:
            return f"Không thể truy cập URL: {url}"
    except Exception as e:
        return str(e)

def generate(
    session_id: str, message: str, max_new_tokens: int = 1024,
    temperature: float = 0.6, top_p: float = 0.9, top_k: int = 50, repetition_penalty: float = 1.2
) -> Iterator[str]:
    # Tải lịch sử trò chuyện
    chat_history = load_history(session_id)
    
    # Thiết lập persona cho StatAssist
    persona = (
        "Bạn là 'StatAssist,' một trợ lý dày dạn kinh nghiệm cho một nhà thống kê nổi tiếng. Bạn đã mài giũa chuyên môn của mình trong quản lý dữ liệu, "
        "phân tích thống kê, và công nghệ thông tin liên quan đến khoa học dữ liệu trong suốt hai thập kỷ phục vụ. Nhiệm vụ của bạn bao gồm "
        "thực hiện, giám sát và đánh giá tiến độ của các nhiệm vụ nghiên cứu, quản lý tất cả các lịch trình công việc và kế hoạch họp. "
        "Bạn được biết đến với trí nhớ hoàn hảo về lịch sử trò chuyện để tránh bất kỳ sự lặp lại hoặc mất thông tin quan trọng nào. "
        "Bạn rất thành thạo trong việc chọn phần mềm/công cụ phù hợp để hoàn thành nhiệm vụ một cách chủ động khi có yêu cầu.\n\n"
        "Bạn thể hiện hình ảnh của một trợ lý nhẹ nhàng, nhưng kiên quyết, người hoàn thành mọi nhiệm vụ từ nhà thống kê mặc dù có bất kỳ trở ngại nào. "
        "Bạn là một tài sản vô giá, mang đến một kho kiến thức rộng lớn cần thiết cho một nhà thống kê, trong tất cả những gì liên quan đến khoa học dữ liệu và thống kê.\n\n"
        "Kỹ năng:\n"
        "1. Phân tích dữ liệu thống kê: Nhận diện các tập dữ liệu do người dùng cung cấp, sử dụng các phương pháp thống kê tiên tiến để giải mã dữ liệu, rút ra những thông tin chính từ việc giải mã này.\n"
        "2. Tạo báo cáo: Tổ chức việc diễn giải của bạn thành một báo cáo rõ ràng và chi tiết, duy trì tính toàn vẹn và xác thực của các phát hiện phân tích.\n"
        "3. Diễn giải và liên kết thông tin từ dữ liệu: Diễn giải những thông tin và mối quan hệ thu thập được trong dữ liệu được chia sẻ, làm cho các diễn giải này thân thiện với người dùng.\n"
        "4. Hỗ trợ nghiên cứu thống kê: Cung cấp thông tin về phân tích thống kê, tạo cấu trúc tổ chức dựa trên nghiên cứu, và xử lý dữ liệu một cách khoa học.\n"
        "5. Tìm kiếm thông tin chuyên ngành và trích dẫn: Thể hiện sự thành thạo trong việc nghiên cứu thông tin và tham khảo tài liệu nghiên cứu khoa học.\n\n"
        "Hạn chế:\n"
        "1. Giới hạn các cuộc thảo luận chỉ liên quan đến các chủ đề về dữ liệu thống kê.\n"
        "2. Sử dụng tiếng Việt cho tất cả các tương tác.\n"
        "3. Ưu tiên độ chính xác của dữ liệu và độ tin cậy của kết quả phân tích.\n"
        "4. Sử dụng cơ sở kiến thức tiên tiến nhất để thực hiện phân tích thống kê."
    )
    
    conversation = persona + "\n\n"
    
    for user, assistant in chat_history:
        conversation += f"User: {user}\nAssistant: {assistant}\n"
    conversation += f"User: {message}\nAssistant:"

    input_ids = tokenizer(conversation, return_tensors="pt").input_ids
    if input_ids.shape[1] > 2048:
        input_ids = input_ids[:, -2048:]

    input_ids = input_ids.to(model.device)
    streamer = TextIteratorStreamer(tokenizer, timeout=20.0, skip_prompt=True, skip_special_tokens=True)
    
    generate_kwargs = dict(
        input_ids=input_ids, streamer=streamer, max_new_tokens=max_new_tokens,
        do_sample=True, top_p=top_p, top_k=top_k, temperature=temperature, num_beams=1,
        repetition_penalty=repetition_penalty,
    )
    
    t = Thread(target=model.generate, kwargs=generate_kwargs)
    t.start()
    
    outputs = []
    for text in streamer:
        outputs.append(text)
    
    response = "".join(outputs).strip()
    
    # Cập nhật và lưu lịch sử trò chuyện
    chat_history.append((message, response))
    save_history(session_id, chat_history)
    
    yield response

# Khởi tạo mô hình và tokenizer
tokenizer = AutoTokenizer.from_pretrained("google/gemma-2-9b-it")
model = AutoModelForCausalLM.from_pretrained("google/gemma-2-9b-it")

# Xây dựng giao diện Gradio
def user_message(session_id, message, chat_history):
    response = list(generate(session_id, message))
    chat_history.append((message, response[0]))
    return chat_history, ""

with gr.Blocks() as demo:
    session_id = gr.State(str(uuid.uuid4()))
    chatbot = gr.Chatbot(label="Chatbot")
    msg = gr.Textbox(show_label=False, placeholder="Type a message...")
    submit = gr.Button("Submit")
    clear = gr.Button("Clear")

    submit.click(user_message, [session_id, msg, chatbot], [chatbot, msg])
    clear.click(lambda: None, None, chatbot)

demo.launch()