Add model auto-download and English logging

2026-06-06 12:47:59 +02:00
parent d8b92d5025
commit 0c7f93d067
17 changed files with 204 additions and 77 deletions
@@ -27,11 +27,11 @@ def load_config(config_path: str) -> Dict:
    try:
        with open(config_path, "r", encoding="utf-8") as f:
            config = json.load(f)
-        logger.info("Конфигурация успешно загружена из %s", config_path)
+        logger.info("Configuration loaded successfully from %s", config_path)
        return config
    except FileNotFoundError as e:
-        logger.error("Файл конфигурации не найден: %s", e)
+        logger.error("Configuration file not found: %s", e)
        raise
    except json.JSONDecodeError as e:
-        logger.error("Ошибка при загрузке конфигурации: %s", e)
-        raise
+        logger.error("Failed to load configuration: %s", e)
+        raise
@@ -6,6 +6,7 @@ OpenAI для транскрибации аудиофайлов в текст.
 возможность использования Flash Attention 2 для ускорения работы модели на поддерживаемых GPU.
 """

+import os
 import time
 import threading
 import traceback
@@ -14,6 +15,7 @@ import logging

 import numpy as np
 import torch
+from huggingface_hub import snapshot_download
 from transformers import (
    WhisperForConditionalGeneration,
    WhisperProcessor,
@@ -94,19 +96,19 @@ class WhisperTranscriber:
            
            # Проверяем, что device_id является целым числом
            if not isinstance(device_id, int):
-                logger.warning("device_id должен быть целым числом, получено: %s. Используем значение по умолчанию 0", device_id)
+                logger.warning("device_id must be an integer, got: %s. Using default value 0", device_id)
                device_id = 0
            
            # Проверяем, доступен ли запрошенный GPU
            device_count = torch.cuda.device_count()
            if device_id >= device_count:
-                logger.warning("Запрошенный GPU с индексом %s недоступен. Доступно GPU: %s. Используем GPU с индексом 0", device_id, device_count)
+                logger.warning("Requested GPU index %s is not available. Available GPU count: %s. Using GPU index 0", device_id, device_count)
                device_id = 0
            
-            logger.info("Используется CUDA GPU с индексом %s для вычислений", device_id)
+            logger.info("Using CUDA GPU with index %s for inference", device_id)
            return torch.device(f"cuda:{device_id}")
        elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
-            logger.info("Используется MPS (Apple Silicon) для вычислений")
+            logger.info("Using MPS (Apple Silicon) for inference")
            # Обходное решение для MPS: PyTorch проверяет is_initialized()
            # при создании тензоров на MPS-устройстве, что вызывает ошибку
            # в однопроцессном режиме.
@@ -114,7 +116,7 @@ class WhisperTranscriber:
            setattr(torch.distributed, "is_initialized", lambda: False)
            return torch.device("mps")
        else:
-            logger.info("Используется CPU для вычислений")
+            logger.info("Using CPU for inference")
            return torch.device("cpu")

    def _get_torch_dtype(self) -> torch.dtype:
@@ -135,7 +137,8 @@ class WhisperTranscriber:
        Raises:
            Exception: Если не удалось загрузить модель.
        """
-        logger.info("Загрузка модели из %s", self.model_path)
+        logger.info("Loading model from %s", self.model_path)
+        self.model_path = self._prepare_model_path()

        model_kwargs = dict(
            torch_dtype=self.torch_dtype,
@@ -149,9 +152,9 @@ class WhisperTranscriber:
            capability = torch.cuda.get_device_capability(self.device.index)
            if capability[0] >= 8:
                use_flash_attn = True
-                logger.info("GPU поддерживает Flash Attention 2 (compute capability: %d.%d)", *capability)
+                logger.info("GPU supports Flash Attention 2 (compute capability: %d.%d)", *capability)
            else:
-                logger.info("GPU не поддерживает Flash Attention 2 (compute capability: %d.%d), используется стандартный режим", *capability)
+                logger.info("GPU does not support Flash Attention 2 (compute capability: %d.%d), using standard attention", *capability)

        try:
            if use_flash_attn:
@@ -160,9 +163,9 @@ class WhisperTranscriber:
                self.model_path, **model_kwargs
            ).to(self.device)
            if use_flash_attn:
-                logger.info("Используется Flash Attention 2")
+                logger.info("Using Flash Attention 2")
        except Exception as e:
-            logger.warning("Не удалось загрузить модель с Flash Attention: %s", e)
+            logger.warning("Failed to load model with Flash Attention: %s", e)
            model_kwargs.pop("attn_implementation", None)
            self.model = WhisperForConditionalGeneration.from_pretrained(
                self.model_path, **model_kwargs
@@ -182,7 +185,89 @@ class WhisperTranscriber:
            device=self.device,
        )

-        logger.info("Модель успешно загружена и готова к использованию")
+        logger.info("Model loaded successfully and is ready to use")
+
+    def _prepare_model_path(self) -> str:
+        """
+        Ensures that the configured model path is usable.
+
+        If ``model_path`` looks like a local path and the model is missing,
+        this method can optionally download a fallback Whisper model into that
+        directory. If ``model_path`` is a Hugging Face repo ID, it is returned
+        unchanged and transformers will resolve it normally.
+        """
+        model_path = self.model_path
+
+        looks_like_local_path = (
+            os.path.isabs(model_path)
+            or model_path.startswith(".")
+            or os.path.sep in model_path
+        )
+        if not looks_like_local_path:
+            return model_path
+
+        if self._is_valid_local_model_dir(model_path):
+            return model_path
+
+        if self.config.get("auto_download_missing_model", False):
+            return self._download_missing_model(model_path)
+
+        if not os.path.isdir(model_path):
+            raise FileNotFoundError(
+                "Local model directory not found: "
+                f"'{model_path}'. Mount the Whisper model into the container, "
+                "set 'model_path' to that folder, or enable "
+                "'auto_download_missing_model'."
+            )
+
+        config_file = os.path.join(model_path, "config.json")
+        raise FileNotFoundError(
+            "The configured model directory exists but does not look like a "
+            f"Hugging Face model folder: '{model_path}'. Missing required file "
+            f"'{config_file}'. Make sure the directory contains files like "
+            "'config.json', tokenizer files, and model weights, or enable "
+            "'auto_download_missing_model'."
+        )
+
+    def _is_valid_local_model_dir(self, model_path: str) -> bool:
+        """Returns True when the directory looks like a local HF model folder."""
+        return os.path.isdir(model_path) and os.path.isfile(
+            os.path.join(model_path, "config.json")
+        )
+
+    def _download_missing_model(self, model_path: str) -> str:
+        """
+        Downloads a fallback Whisper model into the configured local directory.
+        """
+        model_repo_id = self.config.get("model_repo_id", "openai/whisper-large-v3")
+        model_revision = self.config.get("model_revision")
+
+        logger.warning(
+            "Local model was not found at %s. Downloading fallback model %s",
+            model_path,
+            model_repo_id,
+        )
+
+        os.makedirs(model_path, exist_ok=True)
+        snapshot_download(
+            repo_id=model_repo_id,
+            local_dir=model_path,
+            local_dir_use_symlinks=False,
+            revision=model_revision,
+        )
+
+        if not self._is_valid_local_model_dir(model_path):
+            raise FileNotFoundError(
+                "Model download completed but the target directory still does not "
+                f"look valid: '{model_path}'. Expected to find '{os.path.join(model_path, 'config.json')}'."
+            )
+
+        logger.info(
+            "Downloaded Whisper model %s into %s",
+            model_repo_id,
+            model_path,
+        )
+        return model_path

    def transcribe(self, audio_path: str, return_timestamps: bool = None,
                   language: str = None, temperature: float = None,
@@ -209,7 +294,7 @@ class WhisperTranscriber:
        if temperature is None:
            temperature = self.temperature

-        logger.info("Начало транскрибации файла: %s", audio_path)
+        logger.info("Starting transcription for file: %s", audio_path)

        try:
            # Загрузка аудио в формате numpy array
@@ -236,7 +321,7 @@ class WhisperTranscriber:
            # Если временные метки не запрошены, возвращаем только текст
            if not return_timestamps:
                transcribed_text = result.get("text", "")
-                logger.info("Транскрибация завершена: получено %s символов текста", len(transcribed_text))
+                logger.info("Transcription completed: produced %s text characters", len(transcribed_text))
                return transcribed_text
            
            # Если временные метки запрошены, обрабатываем и форматируем результат
@@ -268,9 +353,9 @@ class WhisperTranscriber:
                        "text": text
                    })
            else:
-                logger.warning("Временные метки запрошены, но не найдены в результате транскрибации")
+                logger.warning("Timestamps were requested but not found in the transcription result")
            
-            logger.info("Транскрибация с временными метками завершена: получено %s сегментов", len(segments))
+            logger.info("Timestamped transcription completed: produced %s segments", len(segments))
            
            # Возвращаем словарь с сегментами и полным текстом
            return {
@@ -279,8 +364,8 @@ class WhisperTranscriber:
            }
            
        except Exception as e:
-            logger.error("Ошибка в процессе транскрибации аудиофайла '%s': %s", audio_path, e)
-            logger.error("Тип исключения: %s", type(e).__name__)
+            logger.error("Error while transcribing audio file '%s': %s", audio_path, e)
+            logger.error("Exception type: %s", type(e).__name__)
            logger.error("Traceback: %s", traceback.format_exc())
            raise

@@ -303,7 +388,7 @@ class WhisperTranscriber:
            - Если return_timestamps=True: словарь с ключами "segments" и "text"
        """
        start_time = time.time()
-        logger.info("Начало обработки файла: %s", input_path)
+        logger.info("Starting file processing: %s", input_path)
        
        temp_files = []
        
@@ -317,17 +402,17 @@ class WhisperTranscriber:
                                     prompt=prompt)
            
            elapsed_time = time.time() - start_time
-            logger.info("Обработка и транскрибация завершены за %.2f секунд", elapsed_time)
+            logger.info("Processing and transcription finished in %.2f seconds", elapsed_time)
            
            return result
            
        except Exception as e:
            elapsed_time = time.time() - start_time
-            logger.error("Ошибка при обработке файла '%s' через %.2f секунд: %s", input_path, elapsed_time, e)
-            logger.error("Тип исключения: %s", type(e).__name__)
+            logger.error("Error while processing file '%s' after %.2f seconds: %s", input_path, elapsed_time, e)
+            logger.error("Exception type: %s", type(e).__name__)
            logger.error("Traceback: %s", traceback.format_exc())
            raise
            
        finally:
            # Очистка временных файлов
-            cleanup_temp_files(temp_files)
+            cleanup_temp_files(temp_files)
@@ -50,8 +50,8 @@ class TranscriptionService:
            try:
                duration = get_audio_duration(file_path)
            except Exception as e:
-                logger.error("Ошибка при определении длительности файла: %s", e)
-                return {"error": f"Не удалось определить длительность аудиофайла: {e}"}, 500
+                logger.error("Failed to determine file duration: %s", e)
+                return {"error": f"Failed to determine audio duration: {e}"}, 500

            start_time = time.time()
            result = self.transcriber.process_file(
@@ -86,6 +86,6 @@ class TranscriptionService:
            return response, 200

        except Exception as e:
-            logger.error("Ошибка при транскрибации файла '%s': %s", filename, e)
+            logger.error("Failed to transcribe file '%s': %s", filename, e)
            logger.error("Traceback: %s", traceback.format_exc())
            return {"error": str(e)}, 500