File size: 1,954 Bytes
4974c8b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 |
from huggingface_hub import PyTorchModelHubMixin
from huggingface_hub.constants import PYTORCH_WEIGHTS_NAME
from huggingface_hub.file_download import hf_hub_download
from unifiedmodel import RRUM
import os
import torch
class YoutubeVideoSimilarityModel(RRUM, PyTorchModelHubMixin):
"""
Hugging Face `PyTorchModelHubMixin` wrapper for RegretsReporter `RRUM` model.
This allows loading, using, and saving the model from Hugging Face model hub
with default Hugging Face methods `from_pretrained` and `save_pretrained`.
"""
@classmethod
def _from_pretrained(
cls,
model_id,
revision,
cache_dir,
force_download,
proxies,
resume_download,
local_files_only,
use_auth_token,
map_location="cpu",
strict=False,
**model_kwargs,
):
map_location = torch.device(map_location)
if os.path.isdir(model_id):
print("Loading weights from local directory")
model_file = os.path.join(model_id, PYTORCH_WEIGHTS_NAME)
else:
model_file = hf_hub_download(
repo_id=model_id,
filename=PYTORCH_WEIGHTS_NAME,
revision=revision,
cache_dir=cache_dir,
force_download=force_download,
proxies=proxies,
resume_download=resume_download,
use_auth_token=use_auth_token,
local_files_only=local_files_only,
)
# convert Huggingface config to RRUM acceptable input parameters
if "config" in model_kwargs:
model_kwargs = {**model_kwargs["config"], **model_kwargs}
del model_kwargs["config"]
model = cls(**model_kwargs)
state_dict = torch.load(model_file, map_location=map_location)
model.load_state_dict(state_dict, strict=strict)
model.eval()
return model
|