Update app.py

app.py

@@ -1,25 +1,43 @@
 # app.py
-
 import gradio as gr
 from transformers import pipeline, set_seed
 import random
+import torch # Может понадобиться для указания устройства
+
+# --- Настройки ---
+# Выбираем русскую модель
+# Попробуем 'sberbank-ai/rugpt3small_based_on_gpt2'.
+# Если Space будет падать из-за нехватки памяти, возможно, придется вернуться к distilgpt2
+# или искать еще более легкую модель.
+MODEL_NAME = 'sberbank-ai/rugpt3small_based_on_gpt2'
+# MODEL_NAME = 'distilgpt2' # Запасной вариант
+
+# Параметры генерации
+MAX_NEW_TOKENS = 100 # Увеличим максимальную длину *нового* текста
+TEMPERATURE = 0.8 # Температура генерации (чуть больше случайности, >1 = больше бреда)
+TOP_P = 0.9 # Ядерная выборка (nucleus sampling)
+# --- Конец Настроек ---
+
+
+print(f"Загрузка модели: {MODEL_NAME}...")
+try:
+    # Попробуем указать device_map='auto' для лучшего распределения на ресурсах Space
+    # или torch_dtype=torch.float16 для экономии памяти (если поддерживается)
+    generator = pipeline(
+        'text-generation',
+        model=MODEL_NAME,
+        # Если есть GPU и CUDA: device=0
+        # Для CPU или автоматического определения: device=-1 (по умолчанию)
+        # Попробуем автоматическое распределение:
+        # device_map="auto", # Может помочь с памятью на много-GPU инстансах (редко в free Spaces)
+        # torch_dtype=torch.float16, # Если модель поддерживает и есть CUDA, ускорит и сэкономит память
+    )
+    print("Модель успешно загружена!")
+except Exception as e:
+    print(f"Ошибка при загрузке модели: {e}")
+    # Завершаем работу, если модель не загрузилась
+    raise SystemExit(f"Не удалось загрузить модель {MODEL_NAME}. Ошибка: {e}")
 
-# Загружаем модель для генерации текста (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):
@@ -27,70 +45,98 @@ def respond(message, chat_history):
     Принимает сообщение пользователя и историю чата, возвращает ответ бота.
     """
     print(f"Получено сообщение: {message}")
-    print(f"История чата: {chat_history}")
+    print(f"История чата (до): {chat_history}")
+
+    # --- Улучшенный Промпт ---
+    # Добавляем инструкцию/роль для бота (можно сделать сложнее)
+    system_prompt = "Ты — дружелюбный и полезный ИИ-ассистент. Отвечай на вопросы пользователя."
+
+    # Формируем историю для промпта (последние N обменов)
+    history_limit = 3 # Сколько последних пар "вопрос-ответ" учитывать
+    history_for_prompt = []
+    for user_msg, bot_msg in chat_history[-history_limit:]:
+        history_for_prompt.append(f"По��ьзователь: {user_msg}")
+        history_for_prompt.append(f"Ассистент: {bot_msg}")
 
-    # Формируем промпт для модели (можно экспериментировать)
-    # Вариант 1: Просто сообщение пользователя
-    prompt = message
+    history_str = "\n".join(history_for_prompt)
 
-    # Вариант 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:"
+    # Собираем финальный промпт
+    prompt = f"{system_prompt}\n\n{history_str}\nПользователь: {message}\nАссистент:"
+    print(f"--- Промпт для модели --- \n{prompt}\n------------------------")
 
     try:
         # Генерируем ответ
-        # Обратите внимание: 'text-generation' часто лучше работает, если ей дать начало фразы
-        # Попробуем сделать так, чтобы она просто продолжила мысль или ответила.
-        # Можно добавить "Ответ:", "Бот:", и т.п. в конец prompt, если модель так лучше реагирует.
+        # Используем max_new_tokens для контроля длины именно ответа
+        generated_outputs = generator(
+            prompt,
+            max_new_tokens=MAX_NEW_TOKENS,
+            num_return_sequences=1,
+            # truncation=True, # Трункация уже не так нужна, если мы контролируем длину истории
+            do_sample=True, # Включаем сэмплирование для более "живых" ответов
+            temperature=TEMPERATURE,
+            top_p=TOP_P,
+            # pad_token_id=generator.tokenizer.eos_token_id # Указываем токен для паддинга
+        )
+        bot_response = generated_outputs[0]['generated_text']
+
+        # --- Очистка ответа ---
+        # Убираем весь промпт из сгенерированного текста
+        if bot_response.startswith(prompt):
+             bot_response = bot_response[len(prompt):].strip()
+        else:
+             # Иногда модель может начать ответ сразу, без повторения промпта,
+             # но может повторить последнюю часть (например, "Ассистент:")
+             # Это требует более сложной очистки, пока оставим базовую.
+             # Можно попробовать убирать только последнюю строку промпта:
+             last_prompt_line = "Ассистент:"
+             if bot_response.strip().startswith(last_prompt_line):
+                 bot_response = bot_response.strip()[len(last_prompt_line):].strip()
 
-        bot_response = generator(prompt)[0]['generated_text']
 
-        # Очистка ответа: убираем исходный промпт из сгенерированного текста
-        # (pipeline часто возвращает промпт + генерацию)
-        if bot_response.startswith(prompt):
-            bot_response = bot_response[len(prompt):].strip()
+        # Дополнительная очистка от незавершенных предложений или артефактов
+        # (можно добавить обрезку по последнему знаку препинания .?!)
+        last_punctuation = max(bot_response.rfind('.'), bot_response.rfind('!'), bot_response.rfind('?'))
+        if last_punctuation != -1:
+             bot_response = bot_response[:last_punctuation+1]
+
 
-        # Иногда модели могут генерировать пустые ответы или спецсимволы, добавим проверку
+        # Проверка на пустой ответ
         if not bot_response or len(bot_response.strip()) == 0:
-            bot_response = "... (модель не сгенерировала ответ)"
+            bot_response = "... (модель не сгенерировала содержательный ответ)"
 
-        print(f"Сгенерированный ответ: {bot_response}")
+        print(f"Сгенерированный ответ (очищенный): {bot_response}")
 
     except Exception as e:
         print(f"Ошибка при генерации: {e}")
-        bot_response = f"Ой, произошла ошибка при генерации ответа: {e}"
+        # Попробуем дать более специфичную информацию, если возможно
+        if "out of memory" in str(e).lower():
+             bot_response = "Ой, кажется, мне не хватило памяти для генерации ответа. Попробуйте более короткий запрос или администратор может попробовать модель поменьше."
+        else:
+             bot_response = f"Ой, произошла ошибка при генерации ответа: {e}"
 
     # Добавляем пару (сообщение пользователя, ответ бота) в историю
     chat_history.append((message, bot_response))
+    print(f"История чата (после): {chat_history}")
 
     # Возвращаем пустую строку для очистки поля ввода и обновленную историю
     return "", chat_history
 
-# --- Создание интерфейса Gradio ---
+# --- Создание интерфейса Gradio (остается почти без изменений) ---
 with gr.Blocks() as demo:
-    gr.Markdown("# Простой Чат-Бот на Hugging Face")
-    gr.Markdown("Введите ваше сообщение и нажмите Enter.")
-
-    # Компонент чат-бота
-    chatbot = gr.Chatbot(label="Диалог")
-
-    # Поле ввода текста
-    msg = gr.Textbox(label="Ваше сообщение", placeholder="Напишите что-нибудь...")
+    gr.Markdown("# Улучшенный Чат-Бот на Hugging Face")
+    gr.Markdown(f"Используемая модель: `{MODEL_NAME}`. Введите ваше сообщение.")
 
-    # Кнопка очистки (опционально)
+    chatbot = gr.Chatbot(label="Диалог", height=500) # Увеличим высоту окна чата
+    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, [msg, chatbot], queue=False) # Очищаем историю на пустой список []
 
-    # Связываем кнопку очистки
-    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) # Для запуска на некоторых платформах
+print("Запуск Gradio интерфейса...")
+# Используйте share=True только если запускаете локально и хотите временную публичную ссылку
+# Для HF Spaces это не нужно.
+demo.launch(debug=True)
+# demo.launch(server_name="0.0.0.0", server_port=7860)