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

# Модель и токен
model_name = "microsoft/DialoGPT-medium"
huggingface_token = os.getenv('HUGGINGFACE_TOKEN')

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

# Функция для общения с моделью
def chat_with_model(input_text, chat_history=[]):
    new_user_input_ids = tokenizer.encode(input_text + tokenizer.eos_token, return_tensors="pt")
    
    # Если есть история чата, объединяем её с новым вводом
    if len(chat_history) > 0:
        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 = model.generate(bot_input_ids, max_length=1000, pad_token_id=tokenizer.pad_token_id)
    
    # Получение текста и вывод ответа
    response = tokenizer.decode(chat_history[:, bot_input_ids.shape[-1]:][0], skip_special_tokens=True)
    
    return response, chat_history

# Интерфейс Gradio
with gr.Blocks() as demo:
    chatbot = gr.Chatbot()
    msg = gr.Textbox()
    state = gr.State([])  # Для сохранения истории чата
    
    def respond(message, chat_history):
        response, chat_history = chat_with_model(message, chat_history)
        return chatbot.update([message, response]), chat_history
    
    msg.submit(respond, [msg, state], [chatbot, state])

demo.launch()