From c11ac7857bf93914f69a209d50a918ad7fe1d9a3 Mon Sep 17 00:00:00 2001 From: Hafedh <70411813+not-lain@users.noreply.github.com> Date: Mon, 20 May 2024 10:38:32 +0100 Subject: [PATCH] fix for custom pipeline configuration (#29004) * fix for custom pipeline configuration * fix for custom pipelines * remove extra exception * added test for custom pipelines extra tag * format with ruff * limit extra tag for first time only * format with ruff * improve tests for custom pipelines --- src/transformers/configuration_utils.py | 5 +++++ src/transformers/feature_extraction_utils.py | 14 ++++++++++---- src/transformers/image_processing_utils.py | 15 ++++++++++----- src/transformers/processing_utils.py | 14 ++++++++++---- src/transformers/tokenization_utils_base.py | 20 +++++++++++++------- src/transformers/utils/__init__.py | 1 + src/transformers/utils/generic.py | 13 +++++++++++++ tests/pipelines/test_pipelines_common.py | 18 ++++++++++++++++++ 8 files changed, 80 insertions(+), 20 deletions(-) diff --git a/src/transformers/configuration_utils.py b/src/transformers/configuration_utils.py index fc7f782e34..86b946d162 100755 --- a/src/transformers/configuration_utils.py +++ b/src/transformers/configuration_utils.py @@ -32,6 +32,7 @@ from .utils import ( CONFIG_NAME, PushToHubMixin, add_model_info_to_auto_map, + add_model_info_to_custom_pipelines, cached_file, copy_func, download_url, @@ -736,6 +737,10 @@ class PretrainedConfig(PushToHubMixin): config_dict["auto_map"] = add_model_info_to_auto_map( config_dict["auto_map"], pretrained_model_name_or_path ) + if "custom_pipelines" in config_dict and not is_local: + config_dict["custom_pipelines"] = add_model_info_to_custom_pipelines( + config_dict["custom_pipelines"], pretrained_model_name_or_path + ) return config_dict, kwargs @classmethod diff --git a/src/transformers/feature_extraction_utils.py b/src/transformers/feature_extraction_utils.py index 76070ebeb8..44e01f8d85 100644 --- a/src/transformers/feature_extraction_utils.py +++ b/src/transformers/feature_extraction_utils.py @@ -31,6 +31,7 @@ from .utils import ( PushToHubMixin, TensorType, add_model_info_to_auto_map, + add_model_info_to_custom_pipelines, cached_file, copy_func, download_url, @@ -539,10 +540,15 @@ class FeatureExtractionMixin(PushToHubMixin): f"loading configuration file {feature_extractor_file} from cache at {resolved_feature_extractor_file}" ) - if "auto_map" in feature_extractor_dict and not is_local: - feature_extractor_dict["auto_map"] = add_model_info_to_auto_map( - feature_extractor_dict["auto_map"], pretrained_model_name_or_path - ) + if not is_local: + if "auto_map" in feature_extractor_dict: + feature_extractor_dict["auto_map"] = add_model_info_to_auto_map( + feature_extractor_dict["auto_map"], pretrained_model_name_or_path + ) + if "custom_pipelines" in feature_extractor_dict: + feature_extractor_dict["custom_pipelines"] = add_model_info_to_custom_pipelines( + feature_extractor_dict["custom_pipelines"], pretrained_model_name_or_path + ) return feature_extractor_dict, kwargs diff --git a/src/transformers/image_processing_utils.py b/src/transformers/image_processing_utils.py index efd1e04a62..c42378d8f3 100644 --- a/src/transformers/image_processing_utils.py +++ b/src/transformers/image_processing_utils.py @@ -31,6 +31,7 @@ from .utils import ( IMAGE_PROCESSOR_NAME, PushToHubMixin, add_model_info_to_auto_map, + add_model_info_to_custom_pipelines, cached_file, copy_func, download_url, @@ -375,11 +376,15 @@ class ImageProcessingMixin(PushToHubMixin): f"loading configuration file {image_processor_file} from cache at {resolved_image_processor_file}" ) - if "auto_map" in image_processor_dict and not is_local: - image_processor_dict["auto_map"] = add_model_info_to_auto_map( - image_processor_dict["auto_map"], pretrained_model_name_or_path - ) - + if not is_local: + if "auto_map" in image_processor_dict: + image_processor_dict["auto_map"] = add_model_info_to_auto_map( + image_processor_dict["auto_map"], pretrained_model_name_or_path + ) + if "custom_pipelines" in image_processor_dict: + image_processor_dict["custom_pipelines"] = add_model_info_to_custom_pipelines( + image_processor_dict["custom_pipelines"], pretrained_model_name_or_path + ) return image_processor_dict, kwargs @classmethod diff --git a/src/transformers/processing_utils.py b/src/transformers/processing_utils.py index 463591032c..a8e47fb683 100644 --- a/src/transformers/processing_utils.py +++ b/src/transformers/processing_utils.py @@ -30,6 +30,7 @@ from .utils import ( PROCESSOR_NAME, PushToHubMixin, add_model_info_to_auto_map, + add_model_info_to_custom_pipelines, cached_file, copy_func, direct_transformers_import, @@ -355,10 +356,15 @@ class ProcessorMixin(PushToHubMixin): else: logger.info(f"loading configuration file {processor_file} from cache at {resolved_processor_file}") - if "auto_map" in processor_dict and not is_local: - processor_dict["auto_map"] = add_model_info_to_auto_map( - processor_dict["auto_map"], pretrained_model_name_or_path - ) + if not is_local: + if "auto_map" in processor_dict: + processor_dict["auto_map"] = add_model_info_to_auto_map( + processor_dict["auto_map"], pretrained_model_name_or_path + ) + if "custom_pipelines" in processor_dict: + processor_dict["custom_pipelines"] = add_model_info_to_custom_pipelines( + processor_dict["custom_pipelines"], pretrained_model_name_or_path + ) return processor_dict, kwargs diff --git a/src/transformers/tokenization_utils_base.py b/src/transformers/tokenization_utils_base.py index 395f9859cd..4cb75c9864 100644 --- a/src/transformers/tokenization_utils_base.py +++ b/src/transformers/tokenization_utils_base.py @@ -42,6 +42,7 @@ from .utils import ( TensorType, add_end_docstrings, add_model_info_to_auto_map, + add_model_info_to_custom_pipelines, cached_file, copy_func, download_url, @@ -2177,13 +2178,18 @@ class PreTrainedTokenizerBase(SpecialTokensMixin, PushToHubMixin): config_tokenizer_class = None init_kwargs = init_configuration - if "auto_map" in init_kwargs and not _is_local: - # For backward compatibility with odl format. - if isinstance(init_kwargs["auto_map"], (tuple, list)): - init_kwargs["auto_map"] = {"AutoTokenizer": init_kwargs["auto_map"]} - init_kwargs["auto_map"] = add_model_info_to_auto_map( - init_kwargs["auto_map"], pretrained_model_name_or_path - ) + if not _is_local: + if "auto_map" in init_kwargs: + # For backward compatibility with odl format. + if isinstance(init_kwargs["auto_map"], (tuple, list)): + init_kwargs["auto_map"] = {"AutoTokenizer": init_kwargs["auto_map"]} + init_kwargs["auto_map"] = add_model_info_to_auto_map( + init_kwargs["auto_map"], pretrained_model_name_or_path + ) + if "custom_pipelines" in init_kwargs: + init_kwargs["custom_pipelines"] = add_model_info_to_custom_pipelines( + init_kwargs["custom_pipelines"], pretrained_model_name_or_path + ) if config_tokenizer_class is None: # Matt: This entire block is only used to decide if the tokenizer class matches the class in the repo. diff --git a/src/transformers/utils/__init__.py b/src/transformers/utils/__init__.py index 8c91463322..82d0b5001c 100755 --- a/src/transformers/utils/__init__.py +++ b/src/transformers/utils/__init__.py @@ -37,6 +37,7 @@ from .generic import ( PaddingStrategy, TensorType, add_model_info_to_auto_map, + add_model_info_to_custom_pipelines, cached_property, can_return_loss, expand_dims, diff --git a/src/transformers/utils/generic.py b/src/transformers/utils/generic.py index a8277588ff..1f332434a9 100644 --- a/src/transformers/utils/generic.py +++ b/src/transformers/utils/generic.py @@ -721,6 +721,19 @@ def add_model_info_to_auto_map(auto_map, repo_id): return auto_map +def add_model_info_to_custom_pipelines(custom_pipeline, repo_id): + """ + Adds the information of the repo_id to a given custom pipeline. + """ + # {custom_pipelines : {task: {"impl": "path.to.task"},...} } + for task in custom_pipeline.keys(): + if "impl" in custom_pipeline[task]: + module = custom_pipeline[task]["impl"] + if "--" not in module: + custom_pipeline[task]["impl"] = f"{repo_id}--{module}" + return custom_pipeline + + def infer_framework(model_class): """ Infers the framework of a given model without using isinstance(), because we cannot guarantee that the relevant diff --git a/tests/pipelines/test_pipelines_common.py b/tests/pipelines/test_pipelines_common.py index 763c7d1a88..6ede7d1c7a 100644 --- a/tests/pipelines/test_pipelines_common.py +++ b/tests/pipelines/test_pipelines_common.py @@ -925,6 +925,24 @@ class DynamicPipelineTester(unittest.TestCase): # Can't make an isinstance check because the new_classifier is from the PairClassificationPipeline class of a # dynamic module self.assertEqual(new_classifier.__class__.__name__, "PairClassificationPipeline") + # check for tag exitence, tag needs to be added when we are calling a custom pipeline from the hub + # useful for cases such as finetuning + self.assertDictEqual( + new_classifier.model.config.custom_pipelines, + { + "pair-classification": { + "impl": f"{USER}/test-dynamic-pipeline--custom_pipeline.PairClassificationPipeline", + "pt": ("AutoModelForSequenceClassification",), + "tf": (), + } + }, + ) + # test if the pipeline still works after the model is finetuned + # (we are actually testing if the pipeline still works from the final repo) + # this is where the user/repo--module.class is used for + new_classifier.model.push_to_hub(repo_name=f"{USER}/test-pipeline-for-a-finetuned-model", token=self._token) + del new_classifier # free up memory + new_classifier = pipeline(model=f"{USER}/test-pipeline-for-a-finetuned-model", trust_remote_code=True) results = classifier("I hate you", second_text="I love you") new_results = new_classifier("I hate you", second_text="I love you")