chattets / app.py
DakMak's picture
Update app.py
5aa3c7f
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
import gradio as gr
# Загрузка модели DialoGPT и токенизатора
model_name = "microsoft/DialoGPT-large"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)
# Функция, которая генерирует ответ на основе входного текста
def generate_response(prompt):
# Преобразование входного текста в токены
input_ids = tokenizer.encode(prompt + tokenizer.eos_token, return_tensors="pt")
# Генерация ответа на основе входных токенов
chat_response = model.generate(
input_ids,
max_length=1000,
pad_token_id=tokenizer.eos_token_id,
do_sample=True,
top_k=50,
top_p=0.95,
temperature=0.8,
no_repeat_ngram_size=2
)
# Декодирование ответа из токенов в текст
response = tokenizer.decode(chat_response[0], skip_special_tokens=True)
return response
# Создание интерфейса с помощью Gradio
inputs = gr.inputs.Textbox(lines=7, label="Введите сообщение")
outputs = gr.outputs.Textbox(label="Ответ")
title = "Chat.UI на основе DialoGPT"
description = "Простой чат на основе модели DialoGPT из библиотеки transformers"
examples = [
["Привет, как дела?"],
["Что ты думаешь о новом фильме?"]
]
chat_interface = gr.Interface(
fn=generate_response,
inputs=inputs,
outputs=outputs,
title=title,
description=description,
examples=examples
)
# Запуск интерфейса
chat_interface.launch()