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

MODEL_REPO = "memengoc/newchat"
device = "cuda" if torch.cuda.is_available() else "cpu"

# Load model và tokenizer từ Hugging Face
tokenizer = AutoTokenizer.from_pretrained(MODEL_REPO)
model = AutoModelForCausalLM.from_pretrained(MODEL_REPO).to(device)
model.eval()

# Hàm sinh phản hồi
def generate_response(prompt):
    inputs = tokenizer(prompt, return_tensors="pt").to(device)
    outputs = model.generate(
        **inputs,
        max_new_tokens=300,
        do_sample=True,
        temperature=0.7,
        top_p=0.9,
        pad_token_id=tokenizer.eos_token_id,
    )
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

# Giao diện Gradio
def chat(user_input, history):
    if history is None:
        history = ""
    prompt = f"{history}\nỨng viên: {user_input}\nNhà tuyển dụng:"
    response = generate_response(prompt)
    new_history = f"{prompt} {response}"
    return response, new_history

with gr.Blocks() as demo:
    gr.Markdown("## 🤖 Chatbot Phỏng Vấn AI (OpenChat 3.5 đã Fine-tuned)")

    chat_history = gr.Textbox(visible=False)
    user_input = gr.Textbox(label="Câu trả lời của ứng viên")
    response_output = gr.Textbox(label="Phản hồi nhà tuyển dụng")

    send_btn = gr.Button("Gửi")

    send_btn.click(chat, inputs=[user_input, chat_history], outputs=[response_output, chat_history])

demo.launch()