tg-invip / app.py
elasko-aim's picture
Update app.py
fd22c23 verified
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
import time
# Загрузка модели и токенизатора
MODEL_NAME = "microsoft/DialoGPT-medium"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
# Инициализация NPC
npc_names = ["Agent Alpha", "Agent Beta", "Agent Gamma"]
chat_history = []
def generate_response(input_text):
"""
Генерирует ответ от NPC на основе входного текста.
"""
# Токенизация входного текста
input_ids = tokenizer.encode(input_text + tokenizer.eos_token, return_tensors="pt")
# Генерация ответа
response_ids = model.generate(
input_ids,
max_length=1000,
pad_token_id=tokenizer.eos_token_id,
no_repeat_ngram_size=2,
top_p=0.95,
top_k=50
)
# Декодирование ответа
response = tokenizer.decode(response_ids[:, input_ids.shape[-1]:][0], skip_special_tokens=True)
return response
def npc_chat():
"""
Обрабатывает диалог между NPC.
"""
global chat_history
# Начальное сообщение
if not chat_history:
chat_history.append(f"{npc_names[0]}: Welcome to Earth Web Space Networks!")
for i, npc in enumerate(npc_names):
# Получаем последнее сообщение из истории
last_message = chat_history[-1].split(": ", 1)[1]
# Генерируем ответ
response = generate_response(last_message)
new_message = f"{npc}: {response}"
chat_history.append(new_message)
# Ограничение длины истории чата
if len(chat_history) > 20:
chat_history.pop(0)
# Возвращаем текущую историю чата
formatted_chat = "\n".join(chat_history)
return formatted_chat
# Gradio интерфейс
with gr.Blocks() as demo:
gr.Markdown("# Earth Web Space Networks Live Chat")
with gr.Row():
chat_output = gr.Textbox(label="Chat History", lines=10, interactive=False)
# Кнопка для обновления чата
update_button = gr.Button("Update Chat")
update_button.click(npc_chat, inputs=None, outputs=chat_output)
demo.launch()