feat: use_artifact if run existing

src/dalle_mini/model/configuration.py

@@ -18,7 +18,7 @@ import warnings
 from transformers.configuration_utils import PretrainedConfig
 from transformers.utils import logging
 
-from .wandb_pretrained import PretrainedFromWandbMixin
+from .utils import PretrainedFromWandbMixin
 
 logger = logging.get_logger(__name__)
 

src/dalle_mini/model/modeling.py

@@ -46,7 +46,7 @@ from transformers.models.bart.modeling_flax_bart import (
 from transformers.utils import logging
 
 from .configuration import DalleBartConfig
-from .wandb_pretrained import PretrainedFromWandbMixin
+from .utils import PretrainedFromWandbMixin
 
 logger = logging.get_logger(__name__)
 

src/dalle_mini/model/tokenizer.py

@@ -2,7 +2,7 @@
 from transformers import BartTokenizer
 from transformers.utils import logging
 
-from .wandb_pretrained import PretrainedFromWandbMixin
+from .utils import PretrainedFromWandbMixin
 
 logger = logging.get_logger(__name__)
 

src/dalle_mini/model/{wandb_pretrained.py → utils.py}

@@ -13,7 +13,10 @@ class PretrainedFromWandbMixin:
             pretrained_model_name_or_path
         ):
             # wandb artifact
-            artifact = wandb.Api().artifact(pretrained_model_name_or_path)
+            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(