Update app.py

app.py

@@ -1,142 +1,187 @@
 # 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)
-# --- Конец Настроек ---
+import re # Для регулярных выражений (очистка текста)
 
+# --- Загрузка моделей ---
+# 1. Генератор текста (для общего ответа и диалога)
+# Замените 'distilgpt2' на нужную модель, если используете другую
+# Попробуйте 'microsoft/DialoGPT-medium' - специально для диалогов (английский)
+# Или русские аналоги, если найдете подходящие по размеру для CPU
+try:
+    # generator = pipeline('text-generation', model='distilgpt2', max_length=100, num_return_sequences=1, truncation=True, pad_token_id=50256) # Указываем pad_token_id явно
+    # generator = pipeline('conversational', model='microsoft/DialoGPT-medium') # Специальная диалоговая модель (англ)
+    generator = pipeline('text-generation', model='sberbank-ai/rugpt3small_based_on_gpt2', max_length=100, num_return_sequences=1, truncation=True, pad_token_id=50256) # Попробуем русскую поменьше
+    print("Генератор текста загружен.")
+    GENERATOR_LOADED = True
+except Exception as e:
+    print(f"ОШИБКА: Не удалось загрузить генератор текста! {e}")
+    generator = None
+    GENERATOR_LOADED = False
 
-print(f"Загрузка модели: {MODEL_NAME}...")
+# 2. Модель для ответов на вопросы (Question Answering)
+# Используем небольшую мультиязычную модель
 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("Модель успешно загружена!")
+    qa_pipeline = pipeline('question-answering', model='bert-large-uncased-whole-word-masking-finetuned-squad') # Пример англоязычной
+    # qa_pipeline = pipeline('question-answering', model='timpal0l/mdeberta-v3-base-squad2') # Пример мультиязычной (может быть лучше для RU)
+    print("QA модель загружена.")
+    QA_LOADED = True
 except Exception as e:
-    print(f"Ошибка при загрузке модели: {e}")
-    # Завершаем работу, если модель не загрузилась
-    raise SystemExit(f"Не удалось загрузить модель {MODEL_NAME}. Ошибка: {e}")
+    print(f"ОШИБКА: Не удалось загрузить QA модель! {e}")
+    qa_pipeline = None
+    QA_LOADED = False
+
+# --- "Встроенные" знания и команды ---
+knowledge_base = {
+    "кто ты": "Я Nova, модель Alpha 0.95, работающая на платформе Hugging Face Spaces. Я использую нейросетевые модели для генерации ответов и ответов на вопросы.",
+    "что ты умеешь": "Я могу пытаться поддерживать диалог, отвечать на вопросы по предоставленному контексту (если удастся его определить), �� генерировать текст на основе предыдущих сообщений. Также я знаю некоторые встроенные команды.",
+    "как дела": "Как у программы, у меня все по плану! Готова обрабатывать ваши запросы.",
+    "помощь": "Спросите меня о чем-нибудь, или дайте команду. Я попробую ответить. Примеры команд: 'кто ты', 'что ты умеешь'.",
+    # Можно добавить больше простых ответов
+}
 
+# --- Функция обработки запроса ---
 
-# Функция, которая будет обрабатывать ввод пользователя и историю чата
 def respond(message, chat_history):
     """
-    Принимает сообщение пользователя и историю чата, возвращает ответ бота.
+    Главная функция обработки. Определяет тип запроса и выбирает стратегию ответа.
     """
     print(f"Получено сообщение: {message}")
-    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}")
-
-    history_str = "\n".join(history_for_prompt)
-
-    # Собираем финальный промпт
-    prompt = f"{system_prompt}\n\n{history_str}\nПользователь: {message}\nАссистент:"
-    print(f"--- Промпт для модели --- \n{prompt}\n------------------------")
-
-    try:
-        # Генерируем ответ
-        # Используем 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()
-
-
-        # Дополнительная очистка от незавершенных предложений или артефактов
-        # (можно добавить обрезку по последнему знаку препинания .?!)
-        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 = "... (модель не сгенерировала содержательный ответ)"
-
-        print(f"Сгенерированный ответ (очищенный): {bot_response}")
-
-    except Exception as e:
-        print(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}")
+    user_message_lower = message.lower().strip().replace("?","").replace(".","").replace("!","") # Очищенный ввод для команд
+
+    # 1. Проверка на встроенные команды/знания
+    if user_message_lower in knowledge_base:
+        bot_response = knowledge_base[user_message_lower]
+        print(f"Ответ из базы знаний: {bot_response}")
+        chat_history.append((message, bot_response))
+        return "", chat_history
+
+    # 2. Попытка ответа на вопрос с помощью QA модели
+    # Очень упрощенная логика: если есть "?" и QA модель загружена
+    if "?" in message and QA_LOADED and qa_pipeline:
+        # Используем историю чата как контекст (последние сообщения)
+        context = ""
+        if chat_history:
+             # Собираем текст из последних N обменов (N=2)
+            context_parts = []
+            for user_msg, bot_msg in chat_history[-2:]:
+                 context_parts.append(f"Пользователь: {user_msg}")
+                 context_parts.append(f"Бот: {bot_msg}")
+            context = "\n".join(context_parts)
+
+        if len(context) > 50: # Нужен минимальный контекст для QA
+             print("Попытка ответа на вопрос через QA модель...")
+             print(f"Контекст: {context[:200]}...") # Логируем начало контекста
+             print(f"Вопрос: {message}")
+             try:
+                 qa_result = qa_pipeline(question=message, context=context)
+                 print(f"Результат QA: {qa_result}")
+                 # Если модель уверена в ответе (score > порога)
+                 if qa_result and qa_result['score'] > 0.1: # Порог уверенности (нужно подбирать)
+                     bot_response = qa_result['answer']
+                     chat_history.append((message, bot_response))
+                     print(f"Ответ от QA модели: {bot_response}")
+                     return "", chat_history
+                 else:
+                      print("QA модель не уверена в ответе.")
+             except Exception as e:
+                  print(f"Ошибка при использовании QA модели: {e}")
+                  # Если QA не сработала, переходим к генерации
+
+    # 3. Генерация ответа с помощью основной модели (если команды/QA не сработали)
+    if GENERATOR_LOADED and generator:
+        print("Генерация ответа с помощью основной модели...")
+        # Формируем промпт с историей
+        prompt_history = []
+        for user_msg, bot_msg in chat_history[-3:]: # Берем последние 3 обмена
+            prompt_history.append(f"Пользователь: {user_msg}")
+            prompt_history.append(f"Бот: {bot_msg}")
+        prompt_history.append(f"Пользователь: {message}")
+        prompt_history.append("Бот:") # Приглашение для генерации ответа
+
+        full_prompt = "\n".join(prompt_history)
+        print(f"Промпт для генератора: {full_prompt[-500:]}") # Логируем конец промпта
+
+        try:
+            # Уменьшаем max_length, если промпт уже длинный
+            # prompt_tokens = len(generator.tokenizer.encode(full_prompt)) # Подсчет токенов может быть медленным
+            # available_tokens = generator.model.config.max_length - prompt_tokens - 10 # Запас
+            # current_max_length = max(10, min(generator.model.config.max_length, prompt_tokens + 50)) # Генерация ~50 токенов
+            # print(f"Макс. длина для генерации: {current_max_length}")
+
+            generated_output = generator(full_prompt)[0]['generated_text']
+
+            # Очистка ответа
+            # Убираем промпт из начала
+            bot_response = generated_output
+            if generated_output.startswith(full_prompt):
+                 bot_response = generated_output[len(full_prompt):].strip()
+            else:
+                 # Если промпт не найден в начале (некоторые модели ведут себя иначе)
+                 # Попробуем убрать последнее сообщение пользователя
+                 last_user_line = f"Пользователь: {message}\nБот:"
+                 if bot_response.strip().startswith(last_user_line.strip()): # Убираем и саму строку "Бот:"
+                     bot_response = bot_response.strip()[len(last_user_line.strip()):].strip()
+                 # Очень грубый способ убрать возможное эхо промпта
+                 elif len(bot_response) > len(message) and message.lower() in bot_response[:len(message)*2].lower():
+                      parts = bot_response.split('\n')
+                      if len(parts)>1:
+                          bot_response = '\n'.join(parts[1:]).strip()
+
+
+            # Дополнительная очистка (убираем возможные "Пользователь:", "Бот:")
+            bot_response = re.sub(r'^\s*(пользователь|user|бот|bot)\s*[:\-]?\s*', '', bot_response, flags=re.IGNORECASE).strip()
+
+            if not bot_response or len(bot_response) < 3: # Проверка на совсем короткий/пустой ответ
+                bot_response = "Хм, не знаю, что на это ответить..."
+                print("Сгенерирован слишком короткий ответ, используется заглушка.")
+
+        except Exception as e:
+            print(f"Ошибка при генерации: {e}")
+            bot_response = f"Упс, произошла ошибка при генерации: {e}"
+
+    else:
+        # Если ни одна модель не загружена
+        bot_response = "К сожалению, мои основные модули сейчас недоступны."
+        print("Ошибка: Генератор текста не загружен.")
 
-    # Возвращаем пустую строку для очистки поля ввода и обновленную историю
+    chat_history.append((message, bot_response))
+    print(f"Финальный ответ: {bot_response}")
     return "", chat_history
 
-# --- Создание интерфейса Gradio (остается почти без изменений) ---
-with gr.Blocks() as demo:
-    gr.Markdown("# Улучшенный Чат-Бот на Hugging Face")
-    gr.Markdown(f"Используемая модель: `{MODEL_NAME}`. Введите ваше сообщение.")
 
-    chatbot = gr.Chatbot(label="Диалог", height=500) # Увеличим высоту окна чата
-    msg = gr.Textbox(label="Ваше сообщение", placeholder="Спроси меня что-нибудь...")
-    clear = gr.Button("Очистить диалог")
+# --- Создание интерфейса Gradio ---
+with gr.Blocks(css=".gradio-container {background-color: #f9f9f9}") as demo: # Добавим немного CSS
+    gr.Markdown(
+        """
+        # Nova Alpha 0.95 🚀
+        Простой чат-бот с базовыми знаниями и попыткой ответа на вопросы.
+        Работает на CPU, ответы могут быть медленными.
+        """
+    )
+
+    chatbot = gr.Chatbot(label="Диалог", height=500) # Увеличим высоту
+    with gr.Row(): # Поместим текстбокс и кнопки в ряд
+        msg = gr.Textbox(
+            label="Ваше сообщение",
+            placeholder="Спросите 'кто ты' или задайте вопрос...",
+            scale=4 # Текстбокс займет больше места
+        )
+        submit_btn = gr.Button("Отправить", variant="primary", scale=1) # Явная кнопка отправки
+        clear_btn = gr.Button("Очистить", scale=1)
+
 
+    # Связываем обработчики:
+    # Отправка по Enter в Textbox
     msg.submit(respond, [msg, chatbot], [msg, chatbot])
-    clear.click(lambda: (None, []), None, [msg, chatbot], queue=False) # Очищаем историю на пустой список []
+    # Отправка по клику на кнопку
+    submit_btn.click(respond, [msg, chatbot], [msg, chatbot])
 
+    # Очистка по клику на кнопку
+    clear_btn.click(lambda: (None, []), None, [msg, chatbot], queue=False) # Возвращаем пустой список для истории
 
-# Запускаем демо
-print("Запуск Gradio интерфейса...")
-# Используйте share=True тол��ко если запускаете локально и хотите временную публичную ссылку
-# Для HF Spaces это не нужно.
-demo.launch(debug=True)
-# demo.launch(server_name="0.0.0.0", server_port=7860)
+# Запуск (важно для Spaces!)
+demo.queue() # Используем очередь для обработки запросов по одному (лучше для CPU)
+demo.launch(debug=True) # debug=True покажет логи и ошибки в интерфейсе