chendl/compositional_test
1
1# coding=utf-82# Copyright 2022 The HuggingFace Inc. team. All rights reserved.3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8# http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15 16import argparse17import collections.abc18import copy19import inspect20import json21import multiprocessing22import os23import shutil24import tempfile25import traceback26from pathlib import Path27 28from check_config_docstrings import get_checkpoint_from_config_class29from datasets import load_dataset30from get_test_info import get_model_to_tester_mapping, get_tester_classes_for_model31from huggingface_hub import Repository, create_repo, hf_api, upload_folder32 33from transformers import (34 CONFIG_MAPPING,35 FEATURE_EXTRACTOR_MAPPING,36 IMAGE_PROCESSOR_MAPPING,37 PROCESSOR_MAPPING,38 TOKENIZER_MAPPING,39 AutoTokenizer,40 LayoutLMv3TokenizerFast,41 PreTrainedTokenizer,42 PreTrainedTokenizerFast,43 logging,44)45from transformers.feature_extraction_utils import FeatureExtractionMixin46from transformers.file_utils import is_tf_available, is_torch_available47from transformers.image_processing_utils import BaseImageProcessor48from transformers.models.auto.configuration_auto import AutoConfig, model_type_to_module_name49from transformers.models.fsmt import configuration_fsmt50from transformers.processing_utils import ProcessorMixin, transformers_module51from transformers.tokenization_utils_base import PreTrainedTokenizerBase52 53 54# make sure tokenizer plays nice with multiprocessing55os.environ["TOKENIZERS_PARALLELISM"] = "false"56 57logging.set_verbosity_error()58logging.disable_progress_bar()59logger = logging.get_logger(__name__)60 61os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"62 63if not is_torch_available():64 raise ValueError("Please install PyTorch.")65 66if not is_tf_available():67 raise ValueError("Please install TensorFlow.")68 69 70FRAMEWORKS = ["pytorch", "tensorflow"]71INVALID_ARCH = []72TARGET_VOCAB_SIZE = 102473 74data = {"training_ds": None, "testing_ds": None}75 76COMPOSITE_MODELS = {77 "EncoderDecoderModel": "EncoderDecoderModel-bert-bert",78 "SpeechEncoderDecoderModel": "SpeechEncoderDecoderModel-wav2vec2-bert",79 "VisionEncoderDecoderModel": "VisionEncoderDecoderModel-vit-gpt2",80 "VisionTextDualEncoderModel": "VisionTextDualEncoderModel-vit-bert",81}82 83# This list contains the model architectures for which a tiny version could not be created.84# Avoid to add new architectures here - unless we have verified carefully that it's (almost) impossible to create them.85# One such case is: no model tester class is implemented for a model type (like `MT5`) because its architecture is86# identical to another one (`MT5` is based on `T5`), but trained on different datasets or with different techniques.87UNCONVERTIBLE_MODEL_ARCHITECTURES = {88 "BertGenerationEncoder",89 "BertGenerationDecoder",90 "CamembertForSequenceClassification",91 "CamembertForMultipleChoice",92 "CamembertForMaskedLM",93 "CamembertForCausalLM",94 "CamembertForTokenClassification",95 "CamembertForQuestionAnswering",96 "CamembertModel",97 "TFCamembertForMultipleChoice",98 "TFCamembertForTokenClassification",99 "TFCamembertForQuestionAnswering",100 "TFCamembertForSequenceClassification",101 "TFCamembertForMaskedLM",102 "TFCamembertModel",103 "TFCamembertForCausalLM",104 "DecisionTransformerModel",105 "GraphormerModel",106 "InformerModel",107 "JukeboxModel",108 "MarianForCausalLM",109 "MaskFormerSwinModel",110 "MaskFormerSwinBackbone",111 "MT5Model",112 "MT5ForConditionalGeneration",113 "TFMT5ForConditionalGeneration",114 "TFMT5Model",115 "QDQBertForSequenceClassification",116 "QDQBertForMaskedLM",117 "QDQBertModel",118 "QDQBertForTokenClassification",119 "QDQBertLMHeadModel",120 "QDQBertForMultipleChoice",121 "QDQBertForQuestionAnswering",122 "QDQBertForNextSentencePrediction",123 "ReformerModelWithLMHead",124 "RetriBertModel",125 "Speech2Text2ForCausalLM",126 "TimeSeriesTransformerModel",127 "TrajectoryTransformerModel",128 "TrOCRForCausalLM",129 "XLMProphetNetForConditionalGeneration",130 "XLMProphetNetForCausalLM",131 "XLMProphetNetModel",132 "XLMRobertaModel",133 "XLMRobertaForTokenClassification",134 "XLMRobertaForMultipleChoice",135 "XLMRobertaForMaskedLM",136 "XLMRobertaForCausalLM",137 "XLMRobertaForSequenceClassification",138 "XLMRobertaForQuestionAnswering",139 "TFXLMRobertaForSequenceClassification",140 "TFXLMRobertaForMaskedLM",141 "TFXLMRobertaForCausalLM",142 "TFXLMRobertaForQuestionAnswering",143 "TFXLMRobertaModel",144 "TFXLMRobertaForMultipleChoice",145 "TFXLMRobertaForTokenClassification",146}147 148 149def get_processor_types_from_config_class(config_class, allowed_mappings=None):150 """Return a tuple of processors for `config_class`.151 152 We use `tuple` here to include (potentially) both slow & fast tokenizers.153 """154 155 # To make a uniform return type156 def _to_tuple(x):157 if not isinstance(x, collections.abc.Sequence):158 x = (x,)159 else:160 x = tuple(x)161 return x162 163 if allowed_mappings is None:164 allowed_mappings = ["processor", "tokenizer", "image_processor", "feature_extractor"]165 166 processor_types = ()167 168 # Check first if a model has `ProcessorMixin`. Otherwise, check if it has tokenizers, and/or an image processor or169 # a feature extractor170 if config_class in PROCESSOR_MAPPING and "processor" in allowed_mappings:171 processor_types = _to_tuple(PROCESSOR_MAPPING[config_class])172 else:173 if config_class in TOKENIZER_MAPPING and "tokenizer" in allowed_mappings:174 processor_types = TOKENIZER_MAPPING[config_class]175 176 if config_class in IMAGE_PROCESSOR_MAPPING and "image_processor" in allowed_mappings:177 processor_types += _to_tuple(IMAGE_PROCESSOR_MAPPING[config_class])178 elif config_class in FEATURE_EXTRACTOR_MAPPING and "feature_extractor" in allowed_mappings:179 processor_types += _to_tuple(FEATURE_EXTRACTOR_MAPPING[config_class])180 181 # Remark: some configurations have no processor at all. For example, generic composite models like182 # `EncoderDecoderModel` is used for any (compatible) text models. Also, `DecisionTransformer` doesn't183 # require any processor.184 185 # We might get `None` for some tokenizers - remove them here.186 processor_types = tuple(p for p in processor_types if p is not None)187 188 return processor_types189 190 191def get_architectures_from_config_class(config_class, arch_mappings, models_to_skip=None):192 """Return a tuple of all possible architectures attributed to a configuration class `config_class`.193 194 For example, BertConfig -> [BertModel, BertForMaskedLM, ..., BertForQuestionAnswering].195 """196 # A model architecture could appear in several mappings. For example, `BartForConditionalGeneration` is in197 # - MODEL_FOR_PRETRAINING_MAPPING_NAMES198 # - MODEL_WITH_LM_HEAD_MAPPING_NAMES199 # - MODEL_FOR_MASKED_LM_MAPPING_NAMES200 # - MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING_NAMES201 # We avoid the duplication.202 architectures = set()203 204 if models_to_skip is None:205 models_to_skip = []206 models_to_skip = UNCONVERTIBLE_MODEL_ARCHITECTURES.union(models_to_skip)207 208 for mapping in arch_mappings:209 if config_class in mapping:210 models = mapping[config_class]211 models = tuple(models) if isinstance(models, collections.abc.Sequence) else (models,)212 for model in models:213 if model.__name__ not in models_to_skip:214 architectures.add(model)215 216 architectures = tuple(architectures)217 218 return architectures219 220 221def get_config_class_from_processor_class(processor_class):222 """Get the config class from a processor class.223 224 Some config/model classes use tokenizers/feature_extractors from other models. For example, `GPT-J` uses225 `GPT2Tokenizer`. If no checkpoint is found for a config class, or a checkpoint is found without necessary file(s) to226 create the processor for `processor_class`, we get the config class that corresponds to `processor_class` and use it227 to find a checkpoint in order to create the processor.228 """229 230 processor_prefix = processor_class.__name__231 for postfix in ["TokenizerFast", "Tokenizer", "ImageProcessor", "FeatureExtractor", "Processor"]:232 processor_prefix = processor_prefix.replace(postfix, "")233 234 # `Wav2Vec2CTCTokenizer` -> `Wav2Vec2Config`235 if processor_prefix == "Wav2Vec2CTC":236 processor_prefix = "Wav2Vec2"237 238 # Find the new configuration class239 new_config_name = f"{processor_prefix}Config"240 new_config_class = getattr(transformers_module, new_config_name)241 242 return new_config_class243 244 245def build_processor(config_class, processor_class, allow_no_checkpoint=False):246 """Create a processor for `processor_class`.247 248 If a processor is not able to be built with the original arguments, this method tries to change the arguments and249 call itself recursively, by inferring a new `config_class` or a new `processor_class` from another one, in order to250 find a checkpoint containing the necessary files to build a processor.251 252 The processor is not saved here. Instead, it will be saved in `convert_processors` after further changes in253 `convert_processors`. For each model architecture`, a copy will be created and saved along the built model.254 """255 # Currently, this solely uses the docstring in the source file of `config_class` to find a checkpoint.256 checkpoint = get_checkpoint_from_config_class(config_class)257 258 if checkpoint is None:259 # try to get the checkpoint from the config class for `processor_class`.260 # This helps cases like `XCLIPConfig` and `VideoMAEFeatureExtractor` to find a checkpoint from `VideoMAEConfig`.261 config_class_from_processor_class = get_config_class_from_processor_class(processor_class)262 checkpoint = get_checkpoint_from_config_class(config_class_from_processor_class)263 264 processor = None265 try:266 processor = processor_class.from_pretrained(checkpoint)267 except Exception as e:268 logger.error(f"{e.__class__.__name__}: {e}")269 270 # Try to get a new processor class from checkpoint. This is helpful for a checkpoint without necessary file to load271 # processor while `processor_class` is an Auto class. For example, `sew` has `Wav2Vec2Processor` in272 # `PROCESSOR_MAPPING_NAMES`, its `tokenizer_class` is `AutoTokenizer`, and the checkpoint273 # `https://huggingface.co/asapp/sew-tiny-100k` has no tokenizer file, but we can get274 # `tokenizer_class: Wav2Vec2CTCTokenizer` from the config file. (The new processor class won't be able to load from275 # `checkpoint`, but it helps this recursive method to find a way to build a processor).276 if (277 processor is None278 and checkpoint is not None279 and issubclass(processor_class, (PreTrainedTokenizerBase, AutoTokenizer))280 ):281 try:282 config = AutoConfig.from_pretrained(checkpoint)283 except Exception as e:284 logger.error(f"{e.__class__.__name__}: {e}")285 config = None286 if config is not None:287 if not isinstance(config, config_class):288 raise ValueError(289 f"`config` (which is of type {config.__class__.__name__}) should be an instance of `config_class`"290 f" ({config_class.__name__})!"291 )292 tokenizer_class = config.tokenizer_class293 new_processor_class = None294 if tokenizer_class is not None:295 new_processor_class = getattr(transformers_module, tokenizer_class)296 if new_processor_class != processor_class:297 processor = build_processor(config_class, new_processor_class)298 # If `tokenizer_class` is not specified in `config`, let's use `config` to get the process class via auto299 # mappings, but only allow the tokenizer mapping being used. This is to make `Wav2Vec2Conformer` build300 if processor is None:301 new_processor_classes = get_processor_types_from_config_class(302 config.__class__, allowed_mappings=["tokenizer"]303 )304 # Used to avoid infinite recursion between a pair of fast/slow tokenizer types305 names = [306 x.__name__.replace("Fast", "") for x in [processor_class, new_processor_class] if x is not None307 ]308 new_processor_classes = [309 x for x in new_processor_classes if x is not None and x.__name__.replace("Fast", "") not in names310 ]311 if len(new_processor_classes) > 0:312 new_processor_class = new_processor_classes[0]313 # Let's use fast tokenizer if there is any314 for x in new_processor_classes:315 if x.__name__.endswith("Fast"):316 new_processor_class = x317 break318 processor = build_processor(config_class, new_processor_class)319 320 if processor is None:321 # Try to build each component (tokenizer & feature extractor) of a `ProcessorMixin`.322 if issubclass(processor_class, ProcessorMixin):323 attrs = {}324 for attr_name in processor_class.attributes:325 attrs[attr_name] = []326 # This could be a tuple (for tokenizers). For example, `CLIPProcessor` has327 # - feature_extractor_class = "CLIPFeatureExtractor"328 # - tokenizer_class = ("CLIPTokenizer", "CLIPTokenizerFast")329 attr_class_names = getattr(processor_class, f"{attr_name}_class")330 if not isinstance(attr_class_names, tuple):331 attr_class_names = (attr_class_names,)332 333 for name in attr_class_names:334 attr_class = getattr(transformers_module, name)335 attr = build_processor(config_class, attr_class)336 if attr is not None:337 attrs[attr_name].append(attr)338 339 # try to build a `ProcessorMixin`, so we can return a single value340 if all(len(v) > 0 for v in attrs.values()):341 try:342 processor = processor_class(**{k: v[0] for k, v in attrs.items()})343 except Exception as e:344 logger.error(f"{e.__class__.__name__}: {e}")345 else:346 # `checkpoint` might lack some file(s) to load a processor. For example, `facebook/hubert-base-ls960`347 # has no tokenizer file to load `Wav2Vec2CTCTokenizer`. In this case, we try to build a processor348 # with the configuration class (for example, `Wav2Vec2Config`) corresponding to `processor_class`.349 config_class_from_processor_class = get_config_class_from_processor_class(processor_class)350 if config_class_from_processor_class != config_class:351 processor = build_processor(config_class_from_processor_class, processor_class)352 353 # Try to create an image processor or a feature extractor without any checkpoint354 if (355 processor is None356 and allow_no_checkpoint357 and (issubclass(processor_class, BaseImageProcessor) or issubclass(processor_class, FeatureExtractionMixin))358 ):359 try:360 processor = processor_class()361 except Exception as e:362 logger.error(f"{e.__class__.__name__}: {e}")363 364 # validation365 if processor is not None:366 if not (isinstance(processor, processor_class) or processor_class.__name__.startswith("Auto")):367 raise ValueError(368 f"`processor` (which is of type {processor.__class__.__name__}) should be an instance of"369 f" {processor_class.__name__} or an Auto class!"370 )371 372 return processor373 374 375def get_tiny_config(config_class, model_class=None, **model_tester_kwargs):376 """Retrieve a tiny configuration from `config_class` using each model's `ModelTester`.377 378 Args:379 config_class: Subclass of `PreTrainedConfig`.380 381 Returns:382 An instance of `config_class` with tiny hyperparameters383 """384 model_type = config_class.model_type385 386 # For model type like `data2vec-vision` and `donut-swin`, we can't get the config/model file name directly via387 # `model_type` as it would be sth. like `configuration_data2vec_vision.py`.388 # A simple way is to use `inspect.getsourcefile(config_class)`.389 config_source_file = inspect.getsourcefile(config_class)390 # The modeling file name without prefix (`modeling_`) and postfix (`.py`)391 modeling_name = config_source_file.split(os.path.sep)[-1].replace("configuration_", "").replace(".py", "")392 393 try:394 print("Importing", model_type_to_module_name(model_type))395 module_name = model_type_to_module_name(model_type)396 if not modeling_name.startswith(module_name):397 raise ValueError(f"{modeling_name} doesn't start with {module_name}!")398 test_file = os.path.join("tests", "models", module_name, f"test_modeling_{modeling_name}.py")399 models_to_model_testers = get_model_to_tester_mapping(test_file)400 # Find the model tester class401 model_tester_class = None402 tester_classes = []403 if model_class is not None:404 tester_classes = get_tester_classes_for_model(test_file, model_class)405 else:406 for _tester_classes in models_to_model_testers.values():407 tester_classes.extend(_tester_classes)408 if len(tester_classes) > 0:409 # sort with the length of the class names first, then the alphabetical order410 # This is to avoid `T5EncoderOnlyModelTest` is used instead of `T5ModelTest`, which has411 # `is_encoder_decoder=False` and causes some pipeline tests failing (also failures in `Optimum` CI).412 # TODO: More fine grained control of the desired tester class.413 model_tester_class = sorted(tester_classes, key=lambda x: (len(x.__name__), x.__name__))[0]414 except ModuleNotFoundError:415 error = f"Tiny config not created for {model_type} - cannot find the testing module from the model name."416 raise ValueError(error)417 418 if model_tester_class is None:419 error = f"Tiny config not created for {model_type} - no model tester is found in the testing module."420 raise ValueError(error)421 422 # `parent` is an instance of `unittest.TestCase`, but we don't need it here.423 model_tester = model_tester_class(parent=None, **model_tester_kwargs)424 425 if hasattr(model_tester, "get_pipeline_config"):426 return model_tester.get_pipeline_config()427 elif hasattr(model_tester, "prepare_config_and_inputs"):428 # `PoolFormer` has no `get_config` defined. Furthermore, it's better to use `prepare_config_and_inputs` even if429 # `get_config` is defined, since there might be some extra changes in `prepare_config_and_inputs`.430 return model_tester.prepare_config_and_inputs()[0]431 elif hasattr(model_tester, "get_config"):432 return model_tester.get_config()433 else:434 error = (435 f"Tiny config not created for {model_type} - the model tester {model_tester_class.__name__} lacks"436 " necessary method to create config."437 )438 raise ValueError(error)439 440 441def convert_tokenizer(tokenizer_fast: PreTrainedTokenizerFast):442 new_tokenizer = tokenizer_fast.train_new_from_iterator(443 data["training_ds"]["text"], TARGET_VOCAB_SIZE, show_progress=False444 )445 446 # Make sure it at least runs447 if not isinstance(new_tokenizer, LayoutLMv3TokenizerFast):448 new_tokenizer(data["testing_ds"]["text"])449 450 return new_tokenizer451 452 453def convert_feature_extractor(feature_extractor, tiny_config):454 to_convert = False455 kwargs = {}456 if hasattr(tiny_config, "image_size"):457 kwargs["size"] = tiny_config.image_size458 kwargs["crop_size"] = tiny_config.image_size459 to_convert = True460 elif (461 hasattr(tiny_config, "vision_config")462 and tiny_config.vision_config is not None463 and hasattr(tiny_config.vision_config, "image_size")464 ):465 kwargs["size"] = tiny_config.vision_config.image_size466 kwargs["crop_size"] = tiny_config.vision_config.image_size467 to_convert = True468 469 # Speech2TextModel specific.470 if hasattr(tiny_config, "input_feat_per_channel"):471 kwargs["feature_size"] = tiny_config.input_feat_per_channel472 kwargs["num_mel_bins"] = tiny_config.input_feat_per_channel473 to_convert = True474 475 if to_convert:476 feature_extractor = feature_extractor.__class__(**kwargs)477 478 return feature_extractor479 480 481def convert_processors(processors, tiny_config, output_folder, result):482 """Change a processor to work with smaller inputs.483 484 For tokenizers, we try to reduce their vocabulary size.485 486 For feature extractor, we use smaller image size or change487 other attributes using the values from `tiny_config`. See `convert_feature_extractor`.488 489 This method should not fail: we catch the errors and put them in `result["warnings"]` with descriptive messages.490 """491 492 def _sanity_check(fast_tokenizer, slow_tokenizer, keep_fast_tokenizer=False):493 """Set tokenizer(s) to `None` if the fast/slow tokenizers have different values for `vocab_size` or `length`.494 495 If `keep_fast_tokenizer=True`, the fast tokenizer will be kept.496 """497 # sanity check 1: fast and slow tokenizers should be compatible (vocab_size)498 if fast_tokenizer is not None and slow_tokenizer is not None:499 if fast_tokenizer.vocab_size != slow_tokenizer.vocab_size:500 warning_messagae = (501 "The fast/slow tokenizers "502 f"({fast_tokenizer.__class__.__name__}/{slow_tokenizer.__class__.__name__}) have different "503 "vocabulary size: "504 f"fast_tokenizer.vocab_size = {fast_tokenizer.vocab_size} and "505 f"slow_tokenizer.vocab_size = {slow_tokenizer.vocab_size}."506 )507 result["warnings"].append(warning_messagae)508 if not keep_fast_tokenizer:509 fast_tokenizer = None510 slow_tokenizer = None511 512 # sanity check 2: fast and slow tokenizers should be compatible (length)513 if fast_tokenizer is not None and slow_tokenizer is not None:514 if len(fast_tokenizer) != len(slow_tokenizer):515 warning_messagae = (516 f"The fast/slow tokenizers () have different length: "517 f"len(fast_tokenizer) = {len(fast_tokenizer)} and "518 f"len(slow_tokenizer) = {len(slow_tokenizer)}."519 )520 result["warnings"].append(warning_messagae)521 if not keep_fast_tokenizer:522 fast_tokenizer = None523 slow_tokenizer = None524 525 return fast_tokenizer, slow_tokenizer526 527 tokenizers = []528 feature_extractors = []529 for processor in processors:530 if isinstance(processor, PreTrainedTokenizerBase):531 if processor.__class__.__name__ not in {x.__class__.__name__ for x in tokenizers}:532 tokenizers.append(processor)533 elif isinstance(processor, BaseImageProcessor):534 if processor.__class__.__name__ not in {x.__class__.__name__ for x in feature_extractors}:535 feature_extractors.append(processor)536 elif isinstance(processor, FeatureExtractionMixin):537 if processor.__class__.__name__ not in {x.__class__.__name__ for x in feature_extractors}:538 feature_extractors.append(processor)539 elif isinstance(processor, ProcessorMixin):540 if hasattr(processor, "tokenizer"):541 if processor.tokenizer.__class__.__name__ not in {x.__class__.__name__ for x in tokenizers}:542 tokenizers.append(processor.tokenizer)543 # Currently, we only have these 2 possibilities544 if hasattr(processor, "image_processor"):545 if processor.image_processor.__class__.__name__ not in {546 x.__class__.__name__ for x in feature_extractors547 }:548 feature_extractors.append(processor.image_processor)549 elif hasattr(processor, "feature_extractor"):550 if processor.feature_extractor.__class__.__name__ not in {551 x.__class__.__name__ for x in feature_extractors552 }:553 feature_extractors.append(processor.feature_extractor)554 555 # check the built processors have the unique type556 num_types = len({x.__class__.__name__ for x in feature_extractors})557 if num_types >= 2:558 raise ValueError(f"`feature_extractors` should contain at most 1 type, but it contains {num_types} types!")559 num_types = len({x.__class__.__name__.replace("Fast", "") for x in tokenizers})560 if num_types >= 2:561 raise ValueError(f"`tokenizers` should contain at most 1 tokenizer type, but it contains {num_types} types!")562 563 fast_tokenizer = None564 slow_tokenizer = None565 566 for tokenizer in tokenizers:567 if isinstance(tokenizer, PreTrainedTokenizerFast):568 fast_tokenizer = tokenizer569 else:570 slow_tokenizer = tokenizer571 572 # If the (original) fast/slow tokenizers don't correspond, keep only the fast tokenizer.573 # This doesn't necessarily imply the fast/slow tokenizers in a single Hub repo. has issues.574 # It's more of an issue in `build_processor` which tries to get a checkpoint with as much effort as possible.575 # For `YosoModel` (which uses `AlbertTokenizer(Fast)`), its real (Hub) checkpoint doesn't contain valid files to576 # load the slower tokenizer (`AlbertTokenizer`), and it ends up finding the (canonical) checkpoint of `AlbertModel`,577 # which has different vocabulary.578 # TODO: Try to improve `build_processor`'s definition and/or usage to avoid the above situation in the first place.579 fast_tokenizer, slow_tokenizer = _sanity_check(fast_tokenizer, slow_tokenizer, keep_fast_tokenizer=True)580 original_fast_tokenizer, original_slow_tokenizer = fast_tokenizer, slow_tokenizer581 582 if fast_tokenizer:583 try:584 # Wav2Vec2ForCTC , ByT5Tokenizer etc. all are already small enough and have no fast version that can585 # be retrained586 if fast_tokenizer.vocab_size > TARGET_VOCAB_SIZE:587 fast_tokenizer = convert_tokenizer(fast_tokenizer)588 except Exception:589 result["warnings"].append(590 (591 f"Failed to convert the fast tokenizer for {fast_tokenizer.__class__.__name__}.",592 traceback.format_exc(),593 )594 )595 596 # If `fast_tokenizer` exists, `slow_tokenizer` should correspond to it.597 if fast_tokenizer:598 # Make sure the fast tokenizer can be saved599 try:600 # We don't save it to `output_folder` at this moment - only at the end of this function.601 with tempfile.TemporaryDirectory() as tmpdir:602 fast_tokenizer.save_pretrained(tmpdir)603 try:604 slow_tokenizer = AutoTokenizer.from_pretrained(tmpdir, use_fast=False)605 except Exception:606 result["warnings"].append(607 (608 f"Failed to load the slow tokenizer saved from {fast_tokenizer.__class__.__name__}.",609 traceback.format_exc(),610 )611 )612 # Let's just keep the fast version613 slow_tokenizer = None614 except Exception:615 result["warnings"].append(616 (617 f"Failed to save the fast tokenizer for {fast_tokenizer.__class__.__name__}.",618 traceback.format_exc(),619 )620 )621 fast_tokenizer = None622 623 # If the (possibly converted) fast/slow tokenizers don't correspond, set them to `None`, and use the original624 # tokenizers.625 fast_tokenizer, slow_tokenizer = _sanity_check(fast_tokenizer, slow_tokenizer, keep_fast_tokenizer=False)626 627 # If there is any conversion failed, we keep the original tokenizers.628 if (original_fast_tokenizer is not None and fast_tokenizer is None) or (629 original_slow_tokenizer is not None and slow_tokenizer is None630 ):631 warning_messagae = (632 "There are some issues when converting the fast/slow tokenizers. The original tokenizers from the Hub "633 " will be used instead."634 )635 result["warnings"].append(warning_messagae)636 # Let's use the original version at the end (`original_fast_tokenizer` and `original_slow_tokenizer`)637 fast_tokenizer = original_fast_tokenizer638 slow_tokenizer = original_slow_tokenizer639 640 # Make sure the fast tokenizer can be saved641 if fast_tokenizer:642 # We don't save it to `output_folder` at this moment - only at the end of this function.643 with tempfile.TemporaryDirectory() as tmpdir:644 try:645 fast_tokenizer.save_pretrained(tmpdir)646 except Exception:647 result["warnings"].append(648 (649 f"Failed to save the fast tokenizer for {fast_tokenizer.__class__.__name__}.",650 traceback.format_exc(),651 )652 )653 fast_tokenizer = None654 # Make sure the slow tokenizer can be saved655 if slow_tokenizer:656 # We don't save it to `output_folder` at this moment - only at the end of this function.657 with tempfile.TemporaryDirectory() as tmpdir:658 try:659 slow_tokenizer.save_pretrained(tmpdir)660 except Exception:661 result["warnings"].append(662 (663 f"Failed to save the slow tokenizer for {slow_tokenizer.__class__.__name__}.",664 traceback.format_exc(),665 )666 )667 slow_tokenizer = None668 669 # update feature extractors using the tiny config670 try:671 feature_extractors = [convert_feature_extractor(p, tiny_config) for p in feature_extractors]672 except Exception:673 result["warnings"].append(674 (675 "Failed to convert feature extractors.",676 traceback.format_exc(),677 )678 )679 feature_extractors = []680 681 if hasattr(tiny_config, "max_position_embeddings") and tiny_config.max_position_embeddings > 0:682 if fast_tokenizer is not None:683 if fast_tokenizer.__class__.__name__ in [684 "RobertaTokenizerFast",685 "XLMRobertaTokenizerFast",686 "LongformerTokenizerFast",687 "MPNetTokenizerFast",688 ]:689 fast_tokenizer.model_max_length = tiny_config.max_position_embeddings - 2690 else:691 fast_tokenizer.model_max_length = tiny_config.max_position_embeddings692 if slow_tokenizer is not None:693 if slow_tokenizer.__class__.__name__ in [694 "RobertaTokenizer",695 "XLMRobertaTokenizer",696 "LongformerTokenizer",697 "MPNetTokenizer",698 ]:699 slow_tokenizer.model_max_length = tiny_config.max_position_embeddings - 2700 else:701 slow_tokenizer.model_max_length = tiny_config.max_position_embeddings702 703 processors = [fast_tokenizer, slow_tokenizer] + feature_extractors704 processors = [p for p in processors if p is not None]705 for p in processors:706 p.save_pretrained(output_folder)707 708 return processors709 710 711def get_checkpoint_dir(output_dir, model_arch):712 """Get framework-agnostic architecture name. Used to save all PT/TF/Flax models into the same directory."""713 714 arch_name = model_arch.__name__715 if arch_name.startswith("TF"):716 arch_name = arch_name[2:]717 elif arch_name.startswith("Flax"):718 arch_name = arch_name[4:]719 720 return os.path.join(output_dir, arch_name)721 722 723def build_model(model_arch, tiny_config, output_dir):724 """Create and save a model for `model_arch`.725 726 Also copy the set of processors to each model (under the same model type) output folder.727 """728 729 checkpoint_dir = get_checkpoint_dir(output_dir, model_arch)730 731 processor_output_dir = os.path.join(output_dir, "processors")732 # copy the (same set of) processors (for a model type) to the model arch. specific folder733 if os.path.isdir(processor_output_dir):734 shutil.copytree(processor_output_dir, checkpoint_dir, dirs_exist_ok=True)735 736 tiny_config = copy.deepcopy(tiny_config)737 738 if any([model_arch.__name__.endswith(x) for x in ["ForCausalLM", "LMHeadModel"]]):739 tiny_config.is_encoder_decoder = False740 tiny_config.is_decoder = True741 742 model = model_arch(config=tiny_config)743 model.save_pretrained(checkpoint_dir)744 model.from_pretrained(checkpoint_dir)745 746 return model747 748 749def fill_result_with_error(result, error, trace, models_to_create):750 """Fill `result` with errors for all target model arch if we can't build processor"""751 error = (error, trace)752 result["error"] = error753 for framework in FRAMEWORKS:754 if framework in models_to_create:755 result[framework] = {}756 for model_arch in models_to_create[framework]:757 result[framework][model_arch.__name__] = {"model": None, "checkpoint": None, "error": error}758 759 result["processor"] = {p.__class__.__name__: p.__class__.__name__ for p in result["processor"].values()}760 761 762def upload_model(model_dir, organization, token):763 """Upload the tiny models"""764 765 arch_name = model_dir.split(os.path.sep)[-1]766 repo_name = f"tiny-random-{arch_name}"767 repo_id = f"{organization}/{repo_name}"768 769 repo_exist = False770 error = None771 try:772 create_repo(repo_id=repo_id, exist_ok=False, repo_type="model", token=token)773 except Exception as e:774 error = e775 if "You already created" in str(e):776 error = None777 logger.warning("Remote repository exists and will be cloned.")778 repo_exist = True779 try:780 create_repo(repo_id=repo_id, exist_ok=True, repo_type="model", token=token)781 except Exception as e:782 error = e783 if error is not None:784 raise error785 786 with tempfile.TemporaryDirectory() as tmpdir:787 repo = Repository(local_dir=tmpdir, clone_from=repo_id, token=token)788 repo.git_pull()789 shutil.copytree(model_dir, tmpdir, dirs_exist_ok=True)790 791 if repo_exist:792 # Open a PR on the existing Hub repo.793 hub_pr_url = upload_folder(794 folder_path=model_dir,795 repo_id=repo_id,796 repo_type="model",797 commit_message=f"Update tiny models for {arch_name}",798 commit_description=f"Upload tiny models for {arch_name}",799 create_pr=True,800 token=token,801 )802 logger.warning(f"PR open in {hub_pr_url}.")803 # TODO: We need this information?804 else:805 # Push to Hub repo directly806 repo.git_add(auto_lfs_track=True)807 repo.git_commit(f"Upload tiny models for {arch_name}")808 repo.git_push(blocking=True) # this prints a progress bar with the upload809 logger.warning(f"Tiny models {arch_name} pushed to {repo_id}.")810 811 812def build_composite_models(config_class, output_dir):813 import tempfile814 815 from transformers import (816 BertConfig,817 BertLMHeadModel,818 BertModel,819 BertTokenizer,820 BertTokenizerFast,821 EncoderDecoderModel,822 GPT2Config,823 GPT2LMHeadModel,824 GPT2Tokenizer,825 GPT2TokenizerFast,826 SpeechEncoderDecoderModel,827 TFEncoderDecoderModel,828 TFVisionEncoderDecoderModel,829 TFVisionTextDualEncoderModel,830 VisionEncoderDecoderModel,831 VisionTextDualEncoderModel,832 ViTConfig,833 ViTFeatureExtractor,834 ViTModel,835 Wav2Vec2Config,836 Wav2Vec2Model,837 Wav2Vec2Processor,838 )839 840 # These will be removed at the end if they are empty841 result = {"error": None, "warnings": []}842 843 if config_class.model_type == "encoder-decoder":844 encoder_config_class = BertConfig845 decoder_config_class = BertConfig846 encoder_processor = (BertTokenizerFast, BertTokenizer)847 decoder_processor = (BertTokenizerFast, BertTokenizer)848 encoder_class = BertModel849 decoder_class = BertLMHeadModel850 model_class = EncoderDecoderModel851 tf_model_class = TFEncoderDecoderModel852 elif config_class.model_type == "vision-encoder-decoder":853 encoder_config_class = ViTConfig854 decoder_config_class = GPT2Config855 encoder_processor = (ViTFeatureExtractor,)856 decoder_processor = (GPT2TokenizerFast, GPT2Tokenizer)857 encoder_class = ViTModel858 decoder_class = GPT2LMHeadModel859 model_class = VisionEncoderDecoderModel860 tf_model_class = TFVisionEncoderDecoderModel861 elif config_class.model_type == "speech-encoder-decoder":862 encoder_config_class = Wav2Vec2Config863 decoder_config_class = BertConfig864 encoder_processor = (Wav2Vec2Processor,)865 decoder_processor = (BertTokenizerFast, BertTokenizer)866 encoder_class = Wav2Vec2Model867 decoder_class = BertLMHeadModel868 model_class = SpeechEncoderDecoderModel869 tf_model_class = None870 elif config_class.model_type == "vision-text-dual-encoder":871 # Not encoder-decoder, but encoder-encoder. We just keep the same name as above to make code easier872 encoder_config_class = ViTConfig873 decoder_config_class = BertConfig874 encoder_processor = (ViTFeatureExtractor,)875 decoder_processor = (BertTokenizerFast, BertTokenizer)876 encoder_class = ViTModel877 decoder_class = BertModel878 model_class = VisionTextDualEncoderModel879 tf_model_class = TFVisionTextDualEncoderModel880 881 with tempfile.TemporaryDirectory() as tmpdir:882 try:883 # build encoder884 models_to_create = {"processor": encoder_processor, "pytorch": (encoder_class,), "tensorflow": []}885 encoder_output_dir = os.path.join(tmpdir, "encoder")886 build(encoder_config_class, models_to_create, encoder_output_dir)887 888 # build decoder889 models_to_create = {"processor": decoder_processor, "pytorch": (decoder_class,), "tensorflow": []}890 decoder_output_dir = os.path.join(tmpdir, "decoder")891 build(decoder_config_class, models_to_create, decoder_output_dir)892 893 # build encoder-decoder894 encoder_path = os.path.join(encoder_output_dir, encoder_class.__name__)895 decoder_path = os.path.join(decoder_output_dir, decoder_class.__name__)896 897 if config_class.model_type != "vision-text-dual-encoder":898 # Specify these explicitly for encoder-decoder like models, but not for `vision-text-dual-encoder` as it899 # has no decoder.900 decoder_config = decoder_config_class.from_pretrained(decoder_path)901 decoder_config.is_decoder = True902 decoder_config.add_cross_attention = True903 model = model_class.from_encoder_decoder_pretrained(904 encoder_path,905 decoder_path,906 decoder_config=decoder_config,907 )908 elif config_class.model_type == "vision-text-dual-encoder":909 model = model_class.from_vision_text_pretrained(encoder_path, decoder_path)910 911 model_path = os.path.join(912 output_dir,913 f"{model_class.__name__}-{encoder_config_class.model_type}-{decoder_config_class.model_type}",914 )915 model.save_pretrained(model_path)916 917 if tf_model_class is not None:918 model = tf_model_class.from_pretrained(model_path, from_pt=True)919 model.save_pretrained(model_path)920 921 # copy the processors922 encoder_processor_path = os.path.join(encoder_output_dir, "processors")923 decoder_processor_path = os.path.join(decoder_output_dir, "processors")924 if os.path.isdir(encoder_processor_path):925 shutil.copytree(encoder_processor_path, model_path, dirs_exist_ok=True)926 if os.path.isdir(decoder_processor_path):927 shutil.copytree(decoder_processor_path, model_path, dirs_exist_ok=True)928 929 # fill `result`930 result["processor"] = {x.__name__: x.__name__ for x in encoder_processor + decoder_processor}931 932 result["pytorch"] = {model_class.__name__: {"model": model_class.__name__, "checkpoint": model_path}}933 934 result["tensorflow"] = {}935 if tf_model_class is not None:936 result["tensorflow"] = {937 tf_model_class.__name__: {"model": tf_model_class.__name__, "checkpoint": model_path}938 }939 except Exception:940 result["error"] = (941 f"Failed to build models for {config_class.__name__}.",942 traceback.format_exc(),943 )944 945 if not result["error"]:946 del result["error"]947 if not result["warnings"]:948 del result["warnings"]949 950 return result951 952 953def get_token_id_from_tokenizer(token_id_name, tokenizer, original_token_id):954 """Use `tokenizer` to get the values of `bos_token_id`, `eos_token_ids`, etc.955 956 The argument `token_id_name` should be a string ending with `_token_id`, and `original_token_id` should be an957 integer that will be return if `tokenizer` has no token corresponding to `token_id_name`.958 """959 960 token_id = original_token_id961 962 if not token_id_name.endswith("_token_id"):963 raise ValueError(f"`token_id_name` is {token_id_name}, which doesn't end with `_token_id`!")964 965 token = getattr(tokenizer, token_id_name.replace("_token_id", "_token"), None)966 if token is not None:967 if isinstance(tokenizer, PreTrainedTokenizerFast):968 token_id = tokenizer._convert_token_to_id_with_added_voc(token)969 else:970 token_id = tokenizer._convert_token_to_id(token)971 972 return token_id973 974 975def get_config_overrides(config_class, processors):976 config_overrides = {}977 978 # Check if there is any tokenizer (prefer fast version if any)979 tokenizer = None980 for processor in processors:981 if isinstance(processor, PreTrainedTokenizerFast):982 tokenizer = processor983 break984 elif isinstance(processor, PreTrainedTokenizer):985 tokenizer = processor986 987 if tokenizer is None:988 return config_overrides989 990 # Get some properties of the (already converted) tokenizer (smaller vocab size, special token ids, etc.)991 # We use `len(tokenizer)` instead of `tokenizer.vocab_size` to avoid potential issues for tokenizers with non-empty992 # `added_tokens_encoder`. One example is the `DebertaV2Tokenizer` where the mask token is the extra token.993 vocab_size = len(tokenizer)994 995 # The original checkpoint has length `35998`, but it doesn't have ids `30400` and `30514` but instead `35998` and996 # `35999`.997 if config_class.__name__ == "GPTSanJapaneseConfig":998 vocab_size += 2999 1000 config_overrides["vocab_size"] = vocab_size1001 1002 # Used to create a new model tester with `tokenizer.vocab_size` in order to get the (updated) special token ids.1003 model_tester_kwargs = {"vocab_size": vocab_size}1004 # CLIP-like models have `text_model_tester` and `vision_model_tester`, and we need to pass `vocab_size` to1005 # `text_model_tester` via `text_kwargs`. The same trick is also necessary for `Flava`.1006 if config_class.__name__ in [1007 "AlignConfig",1008 "AltCLIPConfig",1009 "ChineseCLIPConfig",1010 "CLIPSegConfig",1011 "ClapConfig",1012 "CLIPConfig",1013 "GroupViTConfig",1014 "OwlViTConfig",1015 "XCLIPConfig",1016 "FlavaConfig",1017 "BlipConfig",1018 "Blip2Config",1019 ]:1020 del model_tester_kwargs["vocab_size"]1021 model_tester_kwargs["text_kwargs"] = {"vocab_size": vocab_size}1022 # `FSMTModelTester` accepts `src_vocab_size` and `tgt_vocab_size` but not `vocab_size`.1023 elif config_class.__name__ == "FSMTConfig":1024 del model_tester_kwargs["vocab_size"]1025 model_tester_kwargs["src_vocab_size"] = tokenizer.src_vocab_size1026 model_tester_kwargs["tgt_vocab_size"] = tokenizer.tgt_vocab_size1027 1028 _tiny_config = get_tiny_config(config_class, **model_tester_kwargs)1029 1030 # handle the possibility of `text_config` inside `_tiny_config` for clip-like models (`owlvit`, `groupvit`, etc.)1031 if hasattr(_tiny_config, "text_config"):1032 _tiny_config = _tiny_config.text_config1033 1034 # Collect values of some special token ids1035 for attr in dir(_tiny_config):1036 if attr.endswith("_token_id"):1037 token_id = getattr(_tiny_config, attr)1038 if token_id is not None:1039 # Using the token id values from `tokenizer` instead of from `_tiny_config`.1040 token_id = get_token_id_from_tokenizer(attr, tokenizer, original_token_id=token_id)1041 config_overrides[attr] = token_id1042 1043 if config_class.__name__ == "FSMTConfig":1044 config_overrides["src_vocab_size"] = tokenizer.src_vocab_size1045 config_overrides["tgt_vocab_size"] = tokenizer.tgt_vocab_size1046 # `FSMTConfig` has `DecoderConfig` as `decoder` attribute.1047 config_overrides["decoder"] = configuration_fsmt.DecoderConfig(1048 vocab_size=tokenizer.tgt_vocab_size, bos_token_id=config_overrides["eos_token_id"]1049 )1050 1051 return config_overrides1052 1053 1054def build(config_class, models_to_create, output_dir):1055 """Create all models for a certain model type.1056 1057 Args:1058 config_class (`PretrainedConfig`):1059 A subclass of `PretrainedConfig` that is used to determine `models_to_create`.1060 models_to_create (`dict`):1061 A dictionary containing the processor/model classes that we want to create the instances. These models are1062 of the same model type which is associated to `config_class`.1063 output_dir (`str`):1064 The directory to save all the checkpoints. Each model architecture will be saved in a subdirectory under1065 it. Models in different frameworks with the same architecture will be saved in the same subdirectory.1066 """1067 if data["training_ds"] is None or data["testing_ds"] is None:1068 ds = load_dataset("wikitext", "wikitext-2-raw-v1")1069 data["training_ds"] = ds["train"]1070 data["testing_ds"] = ds["test"]1071 1072 if config_class.model_type in [1073 "encoder-decoder",1074 "vision-encoder-decoder",1075 "speech-encoder-decoder",1076 "vision-text-dual-encoder",1077 ]:1078 return build_composite_models(config_class, output_dir)1079 1080 result = {k: {} for k in models_to_create}1081 1082 # These will be removed at the end if they are empty1083 result["error"] = None1084 result["warnings"] = []1085 1086 # Build processors1087 processor_classes = models_to_create["processor"]1088 1089 if len(processor_classes) == 0:1090 error = f"No processor class could be found in {config_class.__name__}."1091 fill_result_with_error(result, error, None, models_to_create)1092 logger.error(result["error"][0])1093 return result1094 1095 for processor_class in processor_classes:1096 try:1097 processor = build_processor(config_class, processor_class, allow_no_checkpoint=True)1098 if processor is not None:1099 result["processor"][processor_class] = processor1100 except Exception:1101 error = f"Failed to build processor for {processor_class.__name__}."1102 trace = traceback.format_exc()1103 fill_result_with_error(result, error, trace, models_to_create)1104 logger.error(result["error"][0])1105 return result1106 1107 if len(result["processor"]) == 0:1108 error = f"No processor could be built for {config_class.__name__}."1109 fill_result_with_error(result, error, None, models_to_create)1110 logger.error(result["error"][0])1111 return result1112 1113 try:1114 tiny_config = get_tiny_config(config_class)1115 except Exception as e:1116 error = f"Failed to get tiny config for {config_class.__name__}: {e}"1117 trace = traceback.format_exc()1118 fill_result_with_error(result, error, trace, models_to_create)1119 logger.error(result["error"][0])1120 return result1121 1122 # Convert the processors (reduce vocabulary size, smaller image size, etc.)1123 processors = list(result["processor"].values())1124 processor_output_folder = os.path.join(output_dir, "processors")1125 try:1126 processors = convert_processors(processors, tiny_config, processor_output_folder, result)1127 except Exception:1128 error = "Failed to convert the processors."1129 trace = traceback.format_exc()1130 result["warnings"].append((error, trace))1131 1132 if len(processors) == 0:1133 error = f"No processor is returned by `convert_processors` for {config_class.__name__}."1134 fill_result_with_error(result, error, None, models_to_create)1135 logger.error(result["error"][0])1136 return result1137 1138 try:1139 config_overrides = get_config_overrides(config_class, processors)1140 except Exception as e:1141 error = f"Failure occurs while calling `get_config_overrides`: {e}"1142 trace = traceback.format_exc()1143 fill_result_with_error(result, error, trace, models_to_create)1144 logger.error(result["error"][0])1145 return result1146 1147 # Just for us to see this easily in the report1148 if "vocab_size" in config_overrides:1149 result["vocab_size"] = config_overrides["vocab_size"]1150 1151 # Update attributes that `vocab_size` involves1152 for k, v in config_overrides.items():1153 if hasattr(tiny_config, k):1154 setattr(tiny_config, k, v)1155 # So far, we only have to deal with `text_config`, as `config_overrides` contains text-related attributes only.1156 elif (1157 hasattr(tiny_config, "text_config")1158 and tiny_config.text_config is not None1159 and hasattr(tiny_config.text_config, k)1160 ):1161 setattr(tiny_config.text_config, k, v)1162 # If `text_config_dict` exists, we need to update its value here too in order to # make1163 # `save_pretrained -> from_pretrained` work.1164 if hasattr(tiny_config, "text_config_dict"):1165 tiny_config.text_config_dict[k] = v1166 1167 if result["warnings"]:1168 logger.warning(result["warnings"][0][0])1169 1170 # update `result["processor"]`1171 result["processor"] = {type(p).__name__: p.__class__.__name__ for p in processors}1172 1173 for pytorch_arch in models_to_create["pytorch"]:1174 result["pytorch"][pytorch_arch.__name__] = {}1175 error = None1176 try:1177 model = build_model(pytorch_arch, tiny_config, output_dir=output_dir)1178 except Exception as e:1179 model = None1180 error = f"Failed to create the pytorch model for {pytorch_arch}: {e}"1181 trace = traceback.format_exc()1182 1183 result["pytorch"][pytorch_arch.__name__]["model"] = model.__class__.__name__ if model is not None else None1184 result["pytorch"][pytorch_arch.__name__]["checkpoint"] = (1185 get_checkpoint_dir(output_dir, pytorch_arch) if model is not None else None1186 )1187 if error is not None:1188 result["pytorch"][pytorch_arch.__name__]["error"] = (error, trace)1189 logger.error(f"{pytorch_arch.__name__}: {error}")1190 1191 for tensorflow_arch in models_to_create["tensorflow"]:1192 # Make PT/TF weights compatible1193 pt_arch_name = tensorflow_arch.__name__[2:] # Remove `TF`1194 pt_arch = getattr(transformers_module, pt_arch_name)1195 1196 result["tensorflow"][tensorflow_arch.__name__] = {}1197 error = None1198 if pt_arch.__name__ in result["pytorch"] and result["pytorch"][pt_arch.__name__]["checkpoint"] is not None:1199 ckpt = get_checkpoint_dir(output_dir, pt_arch)1200 # Use the same weights from PyTorch.