Spaces:
Sleeping
Sleeping
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() | |