boris's picture
feat: use_artifact if run existing
a5ed112
raw
history blame
844 Bytes
import os
import wandb
class PretrainedFromWandbMixin:
@classmethod
def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
"""
Initializes from a wandb artifact, or delegates loading to the superclass.
"""
if ":" in pretrained_model_name_or_path and not os.path.isdir(
pretrained_model_name_or_path
):
# wandb artifact
if wandb.run is not None:
artifact = wandb.run.use_artifact(pretrained_model_name_or_path)
else:
artifact = wandb.Api().artifact(pretrained_model_name_or_path)
pretrained_model_name_or_path = artifact.download()
return super(PretrainedFromWandbMixin, cls).from_pretrained(
pretrained_model_name_or_path, *model_args, **kwargs
)