Improve the offline support for the ONNX/TFLite export

optimum/exporters/onnx/__main__.py

@@ -17,6 +17,7 @@
 import argparse
 from pathlib import Path
 
+from requests.exceptions import ConnectionError as RequestsConnectionError
 from transformers import AutoTokenizer
 from transformers.utils import is_torch_available
 
@@ -166,6 +167,19 @@ def main_export(
         )
 
     torch_dtype = None if fp16 is False else torch.float16
+
+    if task == "auto":
+        try:
+            task = TasksManager.infer_task_from_model(model_name_or_path)
+        except KeyError as e:
+            raise KeyError(
+                f"The task could not be automatically inferred. Please provide the argument --task with the relevant task from {', '.join(TasksManager.get_all_tasks())}. Detailed error: {e}"
+            )
+        except RequestsConnectionError as e:
+            raise RequestsConnectionError(
+                f"The task could not be automatically inferred as this is available only for models hosted on the Hugging Face Hub. Please provide the argument --task with the relevant task from {', '.join(TasksManager.get_all_tasks())}. Detailed error: {e}"
+            )
+
     model = TasksManager.get_model_from_task(
         task,
         model_name_or_path,
@@ -180,14 +194,6 @@ def main_export(
         torch_dtype=torch_dtype,
     )
 
-    if task == "auto":
-        try:
-            task = TasksManager.infer_task_from_model(model_name_or_path)
-        except KeyError as e:
-            raise KeyError(
-                f"The task could not be automatically inferred. Please provide the argument --task with the task from {', '.join(TasksManager.get_all_tasks())}. Detailed error: {e}"
-            )
-
     if task != "stable-diffusion" and task + "-with-past" in TasksManager.get_supported_tasks_for_model_type(
         model.config.model_type.replace("_", "-"), "onnx"
     ):

optimum/exporters/tasks.py

@@ -1125,15 +1125,18 @@ def get_model_class_for_task(
 
     @staticmethod
     def determine_framework(
-        model_name_or_path: Union[str, Path], subfolder: str = "", framework: Optional[str] = None
+        model_name_or_path: Union[str, Path],
+        subfolder: str = "",
+        framework: Optional[str] = None,
+        cache_dir: str = huggingface_hub.constants.HUGGINGFACE_HUB_CACHE,
     ) -> str:
         """
         Determines the framework to use for the export.
 
         The priority is in the following order:
             1. User input via `framework`.
             2. If local checkpoint is provided, use the same framework as the checkpoint.
-            3. If model repo, try to infer the framework from the Hub.
+            3. If model repo, try to infer the framework from the cache if available, else from the Hub.
             4. If could not infer, use available framework in environment, with priority given to PyTorch.
 
         Args:
@@ -1161,11 +1164,26 @@ def determine_framework(
                 for file in filenames
             ]
         else:
-            if not isinstance(model_name_or_path, str):
-                model_name_or_path = str(model_name_or_path)
-            all_files = huggingface_hub.list_repo_files(model_name_or_path, repo_type="model")
-            if subfolder != "":
-                all_files = [file[len(subfolder) + 1 :] for file in all_files if file.startswith(subfolder)]
+            object_id = model_name_or_path.replace("/", "--")
+            full_model_path = Path(cache_dir, f"models--{object_id}")
+            if full_model_path.is_dir():  # explore the cache first
+                # Resolve refs (for instance to convert main to the associated commit sha)
+                revision_file = Path(full_model_path, "refs", "main")
+                if revision_file.is_file():
+                    with open(revision_file) as f:
+                        revision = f.read()
+                cached_path = Path(full_model_path, "snapshots", revision, subfolder)
+                all_files = [
+                    os.path.relpath(os.path.join(dirpath, file), cached_path)
+                    for dirpath, _, filenames in os.walk(cached_path)
+                    for file in filenames
+                ]
+            else:
+                if not isinstance(model_name_or_path, str):
+                    model_name_or_path = str(model_name_or_path)
+                all_files = huggingface_hub.list_repo_files(model_name_or_path, repo_type="model")
+                if subfolder != "":
+                    all_files = [file[len(subfolder) + 1 :] for file in all_files if file.startswith(subfolder)]
 
         weight_name = Path(WEIGHTS_NAME).stem
         weight_extension = Path(WEIGHTS_NAME).suffix

optimum/exporters/tflite/__main__.py

@@ -16,6 +16,8 @@
 
 from argparse import ArgumentParser
 
+from requests.exceptions import ConnectionError as RequestsConnectionError
+
 from ...commands.export.tflite import parse_args_tflite
 from ...utils import logging
 from ...utils.save_utils import maybe_load_preprocessors, maybe_save_preprocessors
@@ -51,6 +53,10 @@ def main():
                 "The task could not be automatically inferred. Please provide the argument --task with the task "
                 f"from {', '.join(TasksManager.get_all_tasks())}. Detailed error: {e}"
             )
+        except RequestsConnectionError as e:
+            raise RequestsConnectionError(
+                f"The task could not be automatically inferred as this is available only for models hosted on the Hugging Face Hub. Please provide the argument --task with the relevant task from {', '.join(TasksManager.get_all_tasks())}. Detailed error: {e}"
+            )
 
     model = TasksManager.get_model_from_task(
         task, args.model, framework="tf", cache_dir=args.cache_dir, trust_remote_code=args.trust_remote_code