src/services/model_manager.py - master - Anna_brv/ASR-Service

src/services/model_manager.py
77 строк2 KB
Anna
first_commit
30 апр 2026, 21:52
30 апр 2026, 21:52283885c
"""Менеджер моделей для управления загрузкой и выгрузкой моделей."""

import os
from pathlib import Path

from loguru import logger

from .asr_service import ASRService


class ModelManager:
    """Менеджер для управления моделями."""

    def __init__(
        self,
        cache_dir: Path,
        default_model_name: str,
        diarization_model_name: str,
        hf_token: str | None = None,
    ):
        self.cache_dir = cache_dir
        self.default_model_name = default_model_name
        self.diarization_model_name = diarization_model_name
        self.hf_token = hf_token or os.getenv("ASR_HF_TOKEN")

        self._asr_service = ASRService(
            model_cache_dir=cache_dir,
            default_model_name=default_model_name,
            diarization_model_name=diarization_model_name,
            hf_token=self.hf_token,
        )

    def load_model(self, model_name: str) -> None:
        """
        Загрузить модель в кэш.

        Args:
            model_name: Название модели.
        """
        logger.info(f"Загрузка модели {model_name}...")
        self._asr_service.load_model(model_name)

    def unload_model(self, model_name: str) -> None:
        """
        Выгрузить модель из памяти.

        Args:
            model_name: Название модели.
        """
        logger.info(f"Выгрузка модели {model_name}...")
        self._asr_service.unload_model(model_name)

    def is_model_loaded(self, model_name: str) -> bool:
        """
        Проверить, загружена ли модель.

        Args:
            model_name: Название модели.

        Returns:
            True если модель загружена.
        """
        return self._asr_service.is_model_loaded(model_name)

    def get_loaded_models(self) -> list[str]:
        """
        Получить список загруженных моделей.

        Returns:
            Список названий загруженных моделей.
        """
        return list(self._asr_service._models.keys())

    def preload_default_model(self) -> None:
        """Предзагрузить модель по умолчанию."""
        logger.info(f"Предзагрузка модели по умолчанию: {self.default_model_name}")
        self.load_model(self.default_model_name)