import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
from peft import PeftModel  # Thư viện PEFT để load Adapter

# Tên Base Model & Adapter Model
base_model_name = "Qwen/Qwen2.5-1.5B-Instruct"
adapter_name = "ShynBui/OU_student_book_v5_1.5"

# Load Base Model
print("🔄 Đang tải Base Model...")
base_model = AutoModelForCausalLM.from_pretrained(
    base_model_name, torch_dtype=torch.float16, device_map="cpu"
)
tokenizer = AutoTokenizer.from_pretrained(base_model_name)

# Load Adapter Model
print("🔄 Đang tải Adapter Model...")
model = PeftModel.from_pretrained(base_model, adapter_name)

# Dùng pipeline để gọi mô hình (chạy trên CPU)
text_generator = pipeline("text-generation", model=model, tokenizer=tokenizer, device=-1)

# Hàm sinh văn bản có thể chỉnh Temperature, Top-P, Top-K
def generate_text(prompt, temperature=0.7, top_p=0.9, top_k=50, max_length=512):
    outputs = text_generator(
        prompt,
        max_length=max_length,  # Giới hạn số token sinh ra
        temperature=temperature,  # Điều chỉnh độ sáng tạo
        top_p=top_p,  # Nucleus Sampling
        top_k=top_k  # Top-k Sampling
    )
    return outputs[0]["generated_text"]

# Giao diện Gradio với thanh trượt điều chỉnh
demo = gr.Interface(
    fn=generate_text,
    inputs=[
        gr.Textbox(label="Nhập câu hỏi"),
        gr.Slider(0.1, 1.5, value=0.7, label="Temperature"),
        gr.Slider(0.1, 1.0, value=0.9, label="Top-p"),
        gr.Slider(1, 100, value=50, label="Top-k"),
        gr.Slider(100, 2048, value=512, label="Max Tokens"),
    ],
    outputs="text",
    title="OU Student Book Chatbot",
    description="Chatbot sử dụng mô hình fine-tuned trên Qwen2.5-1.5B, chạy trên CPU với các tham số điều chỉnh.",
)

# Chạy ứng dụng
demo.launch()