From d6c029f708186c4566e05d9479d20e241981caf1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=93=D0=B0=D0=BB=D0=B8=D0=BD=D0=B3=D0=B5=D1=80=20=D0=A0?= =?UTF-8?q?=2E=D0=A1=2E?= Date: Wed, 8 Jul 2026 09:44:48 +0700 Subject: [PATCH] feat: preload Whisper model at startup with disk space check - Add WHISPER_MODEL_SIZE and WHISPER_MIN_FREE_SPACE_BYTES config options - voice.py now preloads/downloads the Whisper model at bot startup - Detailed logging for model presence, disk space, download progress and readiness - If model is already present locally, preload is skipped - If disk space is insufficient, bot fails fast with a clear error - main.py calls preload_model() after proxy setup and before polling - Document voice recognition model download behavior in README --- README.md | 43 +++++++++++++++++++++++ bot/utils/voice.py | 87 ++++++++++++++++++++++++++++++++++++++++++---- config.py | 5 +++ main.py | 9 +++++ 4 files changed, 137 insertions(+), 7 deletions(-) diff --git a/README.md b/README.md index 6447f1e..9b47454 100644 --- a/README.md +++ b/README.md @@ -95,6 +95,10 @@ ROUTERAI_BASE_URL=https://routerai.ru/api/v1 # Внимание: проверка расходует токены API. По умолчанию отключена. AI_HEALTH_CHECK_ENABLED=false AI_HEALTH_CHECK_INTERVAL=600 + +# Настройки модели Whisper для распознавания голоса. +WHISPER_MODEL_SIZE=base +WHISPER_MIN_FREE_SPACE_BYTES=5368709120 ``` Для включения прокси установите `PROXY_ENABLED=true` и укажите корректный `PROXY_URL`. @@ -114,6 +118,45 @@ docker compose up -d --build Модели Whisper сохраняются в `./models`, база данных и логи — в `./bot/data`. +## 🎙️ Распознавание голосовых сообщений + +Бот автоматически распознаёт все голосовые сообщения с помощью локальной модели [faster-whisper](https://github.com/SYSTRAN/faster-whisper). + +### Поведение загрузки модели + +- При запуске бота проверяется наличие модели в `./models/faster-whisper//`. +- Если модель уже скачана — бот сразу готов к работе. +- Если модели нет — бот проверяет свободное место на диске и **скачивает модель автоматически** перед началом polling. +- Для скачивания используется прокси из `.env`, если `PROXY_ENABLED=true`. + +### Требования к месту + +| Модель | Размер модели | Рекомендуемое свободное место | +|--------|---------------|-------------------------------| +| `base` | ~150 MB архив, ~500 MB на диске | 5 GB (`WHISPER_MIN_FREE_SPACE_BYTES=5368709120`) | +| `small` | ~500 MB архив, ~1.5 GB на диске | 5 GB+ | +| `medium` | ~1.5 GB архив, ~5 GB на диске | 10 GB+ | + +> 💡 По умолчанию используется модель `base`. Для смены модели измените `WHISPER_MODEL_SIZE` в `.env`. + +### Логирование + +В логах будет видно: + +``` +Free disk space: 45.23 GB, required: 5.00 GB +Whisper model 'base' found locally at /app/models/faster-whisper/base +Whisper model 'base' is ready +``` + +или при скачивании: + +``` +Whisper model 'base' not found locally. Starting download to /app/models/faster-whisper... +Whisper model 'base' downloaded and loaded in 125.4s +Whisper model 'base' is ready +``` + ## 🧪 Тесты ```bash diff --git a/bot/utils/voice.py b/bot/utils/voice.py index 00e83b6..d11e16a 100644 --- a/bot/utils/voice.py +++ b/bot/utils/voice.py @@ -2,12 +2,14 @@ import asyncio import logging import os import re +import shutil +import time from pathlib import Path from faster_whisper import WhisperModel from aiogram import Bot -from config import PROXY_ENABLED, PROXY_URL +from config import PROXY_ENABLED, PROXY_URL, WHISPER_MODEL_SIZE, WHISPER_MIN_FREE_SPACE_BYTES logger = logging.getLogger(__name__) @@ -25,26 +27,59 @@ def _setup_proxy() -> None: logger.info("Proxy set for model download: %s...", PROXY_URL[:30]) +def _get_model_path() -> Path: + return MODELS_DIR / WHISPER_MODEL_SIZE + + +def _is_model_downloaded() -> bool: + """Check whether the Whisper model files are already present locally.""" + model_path = _get_model_path() + model_bin = model_path / "model.bin" + return model_path.is_dir() and model_bin.is_file() + + +def _get_free_disk_space_bytes(path: Path) -> int: + """Return free disk space in bytes for the filesystem containing path.""" + path.mkdir(parents=True, exist_ok=True) + return shutil.disk_usage(path).free + + def _load_model() -> WhisperModel: """Synchronous model load/download. Must run in a thread.""" MODELS_DIR.mkdir(parents=True, exist_ok=True) VOICE_DIR.mkdir(parents=True, exist_ok=True) - model_path = MODELS_DIR / "base" + model_path = _get_model_path() model_bin = model_path / "model.bin" if not model_path.is_dir() or not model_bin.is_file(): _setup_proxy() - logger.info("Model not found locally, downloading to %s...", MODELS_DIR) - return WhisperModel( - "base", + logger.info( + "Whisper model '%s' not found locally. Starting download to %s...", + WHISPER_MODEL_SIZE, + MODELS_DIR, + ) + start = time.monotonic() + model = WhisperModel( + WHISPER_MODEL_SIZE, device="cpu", cpu_threads=4, compute_type="int8", download_root=str(MODELS_DIR), ) + elapsed = time.monotonic() - start + logger.info( + "Whisper model '%s' downloaded and loaded in %.1fs", + WHISPER_MODEL_SIZE, + elapsed, + ) + return model - logger.info("Loading model from %s...", model_path) + logger.info( + "Loading Whisper model '%s' from %s...", + WHISPER_MODEL_SIZE, + model_path, + ) return WhisperModel( str(model_path), device="cpu", @@ -53,6 +88,44 @@ def _load_model() -> WhisperModel: ) +async def preload_model() -> None: + """Download/load the Whisper model at bot startup if it is not already loaded.""" + global _model + async with _model_lock: + if _model is not None: + logger.info("Whisper model is already loaded, skipping preload") + return + + if _is_model_downloaded(): + logger.info( + "Whisper model '%s' found locally at %s", + WHISPER_MODEL_SIZE, + _get_model_path(), + ) + else: + free_space = _get_free_disk_space_bytes(MODELS_DIR) + required_space = WHISPER_MIN_FREE_SPACE_BYTES + logger.info( + "Free disk space: %.2f GB, required: %.2f GB", + free_space / (1024 ** 3), + required_space / (1024 ** 3), + ) + if free_space < required_space: + raise RuntimeError( + f"Not enough disk space to download Whisper model '{WHISPER_MODEL_SIZE}'. " + f"Free: {free_space / (1024 ** 3):.2f} GB, " + f"required: {required_space / (1024 ** 3):.2f} GB" + ) + logger.info( + "Whisper model '%s' will be downloaded at startup", + WHISPER_MODEL_SIZE, + ) + + loop = asyncio.get_running_loop() + _model = await loop.run_in_executor(None, _load_model) + logger.info("Whisper model '%s' is ready", WHISPER_MODEL_SIZE) + + async def _get_model() -> WhisperModel: """Thread-safe lazy initializer for the Whisper model.""" global _model @@ -61,7 +134,7 @@ async def _get_model() -> WhisperModel: if _model is None: loop = asyncio.get_running_loop() _model = await loop.run_in_executor(None, _load_model) - logger.info("Whisper model loaded") + logger.info("Whisper model '%s' loaded on demand", WHISPER_MODEL_SIZE) return _model diff --git a/config.py b/config.py index ad44a4b..0f591ea 100644 --- a/config.py +++ b/config.py @@ -28,6 +28,11 @@ AI_BLOCK_DEFAULT_DURATION = 86400 AI_HEALTH_CHECK_ENABLED = os.getenv("AI_HEALTH_CHECK_ENABLED", "false").lower() == "true" AI_HEALTH_CHECK_INTERVAL = int(os.getenv("AI_HEALTH_CHECK_INTERVAL", "600")) +# Whisper model settings. +# The base model requires ~2.5 GB for download; 5 GB free space is recommended. +WHISPER_MODEL_SIZE = os.getenv("WHISPER_MODEL_SIZE", "base") +WHISPER_MIN_FREE_SPACE_BYTES = int(os.getenv("WHISPER_MIN_FREE_SPACE_BYTES", str(5 * 1024 * 1024 * 1024))) + ACCESS_KEY = os.getenv("ACCESS_KEY", "").strip() SECRET_KEY = os.getenv("SECRET_KEY", "").strip() BUCKET_NAME = os.getenv("BUCKET_NAME", "").strip() diff --git a/main.py b/main.py index 2b4954e..5f5d529 100644 --- a/main.py +++ b/main.py @@ -13,6 +13,7 @@ from bot.routers import moderation, layout, weather, ai, voice, yadisk, dialogue from bot.utils.database import init_db, save_chat_user from bot.utils.ai_client import start_model_health_check from bot.utils.logging_config import setup_logging +from bot.utils.voice import preload_model from bot.setup_commands import setup_bot_commands from config import PROXY_ENABLED @@ -75,6 +76,14 @@ async def main(): health_check_task.cancel() return + try: + await preload_model() + logging.info("Модель Whisper готова к работе") + except Exception as exc: + logging.error("Не удалось загрузить модель Whisper: %s", exc) + health_check_task.cancel() + return + try: await dp.start_polling(bot) finally: