import gradio as gr
from transformers import pipeline
import spaces

qa_pipeline = pipeline("text-generation", model="tiiuae/falcon-rw-1b")

@spaces.GPU
def respond(message, history):
    prompt = "Ты — дружелюбный бот, знающий всё об Университете Иннополис.\n"
    for user, bot in history:
        prompt += f"Пользователь: {user}\nБот: {bot}\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=["Когда основан Университет Иннополис?", "Как поступить в магистратуру?"],
    cache_examples=False  # 🛠 Отключаем сохранение длинной истории
)

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