Update app.py

app.py

@@ -7,6 +7,7 @@ import logging
 import gc
 import threading
 import json
+from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig
 
 # Настройка логирования
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
@@ -14,119 +15,48 @@ logger = logging.getLogger(__name__)
 
 # Проверка наличия ZERO GPU
 ZERO_GPU_ENABLED = os.environ.get("HF_ZERO_GPU", "0") == "1"
-if ZERO_GPU_ENABLED:
-    logger.info("Обнаружена переменная HF_ZERO_GPU=1, используем Zero GPU режим")
-else:
-    logger.info("Zero GPU не активирован, используем стандартный режим")
+logger.info(f"Zero GPU активирован: {ZERO_GPU_ENABLED}")
 
 # Получаем API токен из переменных окружения
 HF_TOKEN = os.environ.get("HF_TOKEN", None)
-if HF_TOKEN:
-    logger.info("API токен найден")
-else:
-    logger.warning("API токен не найден! Для доступа к закрытой модели необходимо добавить HF_TOKEN в секреты репозитория")
+logger.info("API токен найден" if HF_TOKEN else "API токен не найден! Добавьте HF_TOKEN в секреты репозитория")
 
 # Информация о системе и CUDA
 logger.info("===== Запуск приложения =====")
 logger.info(f"PyTorch: {torch.__version__}")
 
-# Проверяем доступность CUDA более безопасным способом
-try:
-    cuda_available = torch.cuda.is_available()
-    logger.info(f"CUDA доступен: {cuda_available}")
-    
-    # Если Zero GPU включен, но CUDA не обнаружена, выводим предупреждение
-    if ZERO_GPU_ENABLED and not cuda_available:
-        logger.warning("Zero GPU включен, но CUDA не обнаружена. Возможно, требуется перезапуск Space.")
-    
-    if cuda_available:
-        try:
-            cuda_device_count = torch.cuda.device_count()
-            logger.info(f"Количество CUDA устройств: {cuda_device_count}")
-            
-            for i in range(cuda_device_count):
-                try:
-                    logger.info(f"CUDA устройство {i}: {torch.cuda.get_device_name(i)}")
-                    # Проверяем доступную память GPU
-                    try:
-                        free_mem = torch.cuda.get_device_properties(i).total_memory - torch.cuda.memory_allocated(i)
-                        logger.info(f"Устройство {i}: свободно {free_mem / 1024**3:.2f} ГБ")
-                    except:
-                        logger.warning(f"Не удалось определить свободную память на устройстве {i}")
-                except Exception as e:
-                    logger.warning(f"Не удалось получить имя устройства {i}: {e}")
-            
-            try:
-                current_device = torch.cuda.current_device()
-                logger.info(f"Текущее CUDA устройство: {current_device}")
-            except Exception as e:
-                logger.warning(f"Не удалось определить текущее устройство: {e}")
-            
-        except Exception as e:
-            logger.warning(f"Ошибка при получении информации о CUDA: {e}")
-            cuda_available = False
-    else:
-        logger.info("CUDA недоступен, будет использоваться CPU")
-except Exception as e:
-    logger.error(f"Критическая ошибка при проверке CUDA: {e}")
-    cuda_available = False
-    logger.info("Принудительно переключаемся на CPU режим из-за ошибки")
+# Проверка CUDA
+cuda_available = torch.cuda.is_available()
+logger.info(f"CUDA доступен: {cuda_available}")
+if cuda_available:
+    logger.info(f"Количество CUDA устройств: {torch.cuda.device_count()}")
+    for i in range(torch.cuda.device_count()):
+        logger.info(f"CUDA устройство {i}: {torch.cuda.get_device_name(i)}")
+        free_mem = torch.cuda.get_device_properties(i).total_memory - torch.cuda.memory_allocated(i)
+        logger.info(f"Устройство {i}: свободно {free_mem / 1024**3:.2f} ГБ")
+else:
+    logger.info("CUDA недоступен, используется CPU")
 
-# Используем домашнюю директорию пользователя (она всегда должна быть доступна)
+# Настройка директорий
 user_home = os.path.expanduser("~")
 DISK_DIR = os.path.join(user_home, "app_data")
-
-# Создаем директорию, если она не существует
-os.makedirs(DISK_DIR, exist_ok=True)
-logger.info(f"Используем директорию для хр��нения: {DISK_DIR}")
-
-# Настраиваем пути для сохранения моделей
 CACHE_DIR = os.path.join(DISK_DIR, "models_cache")
-TORCH_HOME = os.path.join(DISK_DIR, "torch_home")
-
-# Создаем директории
 os.makedirs(CACHE_DIR, exist_ok=True)
-os.makedirs(TORCH_HOME, exist_ok=True)
-
-# Устанавливаем переменные окружения для управления кэшированием
 os.environ["TRANSFORMERS_CACHE"] = CACHE_DIR
 os.environ["HF_HOME"] = CACHE_DIR
-os.environ["TORCH_HOME"] = TORCH_HOME
-
-# Функция для проверки свободного места на диске
-def check_disk_space(path):
-    try:
-        total, used, free = shutil.disk_usage(path)
-        logger.info(f"Диск {path}: всего {total // (1024**3)} ГБ, свободно {free // (1024**3)} ГБ")
-        return free
-    except Exception as e:
-        logger.warning(f"Не удалось проверить диск {path}: {e}")
-        return None
-
-# Выводим информацию о диске перед загрузкой
-logger.info("Информация о дисках перед загрузкой:")
-check_disk_space("/")
-check_disk_space(DISK_DIR)
-
-# Выбираем модель в зависимости от доступных ресурсов
-if cuda_available:
-    # Для режима GPU используем основную модель
-    model_name = "KoDer123/Nerealnost_phi"
-else:
-    # Для режима CPU - та же модель или можно выбрать модель поменьше
-    model_name = "KoDer123/Nerealnost_phi"
+logger.info(f"Используем директорию для кэша: {CACHE_DIR}")
 
+# Модель
+model_name = "unsloth/Phi-3.5-mini-instruct"
 logger.info(f"Выбрана модель: {model_name}")
 
-# Глобальные переменные для модели
+# Глобальные переменные
 model = None
 tokenizer = None
 is_model_loaded = False
-
-# Переопределяем EOS_TOKEN для случая, когда токенизатор не загружен
 DEFAULT_EOS_TOKEN = "</s>"
 
-# Класс для обработки таймаута, работает в любом потоке
+# Класс для таймаута
 class TimeoutManager:
     def __init__(self, seconds):
         self.seconds = seconds
@@ -153,21 +83,17 @@ class TimeoutManager:
 class TimeoutException(Exception):
     pass
 
-# Функция для очистки памяти
+# Очистка памяти
 def clear_memory():
     if cuda_available:
         torch.cuda.empty_cache()
     gc.collect()
 
-# Функция для загрузки модели
+# Загрузка модели
 def load_model():
     global model, tokenizer, is_model_loaded
     
-    # Подключаем нужные библиотеки при необходимости
-    from transformers import AutoModelForCausalLM, AutoTokenizer
-    
     try:
-        # Очищаем память перед загрузкой
         clear_memory()
         
         logger.info("Загружаем токенизатор...")
@@ -175,140 +101,96 @@ def load_model():
             model_name,
             token=HF_TOKEN,
             cache_dir=CACHE_DIR,
-            local_files_only=False
+            local_files_only=False,
+            revision="main"
         )
-        
-        # Устанавливаем pad_token, если его нет
         if tokenizer.pad_token is None:
             tokenizer.pad_token = tokenizer.eos_token
+        logger.info(f"Токенизатор загружен: vocab_size={tokenizer.vocab_size}")
+        
+        logger.info("Загружаем конфигурацию модели...")
+        config = AutoConfig.from_pretrained(model_name, token=HF_TOKEN, cache_dir=CACHE_DIR)
+        logger.info(f"Конфигурация модели: {config}")
         
         logger.info("Загружаем модель...")
-        # Определяем оптимальный режим загрузки
         model_kwargs = {
             "cache_dir": CACHE_DIR,
             "trust_remote_code": True,
-            "token": HF_TOKEN
+            "token": HF_TOKEN,
+            "config": config
         }
         
-        # Проверяем доступность CUDA
         if cuda_available:
             logger.info("Загружаем модель в режиме GPU...")
             model_kwargs.update({
                 "torch_dtype": torch.float16,
-                "device_map": "auto",  # Автоматически распределить по устройствам
+                "device_map": "auto",
+                "load_in_4bit": True  # Оптимизация от unsloth
             })
         else:
             logger.info("Загружаем модель в режиме CPU...")
             model_kwargs.update({
                 "torch_dtype": torch.float32,
+                "load_in_4bit": False
             })
         
-        # Загружаем модель
         model = AutoModelForCausalLM.from_pretrained(
             model_name,
             **model_kwargs
         )
         
-        # Если GPU недоступен, явно переносим модель на CPU
         if not cuda_available:
             model = model.to("cpu")
         
         device_info = next(model.parameters()).device
-        logger.info(f"Модель успешно загружена на устройство: {device_info}")
+        logger.info(f"Модель загружена на устройство: {device_info}")
         
         is_model_loaded = True
-        return f"Модель успешно загружена на {device_info}"
+        return f"Модель загружена на {device_info}"
     except Exception as e:
-        error_msg = str(e)
-        logger.error(f"Ошибка загрузки модели: {error_msg}")
+        logger.error(f"Ошибка загрузки модели: {str(e)}")
         is_model_loaded = False
-        return f"Ошибка загрузки модели: {error_msg}"
+        return f"Ошибка загрузки модели: {str(e)}"
 
-# Загружаем модель
+# Загружаем модель при запуске
 start_time = time.time()
 load_result = load_model()
-end_time = time.time()
-logger.info(f"Загрузка модели заняла {end_time - start_time:.2f} секунд. Результат: {load_result}")
-
-# Выводим информацию о диске после загрузки
-logger.info("Информация о дисках после загрузки:")
-check_disk_space("/")
-check_disk_space(DISK_DIR)
+logger.info(f"Загрузка заняла {time.time() - start_time:.2f} секунд. Результат: {load_result}")
 
-# Определяем шаблон Q&A, как при обучении
-qa_prompt = "<s>Пользователь: {}\nАссистент: {}"
-EOS_TOKEN = DEFAULT_EOS_TOKEN
-if tokenizer is not None and hasattr(tokenizer, 'eos_token') and tokenizer.eos_token:
-    EOS_TOKEN = tokenizer.eos_token
+# Шаблон для генерации
+EOS_TOKEN = tokenizer.eos_token if tokenizer and tokenizer.eos_token else DEFAULT_EOS_TOKEN
+qa_prompt = "<|user|>{}\n<|assistant|> {}"  # Формат для Phi-3.5-mini-instruct
 
-def respond(
-    message,
-    history,
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-    generation_timeout,
-):
+# Функция генерации ответа
+def respond(message, history, system_message, max_tokens, temperature, top_p, generation_timeout):
     global model, tokenizer, is_model_loaded
     
-    # Проверяем, загружена ли модель
     if not is_model_loaded or model is None or tokenizer is None:
-        if not HF_TOKEN:
-            return "Модель не загружена. Для доступа к закрытой модели требуется добавить HF_TOKEN в секреты репозитория."
-        else:
-            return "Модель не загружена или произошла ошибка при загрузке. Проверьте логи для получения дополнительной информации."
+        return "Модель не загружена. Проверьте логи или добавьте HF_TOKEN."
     
-    # Очищаем память перед генерацией
     clear_memory()
-    
-    # Замеряем время
     start_time = time.time()
     
-    # Преобразование истории в правильный формат
-    formatted_history = []
-    if isinstance(history, list):
-        for item in history:
-            if isinstance(item, tuple) and len(item) == 2:
-                formatted_history.append(item)
-            elif isinstance(item, dict) and "role" in item and "content" in item:
-                # Обрабатываем формат messages
-                if item["role"] == "user":
-                    user_message = item["content"]
-                    assistant_message = None
-                    # Ищем следующее сообщение assistant
-                    idx = history.index(item)
-                    if idx + 1 < len(history) and isinstance(history[idx+1], dict) and history[idx+1].get("role") == "assistant":
-                        assistant_message = history[idx+1].get("content")
-                    if assistant_message:
-                        formatted_history.append((user_message, assistant_message))
-    
-    # Формируем историю в текстовом формате
+    # Форматирование истории
     full_prompt = ""
     if system_message:
         full_prompt += qa_prompt.format(system_message, "") + "\n"
-
-    for user_msg, assistant_msg in formatted_history:
+    
+    for user_msg, assistant_msg in history:
         if user_msg and assistant_msg:
             full_prompt += qa_prompt.format(user_msg, assistant_msg) + EOS_TOKEN + "\n"
-
+    
     full_prompt += qa_prompt.format(message, "")
     
-    logger.info(f"Генерируем ответ на запрос: '{message[:50]}...' (длина промпта: {len(full_prompt)})")
+    logger.info(f"Генерируем ответ на: '{message[:50]}...'")
     
     try:
-        # Настраиваем таймаут
         timeout_mgr = TimeoutManager(generation_timeout)
         timeout_mgr.start()
         
-        # Токенизация входных данных
         inputs = tokenizer(full_prompt, return_tensors="pt").to(model.device)
+        timeout_mgr.check_timeout()
         
-        # Проверяем таймаут
-        if timeout_mgr.timeout_occurred:
-            raise TimeoutException("Timeout during tokenization")
-            
-        # Генерация ответа
         gen_kwargs = {
             "input_ids": inputs.input_ids,
             "max_new_tokens": max_tokens,
@@ -319,162 +201,80 @@ def respond(
         }
         
         outputs = model.generate(**gen_kwargs)
-        
-        # Останавливаем таймаут
         timeout_mgr.stop()
         
-        # Декодирование полного вывода
-        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        
-        # Извлекаем только часть после "Ассистент: "
-        response_start = generated_text.rfind("Ассистент: ") + len("Ассистент: ")
-        if response_start >= len("Ассистент: "):  # Проверяем, что "Ассистент: " найден
-            response = generated_text[response_start:].strip()
-        else:
-            # Если не найдено, возвращаем весь текст
-            response = generated_text.strip()
-        
-        end_time = time.time()
-        generation_time = end_time - start_time
-        logger.info(f"Генерация заняла {generation_time:.2f} секунд. Получен ответ длиной {len(response)} символов")
+        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        response_start = response.rfind("<|assistant|> ") + len("<|assistant|> ")
+        response = response[response_start:].strip() if response_start >= len("<|assistant|> ") else response.strip()
         
+        logger.info(f"Генерация заняла {time.time() - start_time:.2f} секунд")
         return response
         
     except TimeoutException:
-        logger.warning(f"Генерация превысила лимит времени ({generation_timeout} секунд)")
-        return f"Генерация ответа превысила лимит времени ({generation_timeout} секунд). Попробуйте уменьшить количество токенов или задать более простой вопрос."
+        return f"Таймаут генерации ({generation_timeout} секунд)."
     except Exception as e:
-        logger.error(f"Ошибка при генерации ответа: {str(e)}")
-        return f"Произошла ошибка при генерации ответа: {str(e)}"
+        logger.error(f"Ошибка генерации: {str(e)}")
+        return f"Ошибка: {str(e)}"
     finally:
-        # Гарантируем остановку таймера
         if 'timeout_mgr' in locals():
             timeout_mgr.stop()
 
-# Настройка интерфейса Gradio
+# Интерфейс Gradio
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.Markdown("# НереальностьQA - Чат с экспертом по эзотерике")
     
     if not HF_TOKEN:
-        gr.Markdown("""
-        ## ⚠️ Внимание: API токен не найден!
-        
-        Для работы с закрытой моделью необходимо добавить HF_TOKEN в секреты репозитория:
-        1. Settings > Repository secrets > New secret
-        2. Name: HF_TOKEN
-        3. Value: ваш токен доступа с huggingface.co/settings/tokens
-        """, elem_id="warning-box")
+        gr.Markdown("⚠️ Добавьте HF_TOKEN в секреты репозитория!")
     
     with gr.Row():
         with gr.Column(scale=4):
-            chatbot = gr.Chatbot(label="Диалог")  # Используем формат по умолчанию для совместимости
-            user_input = gr.Textbox(
-                placeholder="Введите ваш вопрос здесь...",
-                label="Ваш вопрос",
-                lines=2
-            )
-            
+            chatbot = gr.Chatbot(label="Диалог")
+            user_input = gr.Textbox(placeholder="Введите вопрос...", label="Ваш вопрос", lines=2)
             with gr.Row():
                 submit_btn = gr.Button("Отправить", variant="primary")
-                clear_btn = gr.Button("Очистить историю")
-            
+                clear_btn = gr.Button("Очистить")
+        
         with gr.Column(scale=1):
-            with gr.Accordion("Настройки генерации", open=False):
+            with gr.Accordion("Настройки", open=False):
                 system_msg = gr.Textbox(
-                    value="Твоя задача найти точный ответ на вопрос пользователя.",
+                    value="Твоя задача — дать точный ответ на вопрос пользователя.",
                     label="Системное сообщение",
                     lines=4
                 )
-                max_tokens = gr.Slider(
-                    minimum=1,
-                    maximum=1024,
-                    value=64 if not cuda_available else 256,  # Меньше токенов для CPU
-                    step=1,
-                    label="Максимальное число токенов"
-                )
-                temperature = gr.Slider(
-                    minimum=0.1,
-                    maximum=1.2,
-                    value=0.5,
-                    step=0.1,
-                    label="Температура"
-                )
-                top_p = gr.Slider(
-                    minimum=0.1,
-                    maximum=1.0,
-                    value=0.9,
-                    step=0.05,
-                    label="Top-p"
-                )
-                generation_timeout = gr.Slider(
-                    minimum=10,
-                    maximum=300,
-                    value=60 if cuda_available else 120,  # Больше времени для CPU
-                    step=10,
-                    label="Таймаут генерации (секунды)"
-                )
+                max_tokens = gr.Slider(1, 1024, value=256, step=1, label="Макс. токенов")
+                temperature = gr.Slider(0.1, 1.2, value=0.7, step=0.1, label="Температура")
+                top_p = gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="Top-p")
+                generation_timeout = gr.Slider(10, 300, value=60, step=10, label="Таймаут (с)")
             
-            with gr.Accordion("Информация о системе", open=True):
+            with gr.Accordion("Информация", open=True):
                 system_info = {
                     "Модель": model_name,
-                    "Режим работы": "GPU" if cuda_available else "CPU",
-                    "Zero GPU": "Активирован" if ZERO_GPU_ENABLED else "Не активирован",
-                    "Директория для кэша": CACHE_DIR,
-                    "Статус загрузки": "Успешно" if is_model_loaded else "Ошибка",
+                    "Режим": "GPU" if cuda_available else "CPU",
+                    "Статус": "Успешно" if is_model_loaded else "Ошибка",
                     "API токен": "Настроен" if HF_TOKEN else "Отсутствует"
                 }
-                
-                info_text = gr.Markdown("\n".join([f"* **{k}**: {v}" for k, v in system_info.items()]))
+                gr.Markdown("\n".join([f"* **{k}**: {v}" for k, v in system_info.items()]))
     
-    # Примеры вопросов
     with gr.Accordion("Примеры вопросов", open=True):
-        examples = gr.Examples(
+        gr.Examples(
             examples=[
-                "Что известно о мире отшедших душ и их взаимодействии с нашим миром?",
-                "Что такое энергетическая ось человека и как она связана с его биополем?",
-                "Расскажи о роли энергии мысли и желания в мире отшедших."
+                "Что известно о мире отшедших душ?",
+                "Что такое энергетическая ось человека?",
+                "Роль энергии мысли в мире отшедших?"
             ],
             inputs=user_input
         )
     
-    # Функция обработки отправки сообщения
     def chat(message, history):
-        if message == "":
+        if not message:
             return history, ""
-        
-        # Генерируем ответ
-        bot_message = respond(
-            message, 
-            history, 
-            system_msg.value, 
-            max_tokens.value, 
-            temperature.value, 
-            top_p.value,
-            generation_timeout.value
-        )
-        
-        # Добавляем в историю и возвращаем
-        history = history + [(message, bot_message)]
+        bot_message = respond(message, history, system_msg.value, max_tokens.value, temperature.value, top_p.value, generation_timeout.value)
+        history.append((message, bot_message))
         return history, ""
     
-    # Обработчики событий
-    submit_btn.click(
-        chat,
-        inputs=[user_input, chatbot],
-        outputs=[chatbot, user_input]
-    )
-    
-    user_input.submit(
-        chat,
-        inputs=[user_input, chatbot],
-        outputs=[chatbot, user_input]
-    )
-    
-    clear_btn.click(
-        lambda: ([], ""),
-        outputs=[chatbot, user_input]
-    )
+    submit_btn.click(chat, [user_input, chatbot], [chatbot, user_input])
+    user_input.submit(chat, [user_input, chatbot], [chatbot, user_input])
+    clear_btn.click(lambda: ([], ""), None, [chatbot, user_input])
 
-# Запуск приложения
 if __name__ == "__main__":
     demo.launch()