import os
import gradio as gr
from transformers import DistilGPT2Tokenizer, DistilGPT2LMHeadModel
import torch

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

# Загрузка модели и токенизатора с использованием токена
model_name = "distilgpt2"
tokenizer = DistilGPT2Tokenizer.from_pretrained(model_name, use_auth_token=HUGGINGFACE_TOKEN)
model = DistilGPT2LMHeadModel.from_pretrained(model_name, use_auth_token=HUGGINGFACE_TOKEN)

# Функция для ведения диалога
def chat_with_model(user_input):
    input_ids = tokenizer.encode(user_input + tokenizer.eos_token, return_tensors='pt')
    chat_history_ids = model.generate(input_ids, max_length=100, pad_token_id=tokenizer.eos_token_id)
    bot_response = tokenizer.decode(chat_history_ids[:, input_ids.shape[-1]:][0], skip_special_tokens=True)
    return bot_response

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

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