Update app.py

app.py

@@ -5,6 +5,9 @@ import os
 import shutil
 import time
 import logging
+import gc
+import signal
+from contextlib import contextmanager
 
 # Настройка логирования
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
@@ -20,17 +23,26 @@ else:
 # Информация о системе и CUDA
 logger.info("===== Запуск приложения =====")
 logger.info(f"PyTorch: {torch.__version__}")
-logger.info(f"CUDA доступен: {torch.cuda.is_available()}")
 
-if torch.cuda.is_available():
+# Проверка CUDA и соответствующие логи
+cuda_available = torch.cuda.is_available()
+logger.info(f"CUDA доступен: {cuda_available}")
+
+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):
             logger.info(f"CUDA устройство {i}: {torch.cuda.get_device_name(i)}")
         logger.info(f"Текущее CUDA устройство: {torch.cuda.current_device()}")
+        
+        # Проверка доступной памяти
+        for i in range(cuda_device_count):
+            free_mem = torch.cuda.get_device_properties(i).total_memory - torch.cuda.memory_allocated(i)
+            logger.info(f"Устройство {i}: свободно {free_mem / 1024**3:.2f} ГБ из {torch.cuda.get_device_properties(i).total_memory / 1024**3:.2f} ГБ")
     except Exception as e:
         logger.warning(f"Ошибка при получении информации о CUDA: {e}")
+        cuda_available = False
 else:
     logger.info("CUDA недоступен, будет использоваться CPU")
 
@@ -70,8 +82,15 @@ logger.info("Информация о дисках перед загрузкой:
 check_disk_space("/")
 check_disk_space(DISK_DIR)
 
-# Загрузка модели и токенизатора
-model_name = "KoDer123/Nerealnost_8M"
+# Настройка модели - выбор в зависимости от доступных ресурсов
+if cuda_available:
+    # Для режима GPU используем более крупную модель (если она есть)
+    model_name = "KoDer123/Nerealnost_8M"  # Ваша основная модель
+else:
+    # Для режима CPU можно выбрать более легкую модель
+    model_name = "KoDer123/Nerealnost_8M"  # Можно заменить на более легкую, если нужно
+
+logger.info(f"Выбрана модель: {model_name}")
 
 # Глобальные переменные для модели
 model = None
@@ -81,15 +100,40 @@ is_model_loaded = False
 # Переопределяем EOS_TOKEN для случая, когда токенизатор не загружен
 DEFAULT_EOS_TOKEN = "</s>"
 
+# Класс таймаута для ограничения времени генерации
+class TimeoutException(Exception):
+    pass
+
+@contextmanager
+def time_limit(seconds):
+    def signal_handler(signum, frame):
+        raise TimeoutException("Timeout")
+    
+    signal.signal(signal.SIGALRM, signal_handler)
+    signal.alarm(seconds)
+    try:
+        yield
+    finally:
+        signal.alarm(0)
+
+# Функция для очистки памяти
+def clear_memory():
+    if cuda_available:
+        torch.cuda.empty_cache()
+    gc.collect()
+
 # Функция для загрузки модели
 def load_model():
     global model, tokenizer, is_model_loaded
     
     try:
+        # Очищаем память перед загрузкой
+        clear_memory()
+        
         logger.info("Загружаем токенизатор...")
         tokenizer = AutoTokenizer.from_pretrained(
             model_name,
-            token=HF_TOKEN,  # Добавляем токен для доступа к закрытой модели
+            token=HF_TOKEN,
             cache_dir=CACHE_DIR,
             local_files_only=False
         )
@@ -103,18 +147,21 @@ def load_model():
         model_kwargs = {
             "cache_dir": CACHE_DIR,
             "trust_remote_code": True,
-            "token": HF_TOKEN  # Добавляем токен для доступа к закрытой модели
+            "token": HF_TOKEN
         }
         
         # Проверяем доступность CUDA
-        if torch.cuda.is_available():
+        if cuda_available:
             logger.info("Загружаем модель в режиме GPU...")
             model_kwargs.update({
                 "torch_dtype": torch.float16,
+                "device_map": "auto",  # Автоматически распределить по устройствам
             })
         else:
             logger.info("Загружаем модель в режиме CPU...")
-            # Без параметров device_map и low_cpu_mem_usage
+            model_kwargs.update({
+                "torch_dtype": torch.float32,
+            })
         
         # Загружаем модель
         model = AutoModelForCausalLM.from_pretrained(
@@ -122,14 +169,15 @@ def load_model():
             **model_kwargs
         )
         
-        # Переносим модель на доступное устройство
-        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        model = model.to(device)
+        # Если GPU недоступен, явно переносим модель на CPU
+        if not cuda_available:
+            model = model.to("cpu")
         
-        logger.info(f"Модель успешно загружена на устройство: {device}")
+        device_info = next(model.parameters()).device
+        logger.info(f"Модель успешно загружена на устройство: {device_info}")
         
         is_model_loaded = True
-        return "Модель успешно загружена"
+        return f"Модель успешно загружена на {device_info}"
     except Exception as e:
         error_msg = str(e)
         logger.error(f"Ошибка загрузки модели: {error_msg}")
@@ -160,6 +208,7 @@ def respond(
     max_tokens,
     temperature,
     top_p,
+    generation_timeout,
 ):
     global model, tokenizer, is_model_loaded
     
@@ -170,6 +219,9 @@ def respond(
         else:
             return "Модель не загружена или произошла ошибка при загрузке. Проверьте логи для получения дополнительной информации."
     
+    # Очищаем память перед генерацией
+    clear_memory()
+    
     # Замеряем время
     start_time = time.time()
     
@@ -187,35 +239,41 @@ def respond(
     logger.info(f"Генерируем ответ на запрос: '{message[:50]}...' (длина промпта: {len(full_prompt)})")
     
     try:
-        # Токенизация входных данных
-        inputs = tokenizer(full_prompt, return_tensors="pt").to(model.device)
-        
-        # Генерация ответа
-        outputs = model.generate(
-            inputs.input_ids,  # Передаем только input_ids, не весь словарь
-            max_new_tokens=max_tokens,
-            temperature=temperature,
-            top_p=top_p,
-            do_sample=True,
-            pad_token_id=tokenizer.pad_token_id,
-        )
-        
-        # Декодирование полного вывода
-        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()
+        # Используем таймаут для генерации
+        with time_limit(generation_timeout):
+            # Токенизация входных данных
+            inputs = tokenizer(full_prompt, return_tensors="pt").to(model.device)
+            
+            # Генерация ответа
+            outputs = model.generate(
+                inputs.input_ids,
+                max_new_tokens=max_tokens,
+                temperature=temperature,
+                top_p=top_p,
+                do_sample=True,
+                pad_token_id=tokenizer.pad_token_id,
+            )
+            
+            # Декодирование полного вывода
+            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)} символов")
         
         return response
+        
+    except TimeoutException:
+        logger.warning(f"Генерация превысила лимит времени ({generation_timeout} секунд)")
+        return f"Генерация ответа превысила лимит времени ({generation_timeout} секунд). Попробуйте уменьшить количество токенов или задать более простой вопрос."
     except Exception as e:
         logger.error(f"Ошибка при генерации ответа: {str(e)}")
         return f"Произошла ошибка при генерации ответа: {str(e)}"
@@ -236,7 +294,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
     
     with gr.Row():
         with gr.Column(scale=4):
-            chatbot = gr.Chatbot(label="Диалог", type="messages")  # Исправление типа чатбота
+            chatbot = gr.Chatbot(label="Диалог", type="messages")
             user_input = gr.Textbox(
                 placeholder="Введите ваш вопрос здесь...",
                 label="Ваш вопрос",
@@ -256,14 +314,14 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
                 )
                 max_tokens = gr.Slider(
                     minimum=1,
-                    maximum=2048,
-                    value=512,
+                    maximum=1024,
+                    value=64 if not cuda_available else 256,  # Меньше токенов для CPU
                     step=1,
                     label="Максимальное число токенов"
                 )
                 temperature = gr.Slider(
                     minimum=0.1,
-                    maximum=4.0,
+                    maximum=1.2,
                     value=0.7,
                     step=0.1,
                     label="Температура"
@@ -271,15 +329,22 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
                 top_p = gr.Slider(
                     minimum=0.1,
                     maximum=1.0,
-                    value=0.95,
+                    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="Таймаут генерации (секунды)"
+                )
             
-            with gr.Accordion("Информация", open=True):
+            with gr.Accordion("Информация о системе", open=True):
                 info_text = gr.Markdown(f"""
                 * **Модель**: {model_name}
-                * **Режим работы**: {"GPU" if torch.cuda.is_available() else "CPU"}
+                * **Режим работы**: {"GPU" if cuda_available else "CPU"}
                 * **Директория для кэша**: {CACHE_DIR}
                 * **Статус загрузки**: {"Успешно" if is_model_loaded else "Ошибка"}
                 * **API токен**: {"Настроен" if HF_TOKEN else "Отсутствует"}
@@ -308,7 +373,8 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
             system_msg.value, 
             max_tokens.value, 
             temperature.value, 
-            top_p.value
+            top_p.value,
+            generation_timeout.value
         )
         
         # Добавляем в историю и возвращаем