Repository: huggingface/transformers

Issue:
Loading FlaxHybridCLIP trained model 
## Environment info
<!-- You can run the command `transformers-cli env` and copy-and-paste its output below.
     Don't forget to fill out the missing fields in that output! -->

- `transformers` version: 4.9.0.dev0
- Platform: Linux-5.4.0-1043-gcp-x86_64-with-glibc2.29
- Python version: 3.8.10
- PyTorch version (GPU?): 1.9.0+cu102 (False)
- Tensorflow version (GPU?): 2.5.0 (False)
- Flax version (CPU?/GPU?/TPU?): 0.3.4 (tpu)
- Jax version: 0.2.16
- JaxLib version: 0.1.68
- Using GPU in script?: N/A
- Using distributed or parallel set-up in script?: N/A

Models:

- FlaxHybridCLIP

## Information

I am not sure about how to load a trained  FlaxHybridCLIP model from a folder. We trained using [this](https://github.com/huggingface/transformers/tree/master/examples/research_projects/jax-projects/hybrid_clip).

I tried:  `FlaxHybridCLIP.from_text_vision_pretrained(PATH_TRAINED_MODEL, PATH_TRAINED_MODEL)`, but I got the following error: 

```Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/vinid/transformers/examples/research_projects/jax-projects/hybrid_clip/modeling_hybrid_clip.py", line 333, in from_text_vision_pretrained
    text_config = AutoConfig.from_pretrained(text_model_name_or_path)
  File "/home/raphaelp/transformers/src/transformers/models/auto/configuration_auto.py", line 452, in from_pretrained
    config_class = CONFIG_MAPPING[config_dict["model_type"]]
KeyError: 'hybrid-clip'
```

The folder (PATH_TRAINED_MODEL) contains the two following files:

- `config.json`
- `flax_model.msgpack`

Thank you :) :) 




Discussion:
Sorry, totally forgot to tag you @patrickvonplaten, @patil-suraj :)

Thanks :) 
Hi @vinid 

To load the trained model or if you save the `FlaxHybridCLIP` model using the `save_pretrained` method, you could directly use the `FlaxHybridCLIP.from_pretrained` method. Let me know If that works.
Thanks for your reply! :)

I am getting a new error now, seems like it is looking for something in the configs

```
>>> FlaxHybridCLIP.from_pretrained(PATH_TRAINED_MODEL)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/giuseppe/transformers/src/transformers/modeling_flax_utils.py", line 265, in from_pretrained
    config, model_kwargs = cls.config_class.from_pretrained(
AttributeError: 'NoneType' object has no attribute 'from_pretrained'
```

Aah, I see. That's actually a bug, I will work on a fix.

As a workaround could you try loading the config and then pass it to the `from_pretrained` ?
```python
config = HybridCLIPConfig.from_pretrained("...")
model = FlaxHybridCLIP.from_pretrained("...", config=config)
```

Resolving patch:
diff --git a/examples/research_projects/jax-projects/hybrid_clip/configuration_hybrid_clip.py b/examples/research_projects/jax-projects/hybrid_clip/configuration_hybrid_clip.py
--- a/examples/research_projects/jax-projects/hybrid_clip/configuration_hybrid_clip.py
+++ b/examples/research_projects/jax-projects/hybrid_clip/configuration_hybrid_clip.py
@@ -25,31 +25,58 @@ class HybridCLIPConfig(PretrainedConfig):
             Dimentionality of text and vision projection layers.
         kwargs (`optional`):
             Dictionary of keyword arguments.
+
+    Examples::
+
+        >>> from transformers import BertConfig, CLIPConfig, HybridCLIPConfig, FlaxHybridCLIP
+
+        >>> # Initializing a BERT and CLIP configuration
+        >>> config_text = BertConfig()
+        >>> config_vision = CLIPConfig()
+
+        >>> config = HybridCLIPConfig.from_text_vision_configs(config_text, config_vision, projection_dim=512)
+
+        >>> # Initializing a BERT and CLIPVision model
+        >>> model = EncoderDecoderModel(config=config)
+
+        >>> # Accessing the model configuration
+        >>> config_text = model.config.text_config
+        >>> config_vision  = model.config.vision_config
+
+        >>> # Saving the model, including its configuration
+        >>> model.save_pretrained('my-model')
+
+        >>> # loading model and config from pretrained folder
+        >>> encoder_decoder_config = HybridCLIPConfig.from_pretrained('my-model')
+        >>> model = FlaxHybridCLIP.from_pretrained('my-model', config=encoder_decoder_config)
     """
 
     model_type = "hybrid-clip"
     is_composition = True
 
-    def __init__(self, text_config_dict, vision_config_dict, projection_dim=512, **kwargs):
+    def __init__(self, projection_dim=512, **kwargs):
         super().__init__(**kwargs)
 
-        if text_config_dict is None:
-            raise ValueError("`text_config_dict` can not be `None`.")
+        if "text_config" not in kwargs:
+            raise ValueError("`text_config` can not be `None`.")
+
+        if "vision_config" not in kwargs:
+            raise ValueError("`vision_config` can not be `None`.")
 
-        if vision_config_dict is None:
-            raise ValueError("`vision_config_dict` can not be `None`.")
+        text_config = kwargs.pop("text_config")
+        vision_config = kwargs.pop("vision_config")
 
-        text_model_type = text_config_dict.pop("model_type")
-        vision_model_type = vision_config_dict.pop("model_type")
+        text_model_type = text_config.pop("model_type")
+        vision_model_type = vision_config.pop("model_type")
 
         from transformers import AutoConfig
 
-        self.text_config = AutoConfig.for_model(text_model_type, **text_config_dict)
+        self.text_config = AutoConfig.for_model(text_model_type, **text_config)
 
         if vision_model_type == "clip":
-            self.vision_config = AutoConfig.for_model(vision_model_type, **vision_config_dict).vision_config
+            self.vision_config = AutoConfig.for_model(vision_model_type, **vision_config).vision_config
         else:
-            self.vision_config = AutoConfig.for_model(vision_model_type, **vision_config_dict)
+            self.vision_config = AutoConfig.for_model(vision_model_type, **vision_config)
 
         self.projection_dim = projection_dim
         self.initializer_factor = 1.0
@@ -64,7 +91,7 @@ def from_text_vision_configs(cls, text_config: PretrainedConfig, vision_config:
             :class:`HybridCLIPConfig`: An instance of a configuration object
         """
 
-        return cls(text_config_dict=text_config.to_dict(), vision_config_dict=vision_config.to_dict(), **kwargs)
+        return cls(text_config=text_config.to_dict(), vision_config=vision_config.to_dict(), **kwargs)
 
     def to_dict(self):
         """
diff --git a/examples/research_projects/jax-projects/hybrid_clip/modeling_hybrid_clip.py b/examples/research_projects/jax-projects/hybrid_clip/modeling_hybrid_clip.py
--- a/examples/research_projects/jax-projects/hybrid_clip/modeling_hybrid_clip.py
+++ b/examples/research_projects/jax-projects/hybrid_clip/modeling_hybrid_clip.py
@@ -123,7 +123,7 @@ def __call__(
 
 
 class FlaxHybridCLIP(FlaxPreTrainedModel):
-    config: HybridCLIPConfig
+    config_class = HybridCLIPConfig
     module_class = FlaxHybridCLIPModule
 
     def __init__(
@@ -304,6 +304,58 @@ def from_text_vision_pretrained(
         *model_args,
         **kwargs,
     ) -> FlaxPreTrainedModel:
+        """
+        Params:
+            text_model_name_or_path (:obj: `str`, `optional`):
+                Information necessary to initiate the text model. Can be either:
+
+                    - A string, the `model id` of a pretrained model hosted inside a model repo on huggingface.co.
+                      Valid model ids can be located at the root-level, like ``bert-base-uncased``, or namespaced under
+                      a user or organization name, like ``dbmdz/bert-base-german-cased``.
+                    - A path to a `directory` containing model weights saved using
+                      :func:`~transformers.FlaxPreTrainedModel.save_pretrained`, e.g., ``./my_model_directory/``.
+                    - A path or url to a `PyTorch checkpoint folder` (e.g, ``./pt_model``). In
+                      this case, ``from_pt`` should be set to :obj:`True` and a configuration object should be provided
+                      as ``config`` argument. This loading path is slower than converting the PyTorch checkpoint in
+                      a Flax model using the provided conversion scripts and loading the Flax model afterwards.
+
+            vision_model_name_or_path (:obj: `str`, `optional`, defaults to `None`):
+                Information necessary to initiate the vision model. Can be either:
+
+                    - A string, the `model id` of a pretrained model hosted inside a model repo on huggingface.co.
+                      Valid model ids can be located at the root-level, like ``bert-base-uncased``, or namespaced under
+                      a user or organization name, like ``dbmdz/bert-base-german-cased``.
+                    - A path to a `directory` containing model weights saved using
+                      :func:`~transformers.FlaxPreTrainedModel.save_pretrained`, e.g., ``./my_model_directory/``.
+                    - A path or url to a `PyTorch checkpoint folder` (e.g, ``./pt_model``). In
+                      this case, ``from_pt`` should be set to :obj:`True` and a configuration object should be provided
+                      as ``config`` argument. This loading path is slower than converting the PyTorch checkpoint in
+                      a Flax model using the provided conversion scripts and loading the Flax model afterwards.
+
+            model_args (remaining positional arguments, `optional`):
+                All remaning positional arguments will be passed to the underlying model's ``__init__`` method.
+
+            kwargs (remaining dictionary of keyword arguments, `optional`):
+                Can be used to update the configuration object (after it being loaded) and initiate the model (e.g.,
+                :obj:`output_attentions=True`).
+
+                - To update the text configuration, use the prefix `text_` for each configuration parameter.
+                - To update the vision configuration, use the prefix `vision_` for each configuration parameter.
+                - To update the parent model configuration, do not use a prefix for each configuration parameter.
+
+                Behaves differently depending on whether a :obj:`config` is provided or automatically loaded.
+
+        Example::
+
+            >>> from transformers import FlaxHybridCLIP
+            >>> # initialize a model from pretrained BERT and CLIP models. Note that the projection layers will be randomly initialized.
+            >>> # If using CLIP's vision model the vision projection layer will be initialized using pre-trained weights
+            >>> model = FlaxHybridCLIP.from_text_vision_pretrained('bert-base-uncased', 'openai/clip-vit-base-patch32')
+            >>> # saving model after fine-tuning
+            >>> model.save_pretrained("./bert-clip")
+            >>> # load fine-tuned model
+            >>> model = FlaxHybridCLIP.from_pretrained("./bert-clip")
+        """
 
         kwargs_text = {
             argument[len("text_") :]: value for argument, value in kwargs.items() if argument.startswith("text_")
@@ -333,9 +385,7 @@ def from_text_vision_pretrained(
                 text_config = AutoConfig.from_pretrained(text_model_name_or_path)
                 kwargs_text["config"] = text_config
 
-            text_model = FlaxAutoModel.from_pretrained(
-                text_model_name_or_path, *model_args, from_pt=True, **kwargs_text
-            )
+            text_model = FlaxAutoModel.from_pretrained(text_model_name_or_path, *model_args, **kwargs_text)
 
         vision_model = kwargs_vision.pop("model", None)
         if vision_model is None:
diff --git a/examples/research_projects/jax-projects/hybrid_clip/run_hybrid_clip.py b/examples/research_projects/jax-projects/hybrid_clip/run_hybrid_clip.py
--- a/examples/research_projects/jax-projects/hybrid_clip/run_hybrid_clip.py
+++ b/examples/research_projects/jax-projects/hybrid_clip/run_hybrid_clip.py
@@ -87,6 +87,10 @@ class ModelArguments:
             "Don't set if you want to train a model from scratch."
         },
     )
+    from_pt: bool = field(
+        default=True,
+        metadata={"help": "whether to load the text and vision model using PyTorch checkpoints."},
+    )
     config_name: Optional[str] = field(
         default=None, metadata={"help": "Pretrained config name or path if not the same as model_name"}
     )
@@ -332,6 +336,8 @@ def main():
         model_args.vision_model_name_or_path,
         seed=training_args.seed,
         dtype=getattr(jnp, model_args.dtype),
+        text_from_pt=model_args.from_pt,
+        vision_from_pt=model_args.from_pt,
     )
     config = model.config
     # set seed for torch dataloaders


Test patch: