import gradio as gr
from transformers import pipeline
import spaces  # обязательно для ZeroGPU

# Модель (CPU-friendly)
qa_pipeline = pipeline("text-generation", model="tiiuae/falcon-rw-1b")

# Основной обработчик
@spaces.GPU  # 👈 обязательно для ZeroGPU, даже если модель на CPU
def respond(message, history):
    system_prompt = (
        "Ты — дружелюбный бот, который отлично знает Университет Иннополис. "
        "Отвечай чётко и по делу на вопросы про университет, город Иннополис, обучение и студенческую жизнь.\n"
    )

    prompt = system_prompt
    for user_msg, bot_msg in history:
        prompt += f"Пользователь: {user_msg}\nБот: {bot_msg}\n"
    prompt += f"Пользователь: {message}\nБот:"

    output = qa_pipeline(
        prompt,
        max_new_tokens=200,
        do_sample=True,
        temperature=0.7,
        top_p=0.95,
        pad_token_id=50256,
    )[0]["generated_text"]

    response = output.split("Бот:")[-1].strip().split("Пользователь:")[0].strip()
    history.append((message, response))
    return history

# Интерфейс
demo = gr.ChatInterface(
    fn=respond,
    title="Innopolis Bot",
    theme="soft",
    examples=["Когда основан Университет Иннополис?", "Как поступить в магистратуру?"],
)

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