diff --git a/optimum/utils/save_utils.py b/optimum/utils/save_utils.py index c7d22bee9c..e455b1f752 100644 --- a/optimum/utils/save_utils.py +++ b/optimum/utils/save_utils.py @@ -27,6 +27,31 @@ def maybe_load_preprocessors( src_name_or_path: Union[str, Path], subfolder: str = "", trust_remote_code: bool = False ) -> List: + """Load all available preprocessors (tokenizer, processor, feature extractor, image processor) from a model path or Hub repo. + + Tries to load each preprocessor type in turn and silently skips any that are + not present. This is useful when you want to save preprocessors alongside + an exported model without knowing in advance which types are available. + + Args: + src_name_or_path (`Union[str, Path]`): Local path or Hugging Face Hub + model identifier to load the preprocessors from. + subfolder (`str`, *optional*, defaults to `""`): Subfolder within the + model directory or Hub repo where the preprocessor files are stored. + trust_remote_code (`bool`, *optional*, defaults to `False`): Whether to + allow running arbitrary remote code when loading preprocessors. + + Returns: + `List`: A list containing the successfully loaded preprocessor objects. + May be empty if none of the expected preprocessors are found. + + Example: + ```py + >>> preprocessors = maybe_load_preprocessors("bert-base-uncased") + >>> [type(p).__name__ for p in preprocessors] + ['BertTokenizerFast'] + ``` + """ preprocessors = [] try: preprocessors.append( diff --git a/optimum/utils/testing_utils.py b/optimum/utils/testing_utils.py index 2969c616e6..87a2bf9a83 100644 --- a/optimum/utils/testing_utils.py +++ b/optimum/utils/testing_utils.py @@ -136,18 +136,22 @@ def require_ort_training(test_case): def require_diffusers(test_case): + """Decorator marking a test that requires the `diffusers` package to be installed.""" return unittest.skipUnless(is_diffusers_available(), "test requires diffusers")(test_case) def require_timm(test_case): + """Decorator marking a test that requires the `timm` package to be installed.""" return unittest.skipUnless(is_timm_available(), "test requires timm")(test_case) def require_sentence_transformers(test_case): + """Decorator marking a test that requires the `sentence-transformers` package to be installed.""" return unittest.skipUnless(is_sentence_transformers_available(), "test requires sentence-transformers")(test_case) def require_datasets(test_case): + """Decorator marking a test that requires the `datasets` package to be installed.""" return unittest.skipUnless(is_datasets_available(), "test requires datasets")(test_case)