diff --git a/msclap/CLAPWrapper.py b/msclap/CLAPWrapper.py index 542ab6f..c899be3 100644 --- a/msclap/CLAPWrapper.py +++ b/msclap/CLAPWrapper.py @@ -1,3 +1,5 @@ +from __future__ import annotations + from pathlib import Path import warnings warnings.filterwarnings("ignore") @@ -16,6 +18,7 @@ import torch import argparse import yaml import sys +from huggingface_hub.file_download import hf_hub_download logging.set_verbosity_error() @@ -23,15 +26,30 @@ class CLAPWrapper(): """ A class for interfacing CLAP model. """ + model_repo = "microsoft/msclap" + model_name = { + '2022': 'CLAP_weights_2022.pth', + '2023': 'CLAP_weights_2023.pth', + 'clapcap': 'clapcap_weights_2023.pth' + } + + def __init__(self, model_fp: Path | str | None = None, version: str = '2023', use_cuda=False): + # Check if version is supported + self.supported_versions = self.model_name.keys() + if version not in self.supported_versions: + raise ValueError(f"The version {version} is not supported. The supported versions are {str(self.supported_versions)}") - def __init__(self, model_fp, version, use_cuda=False): - self.supported_versions = ['2022', '2023', 'clapcap'] self.np_str_obj_array_pattern = re.compile(r'[SaUO]') self.file_path = os.path.realpath(__file__) self.default_collate_err_msg_format = ( "default_collate: batch must contain tensors, numpy arrays, numbers, " "dicts or lists; found {}") - self.config_as_str = self.get_config_path(version) + self.config_as_str = (Path(__file__).parent / f"configs/config_{version}.yml").read_text() + + # Automatically download model if not provided + if not model_fp: + model_fp = hf_hub_download(self.model_repo, self.model_name[version]) + self.model_fp = model_fp self.use_cuda = use_cuda if 'clapcap' in version: @@ -39,12 +57,6 @@ class CLAPWrapper(): else: self.clap, self.tokenizer, self.args = self.load_clap() - def get_config_path(self, version): - if version in self.supported_versions: - return (Path(__file__).parent / f"configs/config_{version}.yml").read_text() - else: - raise ValueError(f"The specific version is not supported. The supported versions are {str(self.supported_versions)}") - def read_config_as_args(self,config_path,args=None,is_config_str=False): return_dict = {}