Load from wandb artifact (#121)

Load model & tokenizer from artifacts.
Fixes #116

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

@@ -46,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__)
 
@@ -419,7 +420,9 @@ class FlaxBartForConditionalGenerationModule(FlaxBartForConditionalGenerationMod
         )
 
 
-class DalleBart(FlaxBartPreTrainedModel, FlaxBartForConditionalGeneration):
+class DalleBart(
+    PretrainedFromWandbMixin, FlaxBartPreTrainedModel, FlaxBartForConditionalGeneration
+):
     """
     Edits:
     - renamed from FlaxBartForConditionalGeneration

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,21 @@
+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
+        )

tools/train/train.py

@@ -44,7 +44,7 @@ from tqdm import tqdm
 from transformers import AutoTokenizer, HfArgumentParser
 
 from dalle_mini.data import Dataset
-from dalle_mini.model import DalleBart, DalleBartConfig
+from dalle_mini.model import DalleBart, DalleBartConfig, DalleBartTokenizer
 
 logger = logging.getLogger(__name__)
 
@@ -58,8 +58,9 @@ class ModelArguments:
     model_name_or_path: Optional[str] = field(
         default=None,
         metadata={
-            "help": "The model checkpoint for weights initialization."
-            "Don't set if you want to train a model from scratch."
+            "help": "The model checkpoint for weights initialization. "
+            "Don't set if you want to train a model from scratch. "
+            "W&B artifact references are supported in addition to the sources supported by `PreTrainedModel`."
         },
     )
     config_name: Optional[str] = field(
@@ -482,13 +483,15 @@ def main():
 
         # load model
         model = DalleBart.from_pretrained(
-            artifact_dir, dtype=getattr(jnp, model_args.dtype), abstract_init=True
+            artifact_dir,
+            dtype=getattr(jnp, model_args.dtype),
+            abstract_init=True,
         )
         # avoid OOM on TPU: see https://github.com/google/flax/issues/1658
         print(model.params)
 
         # load tokenizer
-        tokenizer = AutoTokenizer.from_pretrained(
+        tokenizer = DalleBartTokenizer.from_pretrained(
             artifact_dir,
             use_fast=True,
         )
@@ -498,7 +501,7 @@ def main():
         if model_args.config_name:
             config = DalleBartConfig.from_pretrained(model_args.config_name)
         else:
-            config = DalleBartConfig.from_pretrained(model_args.model_name_or_path)
+            config = None
 
         # Load or create new model
         if model_args.model_name_or_path:
@@ -524,7 +527,7 @@ def main():
                 model_args.tokenizer_name, use_fast=True
             )
         else:
-            tokenizer = AutoTokenizer.from_pretrained(
+            tokenizer = DalleBartTokenizer.from_pretrained(
                 model_args.model_name_or_path,
                 use_fast=True,
             )