my_test_bot / app.py
Makaria's picture
token
58e1ab8
raw
history blame
1.8 kB
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()