Download python files

Author: xadupre

_unittests/ut_torch_models/test_hghub_api.py

@@ -1,5 +1,7 @@
+import os
 import unittest
 import pandas
+import transformers
 from onnx_diagnostic.ext_test_case import (
     ExtTestCase,
     hide_stdout,
@@ -13,6 +15,7 @@
     enumerate_model_list,
     get_model_info,
     get_pretrained_config,
+    download_code_modelid,
     task_from_id,
     task_from_arch,
     task_from_tags,
@@ -147,6 +150,24 @@ def test__ccached_config_64(self):
         conf = _ccached_hf_internal_testing_tiny_random_beitforimageclassification()
         self.assertEqual(conf.auxiliary_channels, 256)
 
+    @requires_transformers("4.50")
+    @requires_torch("2.7")
+    @ignore_errors(OSError)  # connectivity issues
+    @hide_stdout()
+    def test_download_code_modelid(self):
+        model_id = "microsoft/Phi-3.5-MoE-instruct"
+        files = download_code_modelid(model_id, verbose=1, add_path_to_sys_path=True)
+        self.assertTrue(all(os.path.exists(f) for f in files))
+        pyf = [os.path.split(name)[-1] for name in files]
+        self.assertEqual(
+            ["configuration_phimoe.py", "modeling_phimoe.py", "sample_finetune.py"], pyf
+        )
+        cls = transformers.dynamic_module_utils.get_class_from_dynamic_module(
+            "modeling_phimoe.Phi4MMImageEmbedding",
+            pretrained_model_name_or_path=os.path.split(files[0])[0],
+        )
+        self.assertEqual(cls.__name__, "Phi4MMImageEmbedding")
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)

onnx_diagnostic/torch_models/hghub/hub_api.py

@@ -3,9 +3,10 @@
 import json
 import os
 import pprint
+import sys
 from typing import Any, Dict, List, Optional, Union
 import transformers
-from huggingface_hub import HfApi, model_info, hf_hub_download
+from huggingface_hub import HfApi, model_info, hf_hub_download, list_repo_files
 from ...helpers.config_helper import update_config
 from . import hub_data_cached_configs
 from .hub_data import __date__, __data_tasks__, load_architecture_task, __data_arch_values__
@@ -327,3 +328,43 @@ def enumerate_model_list(
             n -= 1
             if n == 0:
                 break
+
+
+def download_code_modelid(
+    model_id: str, verbose: int = 0, add_path_to_sys_path: bool = True
+) -> List[str]:
+    """
+    Downloads the code for a given model id.
+
+    :param model_id: model id
+    :param verbose: verbosity
+    :param add_path_to_sys_path: add folder where the files are downloaded to sys.path
+    :return: list of downloaded files
+    """
+    if verbose:
+        print(f"[download_code_modelid] retrieve file list for {model_id!r}")
+    files = list_repo_files(model_id)
+    pyfiles = [name for name in files if os.path.splitext(name)[-1] == ".py"]
+    if verbose:
+        print(f"[download_code_modelid] python files {pyfiles}")
+    absfiles = []
+    pathes = set()
+    for i, name in enumerate(pyfiles):
+        if verbose:
+            print(f"[download_code_modelid] download file {i+1}/{len(pyfiles)}: {name!r}")
+        r = hf_hub_download(repo_id=model_id, filename=name)
+        p = os.path.split(r)[0]
+        pathes.add(p)
+        absfiles.append(r)
+    if add_path_to_sys_path:
+        for p in pathes:
+            init = os.path.join(p, "__init__.py")
+            if not os.path.exists(init):
+                with open(init, "w"):
+                    pass
+            if p in sys.path:
+                continue
+            if verbose:
+                print(f"[download_code_modelid] add {p!r} to 'sys.path'")
+            sys.path.insert(0, p)
+    return absfiles

onnx_diagnostic/torch_models/hghub/model_inputs.py

@@ -6,7 +6,7 @@
 import transformers
 from ...helpers.config_helper import update_config
 from ...tasks import reduce_model_config, random_input_kwargs
-from .hub_api import task_from_arch, task_from_id, get_pretrained_config
+from .hub_api import task_from_arch, task_from_id, get_pretrained_config, download_code_modelid
 
 
 def _code_needing_rewriting(model: Any) -> Any:
@@ -149,7 +149,23 @@ def get_untrained_model_with_inputs(
         model = transformers.AutoModel.from_pretrained(model_id, **mkwargs)
     else:
         if archs is not None:
-            model = getattr(transformers, archs[0])(config)
+            try:
+                model = getattr(transformers, archs[0])(config)
+            except AttributeError as e:
+                # The code of the models is not in transformers but in the
+                # repository of the model. We need to download it.
+                pyfiles = download_code_modelid(model_id, verbose=verbose)
+                if pyfiles:
+                    cls = transformers.dynamic_module_utils.get_class_from_dynamic_module(
+                        archs[0], pretrained_model_name_or_path=os.path.split(pyfiles[0])[0]
+                    )
+                    model = cls(config)
+                else:
+                    raise AttributeError(
+                        f"Unable to find class 'tranformers.{archs[0]}'. "
+                        f"The code needs to be downloaded, config="
+                        f"\n{pprint.pformat(config)}."
+                    ) from e
         else:
             assert same_as_pretrained and use_pretrained, (
                 f"Model {model_id!r} cannot be built, the model cannot be built. "