Repository: huggingface/optimum

Issue:
Add Model loading from subfolders similar to transformers
### Feature request

Transformers is adding support for loading models from subfolders: https://github.com/huggingface/transformers/pull/18184

```python
from transformers import BertModel
BertModel.from_pretrained("fusing/latent-diffusion-text2im-large", revision="d5eab56", subfolder="bert")
```


### Motivation

Makes it easier to manage optimized models in existing repositories

### Your contribution

.


Discussion:
Hi @philschmid,
You mean for `ORTModel` right?
yes! 

Resolving patch:
diff --git a/optimum/modeling_base.py b/optimum/modeling_base.py
--- a/optimum/modeling_base.py
+++ b/optimum/modeling_base.py
@@ -40,6 +40,9 @@
             standard cache should not be used.
         local_files_only(`bool`, *optional*, defaults to `False`):
             Whether or not to only look at local files (i.e., do not try to download the model).
+        subfolder (`str`, *optional*, defaults to `""`):
+            In case the relevant files are located inside a subfolder of the model repo either locally or on huggingface.co, you can
+            specify the folder name here.
 """
 
 
@@ -183,6 +186,7 @@ def _from_pretrained(
         revision: Optional[Union[str, None]] = None,
         force_download: bool = False,
         cache_dir: Optional[str] = None,
+        subfolder: Optional[str] = "",
         **kwargs,
     ):
         """Overwrite this method in subclass to define how to load your model from pretrained"""
@@ -197,6 +201,7 @@ def from_pretrained(
         force_download: bool = False,
         use_auth_token: Optional[str] = None,
         cache_dir: Optional[str] = None,
+        subfolder: Optional[str] = "",
         **model_kwargs,
     ):
         """
@@ -207,7 +212,9 @@ def from_pretrained(
         if len(str(model_id).split("@")) == 2:
             model_id, revision = model_id.split("@")
 
-        if os.path.isdir(model_id) and CONFIG_NAME in os.listdir(model_id):
+        if os.path.isdir(os.path.join(model_id, subfolder)) and CONFIG_NAME in os.listdir(
+            os.path.join(model_id, subfolder)
+        ):
             config = AutoConfig.from_pretrained(os.path.join(model_id, CONFIG_NAME))
         else:
             try:
@@ -217,7 +224,21 @@ def from_pretrained(
                     cache_dir=cache_dir,
                     force_download=force_download,
                     use_auth_token=use_auth_token,
+                    subfolder=subfolder,
                 )
+            except OSError:
+                # if config not found in subfolder, search for it at the top level
+                if subfolder != "":
+                    config = AutoConfig.from_pretrained(
+                        pretrained_model_name_or_path=model_id,
+                        revision=revision,
+                        cache_dir=cache_dir,
+                        force_download=force_download,
+                        use_auth_token=use_auth_token,
+                    )
+                    logger.info(
+                        f"config.json not found in the specified subfolder {subfolder}. Using the top level config.json."
+                    )
             except requests.exceptions.RequestException:
                 logger.warning("config.json NOT FOUND in HuggingFace Hub")
                 config = None
@@ -232,6 +253,7 @@ def from_pretrained(
                 cache_dir=cache_dir,
                 force_download=force_download,
                 use_auth_token=use_auth_token,
+                subfolder=subfolder,
                 **model_kwargs,
             )
         else:
@@ -241,6 +263,7 @@ def from_pretrained(
                 cache_dir=cache_dir,
                 force_download=force_download,
                 use_auth_token=use_auth_token,
+                subfolder=subfolder,
                 **model_kwargs,
             )
 
@@ -252,6 +275,7 @@ def _from_transformers(
         revision: Optional[Union[str, None]] = None,
         force_download: bool = False,
         cache_dir: Optional[str] = None,
+        subfolder: Optional[str] = "",
         **kwargs,
     ):
         """Overwrite this method in subclass to define how to load your model from vanilla transformers model"""
diff --git a/optimum/onnxruntime/modeling_ort.py b/optimum/onnxruntime/modeling_ort.py
--- a/optimum/onnxruntime/modeling_ort.py
+++ b/optimum/onnxruntime/modeling_ort.py
@@ -175,7 +175,7 @@ def load_model(
 
         Arguments:
             path (`str` or `Path`):
-                Directory from which to load the model.
+                Path of the ONNX model.
             provider (`str`, *optional*):
                 ONNX Runtime provider to use for loading the model. See https://onnxruntime.ai/docs/execution-providers/
                 for possible providers. Defaults to `CPUExecutionProvider`.
@@ -230,6 +230,7 @@ def from_pretrained(
         force_download: bool = False,
         use_auth_token: Optional[str] = None,
         cache_dir: Optional[str] = None,
+        subfolder: Optional[str] = "",
         provider: Optional[str] = "CPUExecutionProvider",
         session_options: Optional[ort.SessionOptions] = None,
         provider_options: Optional[Dict] = None,
@@ -255,6 +256,7 @@ def from_pretrained(
             force_download,
             use_auth_token,
             cache_dir,
+            subfolder,
             provider=provider,
             session_options=session_options,
             provider_options=provider_options,
@@ -271,6 +273,7 @@ def _from_pretrained(
         force_download: bool = False,
         cache_dir: Optional[str] = None,
         file_name: Optional[str] = None,
+        subfolder: Optional[str] = "",
         **kwargs,
     ):
         """
@@ -301,9 +304,9 @@ def _from_pretrained(
         config = kwargs.pop("config", {})
         model_file_name = file_name if file_name is not None else ONNX_WEIGHTS_NAME
         # load model from local directory
-        if os.path.isdir(model_id):
-            model = ORTModel.load_model(os.path.join(model_id, model_file_name), **kwargs)
-            kwargs["model_save_dir"] = Path(model_id)
+        if os.path.isdir(os.path.join(model_id, subfolder)):
+            model = ORTModel.load_model(os.path.join(model_id, subfolder, model_file_name), **kwargs)
+            kwargs["model_save_dir"] = Path(model_id).joinpath(subfolder)
             kwargs["latest_model_name"] = model_file_name
         # load model from hub
         else:
@@ -311,6 +314,7 @@ def _from_pretrained(
             model_cache_path = hf_hub_download(
                 repo_id=model_id,
                 filename=model_file_name,
+                subfolder=subfolder,
                 use_auth_token=use_auth_token,
                 revision=revision,
                 cache_dir=cache_dir,
@@ -332,6 +336,7 @@ def _from_transformers(
         revision: Optional[Union[str, None]] = None,
         force_download: bool = False,
         cache_dir: Optional[str] = None,
+        subfolder: Optional[str] = "",
         **kwargs,
     ):
         """
@@ -356,15 +361,17 @@ def _from_transformers(
                 kwargs will be passed to the model during initialization
         """
         # create local save dir in cache dir
-        save_dir = Path(save_dir).joinpath(model_id)
+        save_dir = Path(save_dir).joinpath(model_id, subfolder)
         save_dir.mkdir(parents=True, exist_ok=True)
         kwargs["model_save_dir"] = save_dir
 
+        config = kwargs.get("config", {})
+
         # reads pipeline task from ORTModelForXXX class if available else tries to extract from hub
         if cls.export_feature is not None:
             task = cls.export_feature
         else:
-            task = HfApi().model_info(model_id, revision=revision).pipeline_tag
+            task = HfApi().model_info(model_id, revision=revision).pipeline_tag  # FIXME: load from subfolder?
             if task in ["sentiment-analysis", "text-classification", "zero-shot-classification"]:
                 task = "sequence-classification"
             elif task in ["feature-extraction", "fill-mask"]:
@@ -372,7 +379,11 @@ def _from_transformers(
         # 2. convert to temp dir
         # FIXME: transformers.onnx conversion doesn't support private models
         preprocessor = get_preprocessor(model_id)
-        model = FeaturesManager.get_model_from_feature(task, model_id)
+
+        framework = FeaturesManager.determine_framework(os.path.join(model_id, subfolder))
+        model_class = FeaturesManager.get_model_class_for_feature(task, framework)
+        model = model_class.from_pretrained(model_id, subfolder=subfolder, config=config, cache_dir=cache_dir)
+
         _, model_onnx_config = FeaturesManager.check_supported_model_or_raise(model, feature=task)
         onnx_config = model_onnx_config(model.config)
 
@@ -384,7 +395,7 @@ def _from_transformers(
             opset=onnx_config.default_onnx_opset,
             output=save_dir.joinpath(ONNX_WEIGHTS_NAME),
         )
-        kwargs["config"] = model.config
+
         # 3. load normal model
         return cls._from_pretrained(save_dir.as_posix(), **kwargs)
 
diff --git a/optimum/onnxruntime/modeling_seq2seq.py b/optimum/onnxruntime/modeling_seq2seq.py
--- a/optimum/onnxruntime/modeling_seq2seq.py
+++ b/optimum/onnxruntime/modeling_seq2seq.py
@@ -304,6 +304,7 @@ def _from_pretrained(
         encoder_file_name: Optional[str] = None,
         decoder_file_name: Optional[str] = None,
         decoder_with_past_file_name: Optional[str] = None,
+        subfolder: Optional[str] = "",
         **kwargs,
     ):
         """
@@ -348,15 +349,17 @@ def _from_pretrained(
         decoder_with_past_file_name = decoder_with_past_file_name or ONNX_DECODER_WITH_PAST_NAME
 
         # Load model from a local directory
-        if os.path.isdir(model_id):
-            decoder_with_past_path = os.path.join(model_id, decoder_with_past_file_name) if use_cache else None
+        if os.path.isdir(os.path.join(model_id, subfolder)):
+            decoder_with_past_path = (
+                os.path.join(model_id, subfolder, decoder_with_past_file_name) if use_cache else None
+            )
             model = cls.load_model(
-                encoder_path=os.path.join(model_id, encoder_file_name),
-                decoder_path=os.path.join(model_id, decoder_file_name),
+                encoder_path=os.path.join(model_id, subfolder, encoder_file_name),
+                decoder_path=os.path.join(model_id, subfolder, decoder_file_name),
                 decoder_with_past_path=decoder_with_past_path,
                 **kwargs,
             )
-            kwargs["model_save_dir"] = Path(model_id)
+            kwargs["model_save_dir"] = Path(model_id).joinpath(subfolder)
             kwargs["last_encoder_name"] = encoder_file_name
             kwargs["last_decoder_name"] = decoder_file_name
             kwargs["last_decoder_with_past_name"] = decoder_with_past_file_name
@@ -371,6 +374,7 @@ def _from_pretrained(
             for file_name, default_file_name in zip(model_file_names, default_file_names):
                 model_cache_path = hf_hub_download(
                     repo_id=model_id,
+                    subfolder=subfolder,
                     filename=file_name,
                     use_auth_token=use_auth_token,
                     revision=revision,
@@ -397,6 +401,7 @@ def _from_pretrained(
     def _from_transformers(
         cls,
         model_id: str,
+        subfolder: Optional[str] = "",
         save_dir: Union[str, Path] = default_cache_path,
         use_auth_token: Optional[Union[bool, str, None]] = None,
         revision: Optional[Union[str, None]] = None,
@@ -430,12 +435,18 @@ def _from_transformers(
                 kwargs will be passed to the model during initialization.
         """
         # Create local save dir in cache dir
-        save_dir = Path(save_dir).joinpath(model_id)
+        save_dir = Path(save_dir).joinpath(model_id, subfolder)
         save_dir.mkdir(parents=True, exist_ok=True)
         kwargs["model_save_dir"] = save_dir
+        config = kwargs.get("config", {})
         use_cache = kwargs.get("use_cache", True)
+
         tokenizer = AutoTokenizer.from_pretrained(model_id)
-        model = FeaturesManager.get_model_from_feature(cls.export_feature, model_id)
+
+        framework = FeaturesManager.determine_framework(os.path.join(model_id, subfolder))
+        model_class = FeaturesManager.get_model_class_for_feature(cls.export_feature, framework)
+        model = model_class.from_pretrained(model_id, subfolder=subfolder, config=config, cache_dir=cache_dir)
+
         _, model_onnx_config = FeaturesManager.check_supported_model_or_raise(model, feature=cls.export_feature)
         onnx_config = model_onnx_config(model.config)
         onnx_opset = onnx_config.default_onnx_opset


Test patch:
diff --git a/tests/onnxruntime/test_modeling.py b/tests/onnxruntime/test_modeling.py
--- a/tests/onnxruntime/test_modeling.py
+++ b/tests/onnxruntime/test_modeling.py
@@ -107,6 +107,33 @@ def test_load_model_from_hub(self):
         self.assertIsInstance(model.model, onnxruntime.capi.onnxruntime_inference_collection.InferenceSession)
         self.assertIsInstance(model.config, PretrainedConfig)
 
+    def test_load_model_from_hub_subfolder(self):
+        # does not pass with ORTModel as it does not have export_feature attribute
+        model = ORTModelForSequenceClassification.from_pretrained(
+            "fxmarty/tiny-bert-sst2-distilled-subfolder", subfolder="my_subfolder", from_transformers=True
+        )
+        self.assertIsInstance(model.model, onnxruntime.capi.onnxruntime_inference_collection.InferenceSession)
+        self.assertIsInstance(model.config, PretrainedConfig)
+
+        model = ORTModel.from_pretrained("fxmarty/tiny-bert-sst2-distilled-onnx-subfolder", subfolder="my_subfolder")
+        self.assertIsInstance(model.model, onnxruntime.capi.onnxruntime_inference_collection.InferenceSession)
+        self.assertIsInstance(model.config, PretrainedConfig)
+
+    def test_load_seq2seq_model_from_hub_subfolder(self):
+        model = ORTModelForSeq2SeqLM.from_pretrained(
+            "fxmarty/tiny-mbart-subfolder", subfolder="my_folder", from_transformers=True
+        )
+        self.assertIsInstance(model.encoder, ORTEncoder)
+        self.assertIsInstance(model.decoder, ORTDecoder)
+        self.assertIsInstance(model.decoder_with_past, ORTDecoder)
+        self.assertIsInstance(model.config, PretrainedConfig)
+
+        model = ORTModelForSeq2SeqLM.from_pretrained("fxmarty/tiny-mbart-onnx-subfolder", subfolder="my_folder")
+        self.assertIsInstance(model.encoder, ORTEncoder)
+        self.assertIsInstance(model.decoder, ORTDecoder)
+        self.assertIsInstance(model.decoder_with_past, ORTDecoder)
+        self.assertIsInstance(model.config, PretrainedConfig)
+
     def test_load_model_from_cache(self):
         _ = ORTModel.from_pretrained(self.TINY_ONNX_MODEL_ID)  # caching