Tokenizer, config, model can be loaded from wandb.

src/dalle_mini/model/__init__.py

@@ -1,2 +1,3 @@
 from .configuration import DalleBartConfig
 from .modeling import DalleBart
+from .tokenizer import DalleBartTokenizer

src/dalle_mini/model/configuration.py

@@ -18,10 +18,12 @@ import warnings
 from transformers.configuration_utils import PretrainedConfig
 from transformers.utils import logging
 
+from .wandb_pretrained import PretrainedFromWandbMixin
+
 logger = logging.get_logger(__name__)
 
 
-class DalleBartConfig(PretrainedConfig):
+class DalleBartConfig(PretrainedFromWandbMixin, PretrainedConfig):
     model_type = "dallebart"
     keys_to_ignore_at_inference = ["past_key_values"]
     attribute_map = {

src/dalle_mini/model/modeling.py

@@ -15,14 +15,12 @@
 """ DalleBart model. """
 
 import math
-import os
 from functools import partial
 from typing import Optional, Tuple
 
 import flax.linen as nn
 import jax
 import jax.numpy as jnp
-import wandb
 from flax.core.frozen_dict import unfreeze
 from flax.linen import make_causal_mask
 from flax.traverse_util import flatten_dict
@@ -48,6 +46,7 @@ from transformers.models.bart.modeling_flax_bart import (
 from transformers.utils import logging
 
 from .configuration import DalleBartConfig
+from .wandb_pretrained import PretrainedFromWandbMixin
 
 logger = logging.get_logger(__name__)
 
@@ -421,7 +420,9 @@ class FlaxBartForConditionalGenerationModule(FlaxBartForConditionalGenerationMod
         )
 
 
-class DalleBart(FlaxBartPreTrainedModel, FlaxBartForConditionalGeneration):
+class DalleBart(
+    PretrainedFromWandbMixin, FlaxBartPreTrainedModel, FlaxBartForConditionalGeneration
+):
     """
     Edits:
     - renamed from FlaxBartForConditionalGeneration
@@ -563,24 +564,3 @@ 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 and not os.path.isdir(
-            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()
-
-        model = super(DalleBart, cls).from_pretrained(
-            pretrained_model_name_or_path, *model_args, **kwargs
-        )
-        model.config.resolved_name_or_path = pretrained_model_name_or_path
-        return model

src/dalle_mini/model/tokenizer.py

@@ -0,0 +1,11 @@
+""" DalleBart tokenizer """
+from transformers import BartTokenizer
+from transformers.utils import logging
+
+from .wandb_pretrained import PretrainedFromWandbMixin
+
+logger = logging.get_logger(__name__)
+
+
+class DalleBartTokenizer(PretrainedFromWandbMixin, BartTokenizer):
+    pass

src/dalle_mini/model/wandb_pretrained.py

@@ -0,0 +1,20 @@
+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
+            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
+        )