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

# Указание названия модели
model_name = "cody82/innopolis_bot_model"

# Получение токена из переменной окружения
hf_token = os.getenv("HUGGINGFACE_HUB_TOKEN")
if hf_token is None:
    raise ValueError("Переменная окружения HUGGINGFACE_HUB_TOKEN не установлена.")

# Загрузка токенизатора и модели с авторизацией
tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=hf_token)
model = AutoModelForCausalLM.from_pretrained(model_name, use_auth_token=hf_token)

# Использование CPU
device = torch.device("cpu")
model.to(device)

# Обработка входящих сообщений (новый формат: messages)
def respond(message, history):
    full_input = ""
    for turn in history:
        if turn["role"] == "user":
            full_input += f"User: {turn['content']}\n"
        elif turn["role"] == "assistant":
            full_input += f"Assistant: {turn['content']}\n"
    full_input += f"User: {message}\nAssistant:"

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

# Интерфейс Gradio
chat = gr.ChatInterface(
    fn=respond,
    chatbot=gr.Chatbot(label="Innopolis Bot", type="messages"),
    title="Innopolis Chatbot",
    theme="soft",
    examples=["Когда был основан университет Иннополис?", "Какие программы есть в магистратуре?"],
)

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