load_model.py

from pathlib import Path
from typing import Optional, Union


def _default_model_path() -> str:
    return str(Path(__file__).resolve().parent)


def load_pretrained_encoder(
    model_id_or_path: Optional[Union[str, Path]] = None,
    device: str = "cpu",
    dtype=None,
):
    from transformers import AutoModel

    kwargs = {"trust_remote_code": True}
    if dtype is not None:
        kwargs["torch_dtype"] = dtype
    model = AutoModel.from_pretrained(str(model_id_or_path or _default_model_path()), **kwargs)
    return model.to(device).eval()


def load_audio_processor(model_id_or_path: Optional[Union[str, Path]] = None, device: str = "cpu", **kwargs):
    try:
        from .processing_bat import BatAudioProcessor
    except ImportError:
        from processing_bat import BatAudioProcessor

    processor = BatAudioProcessor.from_pretrained(model_id_or_path or _default_model_path(), **kwargs)
    return processor.to(device)