import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
import spaces

# 1. Triệu hồi bộ não Gemma 2B
model_id = "google/gemma-2b-it"

tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.bfloat16,
    device_map="auto",
)

# 2. Hàm xử lý API - Tập trung vào tốc độ và độ chính xác
@spaces.GPU
def api_chat(message, history):
    # Chuyển đổi format cho Gemma
    conversation = []
    for user, assistant in history:
        conversation.extend([{"role": "user", "content": user}, {"role": "model", "content": assistant}])
    conversation.append({"role": "user", "content": message})

    # Mã hóa đầu vào
    input_ids = tokenizer.apply_chat_template(conversation, return_tensors="pt", add_generation_prompt=True).to(model.device)
    
    # Tạo phản hồi (không dùng Streamer để trả về 1 cục JSON cho dễ xử lý ở Web)
    output_ids = model.generate(
        input_ids,
        max_new_tokens=1024,
        do_sample=True,
        temperature=0.7,
    )
    
    # Giải mã và trả kết quả
    response = tokenizer.decode(output_ids[0][len(input_ids[0]):], skip_special_tokens=True)
    return response

# 3. Khởi tạo interface kiểu "Cổng kết nối"
demo = gr.ChatInterface(fn=api_chat)

if __name__ == "__main__":
    demo.launch()