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

# Загрузка модели и токенизатора
model_name = "DialoGPT-medium"  # Используем Diálogo GPT для диалогов
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# Функция для общения с моделью
def chat_with_model(user_input, chat_history):
    # Токенизация пользовательского ввода
    new_user_input_ids = tokenizer.encode(user_input + tokenizer.eos_token, return_tensors='pt')
    
    # Объединение истории чата и нового ввода
    if chat_history:
        bot_input_ids = torch.cat([torch.tensor(chat_history), new_user_input_ids], dim=-1)
    else:
        bot_input_ids = new_user_input_ids

    # Генерация ответа от модели
    chat_history_ids = model.generate(bot_input_ids, max_length=1000, pad_token_id=tokenizer.eos_token_id,
                                       temperature=0.8, top_p=0.9, do_sample=True)

    # Получаем ответ бота и обновляем историю чата
    bot_output = tokenizer.decode(chat_history_ids[:, bot_input_ids.shape[-1]:][0], skip_special_tokens=True)

    # Возвращаем новый ввод и обновлённую историю чата без лишних цифр
    return bot_output, chat_history_ids.tolist()

# Интерфейс Gradio
iface = gr.Interface(
    fn=chat_with_model,
    inputs=[gr.inputs.Textbox(label="Ваше сообщение"), gr.inputs.State()],
    outputs=[gr.outputs.Textbox(label="Ответ бота"), gr.outputs.State()],
    title="Чат-бот с сарказмом и абсурдом",
    description="Побеседуй со своим ботом, который использует сарказм и немного абсурда. Напиши ему что угодно!",
)

iface.launch()