import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
import spaces  # ← обязательно для ZeroGPU

model_name = "openai-community/gpt2"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)
model.to("cuda" if torch.cuda.is_available() else "cpu")

@spaces.GPU  # 💥 Без этого ZeroGPU не сработает!
def respond(message, history=[]):
    history = history or []
    full_prompt = "Ты — ассистент, который знает всё об Университете Иннополис.\n"

    for user, bot in history:
        full_prompt += f"Пользователь: {user}\nБот: {bot}\n"
    full_prompt += f"Пользователь: {message}\nБот:"

    inputs = tokenizer(full_prompt, return_tensors="pt").to(model.device)
    outputs = model.generate(
        **inputs,
        max_new_tokens=200,
        do_sample=True,
        temperature=0.7,
        top_p=0.9,
        pad_token_id=tokenizer.eos_token_id
    )
    output_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    response = output_text.split("Бот:")[-1].strip()

    history.append((message, response))
    return history

chat = gr.ChatInterface(
    fn=respond,
    title="Innopolis Bot",
    chatbot=gr.Chatbot(label="Unitrip"),
    examples=["Когда основан Университет Иннополис?", "Сколько программ бакалавриата?"],
    cache_examples=False
)

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