import os
import gradio as gr
import torch
from transformers import pipeline
from huggingface_hub import login

# Lấy token từ biến môi trường
token = os.getenv("TOKEN")

if token is None:
    raise ValueError("TOKEN environment variable not set")

# Đăng nhập vào Hugging Face bằng token
login(token=token)

# Tải mô hình từ Hugging Face Hub
model_name = "HuggingFaceH4/zephyr-orpo-141b-A35b-v0.1"
pipe = pipeline(
    "text-generation",
    model=model_name,
    device_map="auto",
    torch_dtype=torch.bfloat16,
)

# Định nghĩa hàm xử lý
def generate_text(user_input):
    messages = [
        {"role": "system", "content": "You are Zephyr, a helpful assistant."},
        {"role": "user", "content": user_input},
    ]
    outputs = pipe(
        messages,
        max_new_tokens=512,
        do_sample=True,
        temperature=0.7,
        top_k=50,
        top_p=0.95,
    )
    return outputs[0]["generated_text"]

# Tạo interface Gradio
iface = gr.Interface(
    fn=generate_text,
    inputs="text",
    outputs="text",
    title="Zephyr - Text Generation",
    description="Generate text responses using the Zephyr model."
)

# Khởi chạy ứng dụng
iface.launch(share=True)