Override from_pretrained to support wandb artifacts.

src/dalle_mini/model/modeling.py

@@ -44,6 +44,7 @@ from transformers.models.bart.modeling_flax_bart import (
     FlaxBartPreTrainedModel,
 )
 from transformers.utils import logging
+import wandb
 
 from .configuration import DalleBartConfig
 
@@ -561,3 +562,18 @@ class DalleBart(FlaxBartPreTrainedModel, FlaxBartForConditionalGeneration):
             outputs = outputs[:1] + (unfreeze(past["cache"]),) + outputs[1:]
 
         return outputs
+
+    @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:
+            # wandb artifact
+            artifact = wandb.Api().artifact(pretrained_model_name_or_path)
+            
+            # we download everything, including opt_state, so we can resume training if needed
+            # see also: #120
+            pretrained_model_name_or_path = artifact.download()
+            
+        return super(DalleBart, cls).from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs)