Chat / app.py
HeavensHackDev's picture
Update app.py
f35bde9 verified
raw
history blame
6.06 kB
# app.py
import gradio as gr
from transformers import pipeline, set_seed
import random
# Загружаем модель для генерации текста (text-generation pipeline)
# Выберите модель:
# 'distilgpt2' - маленькая английская модель GPT-2
# 'sberbank-ai/rugpt3small_based_on_gpt2' - маленькая русская модель (может требовать больше ресурсов)
# 'bigscience/bloomz-560m' - многоязычная модель поменьше (может быть медленной на CPU)
# Попробуем с distilgpt2 для начала, т.к. она легче всего запустится на CPU
# Если хотите русскую, раскомментируйте строку с 'sberbank-ai...' и закомментируйте distilgpt2
# generator = pipeline('text-generation', model='distilgpt2', max_new_tokens=50) # Ограничим длину генерации
# generator = pipeline('text-generation', model='sberbank-ai/rugpt3small_based_on_gpt2', max_new_tokens=50)
generator = pipeline('text-generation', model='distilgpt2', max_length=70, num_return_sequences=1, truncation=True) # Альтернативные параметры
print("Модель загружена!") # Для отладки в логах
# Установка зерна для воспроизводимости (опционально)
# set_seed(random.randint(0, 10000)) # Используем случайное зерно для разнообразия
# Функция, которая будет обрабатывать ввод пользователя и историю чата
def respond(message, chat_history):
"""
Принимает сообщение пользователя и историю чата, возвращает ответ бота.
"""
print(f"Получено сообщение: {message}")
print(f"История чата: {chat_history}")
# Формируем промпт для модели (можно экспериментировать)
# Вариант 1: Просто сообщение пользователя
prompt = message
# Вариант 2: Добавляем немного истории (простой)
# history_text = "\n".join([f"User: {m}\nBot: {b}" for m, b in chat_history[-2:]]) # Последние 2 обмена
# prompt = f"{history_text}\nUser: {message}\nBot:"
try:
# Генерируем ответ
# Обратите внимание: 'text-generation' часто лучше работает, если ей дать начало фразы
# Попробуем сделать так, чтобы она просто продолжила мысль или ответила.
# Можно добавить "Ответ:", "Бот:", и т.п. в конец prompt, если модель так лучше реагирует.
bot_response = generator(prompt)[0]['generated_text']
# Очистка ответа: убираем исходный промпт из сгенерированного текста
# (pipeline часто возвращает промпт + генерацию)
if bot_response.startswith(prompt):
bot_response = bot_response[len(prompt):].strip()
# Иногда модели могут генерировать пустые ответы или спецсимволы, добавим проверку
if not bot_response or len(bot_response.strip()) == 0:
bot_response = "... (модель не сгенерировала ответ)"
print(f"Сгенерированный ответ: {bot_response}")
except Exception as e:
print(f"Ошибка при генерации: {e}")
bot_response = f"Ой, произошла ошибка при генерации ответа: {e}"
# Добавляем пару (сообщение пользователя, ответ бота) в историю
chat_history.append((message, bot_response))
# Возвращаем пустую строку для очистки поля ввода и обновленную историю
return "", chat_history
# --- Создание интерфейса Gradio ---
with gr.Blocks() as demo:
gr.Markdown("# Простой Чат-Бот на Hugging Face")
gr.Markdown("Введите ваше сообщение и нажмите Enter.")
# Компонент чат-бота
chatbot = gr.Chatbot(label="Диалог")
# Поле ввода текста
msg = gr.Textbox(label="Ваше сообщение", placeholder="Напишите что-нибудь...")
# Кнопка очистки (опционально)
clear = gr.Button("Очистить диалог")
# Связываем компоненты:
# Когда пользователь отправляет сообщение (нажатием Enter в Textbox или Button, если бы она была)
# вызывается функция respond с текущим сообщением и историей чата.
# Результат функции (пустая строка и новая история) обновляет Textbox (msg) и Chatbot (chatbot).
msg.submit(respond, [msg, chatbot], [msg, chatbot])
# Связываем кнопку очистки
clear.click(lambda: (None, None), None, [msg, chatbot], queue=False)
# Запускаем демо
# concurrency_count - сколько запросов может обрабатывать одновременно
# share=True - если хотите получить временную публичную ссылку (не нужно для Spaces)
demo.launch(debug=True) # Включаем debug для логов в интерфейсе Spaces
# demo.launch(server_name="0.0.0.0", server_port=7860) # Для запуска на некоторых платформах