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

# Импортируем токены из переменных окружения
HUGGINGFACE_TOKEN = os.getenv("HUGGINGFACE_TOKEN")

# Загрузка модели и токенизатора DialoGPT
model_name = "microsoft/DialoGPT-medium"  # Можно использовать small или large версии
tokenizer = AutoTokenizer.from_pretrained(model_name, token=HUGGINGFACE_TOKEN)
model = AutoModelForCausalLM.from_pretrained(model_name, token=HUGGINGFACE_TOKEN)

# Функция для ведения диалога
def chat_with_model(user_input, chat_history=[]):
    # Кодируем входное сообщение и добавляем к истории
    new_user_input_ids = tokenizer.encode(user_input + tokenizer.eos_token, return_tensors='pt')

    # Объединяем новую информацию с историей
    bot_input_ids = torch.cat([torch.tensor(chat_history), new_user_input_ids], dim=-1) if chat_history else new_user_input_ids

    # Генерируем ответ
    chat_history_ids = model.generate(
        bot_input_ids,
        max_length=100,
        num_return_sequences=1,
        pad_token_id=tokenizer.eos_token_id,
        temperature=0.7,
        top_k=50,
        top_p=0.95
    )

    # Декодируем ответ
    bot_response = tokenizer.decode(chat_history_ids[:, bot_input_ids.shape[-1]:][0], skip_special_tokens=True)

    # Обновляем историю
    chat_history.append(new_user_input_ids)
    chat_history.append(chat_history_ids[:, bot_input_ids.shape[-1]:])

    return bot_response, chat_history

# Создание интерфейса Gradio
iface = gr.Interface(
    fn=chat_with_model,
    inputs="text",
    outputs="text",
    title="Чатбот на DialoGPT",
    description="Поболтай со своим чатботом!"
)

# Запуск интерфейса
iface.launch()