Add files using upload-large-folder tool

docs/transformers/tests/trainer/test_trainer_fsdp.py

@@ -0,0 +1,175 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from transformers import is_torch_available
+from transformers.testing_utils import (
+    TestCasePlus,
+    backend_device_count,
+    execute_subprocess_async,
+    get_torch_dist_unique_port,
+    require_accelerate,
+    require_fp8,
+    require_torch_multi_accelerator,
+    run_first,
+    torch_device,
+)
+
+
+if is_torch_available():
+    import torch
+    import torch.distributed
+    import torch.utils.data
+
+    from transformers import (
+        AutoModelForCausalLM,
+        AutoTokenizer,
+        DataCollatorForSeq2Seq,
+        EvalPrediction,
+        GenerationConfig,
+        HfArgumentParser,
+        PreTrainedTokenizerBase,
+        Seq2SeqTrainer,
+        Seq2SeqTrainingArguments,
+    )
+
+    class DummyTextDataset(torch.utils.data.Dataset[str]):
+        def __init__(self, tokenizer: PreTrainedTokenizerBase) -> None:
+            data = 4 * [
+                "Hello world!",
+                "The quick brown fox jumps over the lazy dog.",
+            ]
+            self.data = [
+                {k: v.squeeze(0) for k, v in tokenizer(item, return_tensors="pt", return_attention_mask=True).items()}
+                for item in data
+            ]
+            for item in self.data:
+                item["labels"] = item["input_ids"]
+
+        def __len__(self) -> int:
+            return len(self.data)
+
+        def __getitem__(self, i: int) -> str:
+            return self.data[i]
+
+
+class TestFSDPTrainer(TestCasePlus):
+    @require_torch_multi_accelerator
+    @require_accelerate
+    @run_first
+    def test_trainer(self):
+        output_dir = self.get_auto_remove_tmp_dir()
+        cmd = [
+            "accelerate",
+            "launch",
+            "--use_fsdp",
+            "--main_process_port",
+            f"{get_torch_dist_unique_port()}",
+            "--num_processes",
+            f"{backend_device_count(torch_device)}",
+            "--fsdp_transformer_layer_cls_to_wrap",
+            "GPT2Block",
+            f"{self.test_file_dir}/test_trainer_fsdp.py",
+            "--output_dir",
+            f"{output_dir}",
+            "--report_to",
+            "none",
+        ]
+        execute_subprocess_async(cmd, env=self.get_env())
+        # successful return here == success - any errors would have caused an error in the sub-call
+
+
+class TestFSDPTrainerFP8(TestCasePlus):
+    @require_torch_multi_accelerator
+    @require_accelerate
+    @require_fp8
+    @run_first
+    def test_trainer(self):
+        output_dir = self.get_auto_remove_tmp_dir()
+        cmd = [
+            "accelerate",
+            "launch",
+            "--use_fsdp",
+            "--main_process_port",
+            f"{get_torch_dist_unique_port()}",
+            "--num_processes",
+            f"{backend_device_count(torch_device)}",
+            "--mixed_precision",
+            "fp8",
+            "--fsdp_transformer_layer_cls_to_wrap",
+            "GPT2Block",
+            f"{self.test_file_dir}/test_trainer_fsdp.py",
+            "--output_dir",
+            f"{output_dir}",
+            "--report_to",
+            "none",
+        ]
+        execute_subprocess_async(cmd, env=self.get_env())
+        # successful return here == success - any errors would have caused an error in the sub-call
+
+
+class TestFSDPTrainerWrap(TestCasePlus):
+    @require_torch_multi_accelerator
+    @require_accelerate
+    @run_first
+    def test_trainer(self):
+        output_dir = self.get_auto_remove_tmp_dir()
+        cmd = [
+            "accelerate",
+            "launch",
+            "--use_fsdp",
+            "--main_process_port",
+            f"{get_torch_dist_unique_port()}",
+            "--num_processes",
+            f"{backend_device_count(torch_device)}",
+            "--fsdp_transformer_layer_cls_to_wrap",
+            "GPT2Block",
+            f"{self.test_file_dir}/test_trainer_fsdp.py",
+            "--output_dir",
+            f"{output_dir}",
+            "--report_to",
+            "none",
+            "--auto_find_batch_size",
+            "True",
+        ]
+        execute_subprocess_async(cmd, env=self.get_env())
+        # successful return here == success - any errors would have caused an error in the sub-call
+
+
+if __name__ == "__main__":
+    parser = HfArgumentParser((Seq2SeqTrainingArguments,))
+    training_args = parser.parse_args_into_dataclasses()[0]
+    training_args.per_device_eval_batch_size = 1
+    training_args.use_legacy_prediction_loop = False
+    training_args.predict_with_generate = True
+    training_args.generation_config = GenerationConfig(max_length=30)
+
+    pretrained_model_name = "hf-internal-testing/tiny-random-gpt2"
+    tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name)
+    tokenizer.pad_token = tokenizer.eos_token
+    device = torch.device(torch.distributed.get_rank())
+    model = AutoModelForCausalLM.from_pretrained(pretrained_model_name).to(device)
+
+    def compute_metrics(p: EvalPrediction) -> dict[str, bool]:
+        return {"accuracy": (p.predictions == p.label_ids).mean()}
+
+    trainer = Seq2SeqTrainer(
+        model=model,
+        args=training_args,
+        data_collator=DataCollatorForSeq2Seq(tokenizer, model),
+        eval_dataset=DummyTextDataset(tokenizer),
+        compute_metrics=compute_metrics,
+    )
+
+    metrics = trainer.evaluate()

docs/transformers/tests/utils/import_structures/failing_export.py

@@ -0,0 +1,23 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# fmt: off
+
+from transformers.utils.import_utils import requires
+
+
+@requires(backends=("random_item_that_should_not_exist",))
+class A0:
+    def __init__(self):
+        pass

docs/transformers/tests/utils/import_structures/import_structure_raw_register.py

@@ -0,0 +1,80 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# fmt: off
+
+from transformers.utils.import_utils import requires
+
+
+@requires()
+class A0:
+    def __init__(self):
+        pass
+
+
+@requires()
+def a0():
+    pass
+
+
+@requires(backends=("torch", "tf"))
+class A1:
+    def __init__(self):
+        pass
+
+
+@requires(backends=("torch", "tf"))
+def a1():
+    pass
+
+
+@requires(
+    backends=("torch", "tf")
+)
+class A2:
+    def __init__(self):
+        pass
+
+
+@requires(
+    backends=("torch", "tf")
+)
+def a2():
+    pass
+
+
+@requires(
+    backends=(
+        "torch",
+        "tf"
+    )
+)
+class A3:
+    def __init__(self):
+        pass
+
+
+@requires(
+    backends=(
+            "torch",
+            "tf"
+    )
+)
+def a3():
+    pass
+
+@requires(backends=())
+class A4:
+    def __init__(self):
+        pass

docs/transformers/tests/utils/import_structures/import_structure_register_with_comments.py

@@ -0,0 +1,79 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# fmt: off
+
+from transformers.utils.import_utils import requires
+
+
+@requires()
+# That's a statement
+class B0:
+    def __init__(self):
+        pass
+
+
+@requires()
+# That's a statement
+def b0():
+    pass
+
+
+@requires(backends=("torch", "tf"))
+# That's a statement
+class B1:
+    def __init__(self):
+        pass
+
+
+@requires(backends=("torch", "tf"))
+# That's a statement
+def b1():
+    pass
+
+
+@requires(backends=("torch", "tf"))
+# That's a statement
+class B2:
+    def __init__(self):
+        pass
+
+
+@requires(backends=("torch", "tf"))
+# That's a statement
+def b2():
+    pass
+
+
+@requires(
+    backends=(
+        "torch",
+        "tf"
+    )
+)
+# That's a statement
+class B3:
+    def __init__(self):
+        pass
+
+
+@requires(
+    backends=(
+        "torch",
+        "tf"
+    )
+)
+# That's a statement
+def b3():
+    pass

docs/transformers/tests/utils/import_structures/import_structure_register_with_duplicates.py

@@ -0,0 +1,77 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# fmt: off
+
+from transformers.utils.import_utils import requires
+
+
+@requires(backends=("torch", "torch"))
+class C0:
+    def __init__(self):
+        pass
+
+
+@requires(backends=("torch", "torch"))
+def c0():
+    pass
+
+
+@requires(backends=("torch", "torch"))
+# That's a statement
+class C1:
+    def __init__(self):
+        pass
+
+
+@requires(backends=("torch", "torch"))
+# That's a statement
+def c1():
+    pass
+
+
+@requires(backends=("torch", "torch"))
+# That's a statement
+class C2:
+    def __init__(self):
+        pass
+
+
+@requires(backends=("torch", "torch"))
+# That's a statement
+def c2():
+    pass
+
+
+@requires(
+    backends=(
+        "torch",
+        "torch"
+    )
+)
+# That's a statement
+class C3:
+    def __init__(self):
+        pass
+
+
+@requires(
+    backends=(
+        "torch",
+        "torch"
+    )
+)
+# That's a statement
+def c3():
+    pass

docs/transformers/tests/utils/test_activations_tf.py

@@ -0,0 +1,60 @@
+# Copyright 2020 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+
+from transformers import is_tf_available
+from transformers.testing_utils import require_tf
+
+
+if is_tf_available():
+    import tensorflow as tf
+
+    from transformers.activations_tf import get_tf_activation
+
+
+@require_tf
+class TestTFActivations(unittest.TestCase):
+    def test_gelu_10(self):
+        x = tf.constant([-100, -1.0, -0.1, 0, 0.1, 1.0, 100.0])
+        gelu = get_tf_activation("gelu")
+        gelu10 = get_tf_activation("gelu_10")
+
+        y_gelu = gelu(x)
+        y_gelu_10 = gelu10(x)
+
+        clipped_mask = tf.where(y_gelu_10 < 10.0, 1.0, 0.0)
+
+        self.assertEqual(tf.math.reduce_max(y_gelu_10).numpy().item(), 10.0)
+        self.assertTrue(np.allclose(y_gelu * clipped_mask, y_gelu_10 * clipped_mask))
+
+    def test_get_activation(self):
+        get_tf_activation("gelu")
+        get_tf_activation("gelu_10")
+        get_tf_activation("gelu_fast")
+        get_tf_activation("gelu_new")
+        get_tf_activation("glu")
+        get_tf_activation("mish")
+        get_tf_activation("quick_gelu")
+        get_tf_activation("relu")
+        get_tf_activation("sigmoid")
+        get_tf_activation("silu")
+        get_tf_activation("swish")
+        get_tf_activation("tanh")
+        with self.assertRaises(KeyError):
+            get_tf_activation("bogus")
+        with self.assertRaises(KeyError):
+            get_tf_activation(None)

docs/transformers/tests/utils/test_add_new_model_like.py

@@ -0,0 +1,1506 @@
+# Copyright 2022 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import os
+import re
+import tempfile
+import unittest
+from pathlib import Path
+
+import transformers
+from transformers.commands.add_new_model_like import (
+    ModelPatterns,
+    _re_class_func,
+    add_content_to_file,
+    add_content_to_text,
+    clean_frameworks_in_init,
+    duplicate_doc_file,
+    duplicate_module,
+    filter_framework_files,
+    find_base_model_checkpoint,
+    get_model_files,
+    get_module_from_file,
+    parse_module_content,
+    replace_model_patterns,
+    retrieve_info_for_model,
+    retrieve_model_classes,
+    simplify_replacements,
+)
+from transformers.testing_utils import require_flax, require_tf, require_torch
+
+
+BERT_MODEL_FILES = {
+    "src/transformers/models/bert/__init__.py",
+    "src/transformers/models/bert/configuration_bert.py",
+    "src/transformers/models/bert/tokenization_bert.py",
+    "src/transformers/models/bert/tokenization_bert_fast.py",
+    "src/transformers/models/bert/tokenization_bert_tf.py",
+    "src/transformers/models/bert/modeling_bert.py",
+    "src/transformers/models/bert/modeling_flax_bert.py",
+    "src/transformers/models/bert/modeling_tf_bert.py",
+    "src/transformers/models/bert/convert_bert_original_tf_checkpoint_to_pytorch.py",
+    "src/transformers/models/bert/convert_bert_original_tf2_checkpoint_to_pytorch.py",
+    "src/transformers/models/bert/convert_bert_pytorch_checkpoint_to_original_tf.py",
+    "src/transformers/models/bert/convert_bert_token_dropping_original_tf2_checkpoint_to_pytorch.py",
+}
+
+VIT_MODEL_FILES = {
+    "src/transformers/models/vit/__init__.py",
+    "src/transformers/models/vit/configuration_vit.py",
+    "src/transformers/models/vit/convert_dino_to_pytorch.py",
+    "src/transformers/models/vit/convert_vit_timm_to_pytorch.py",
+    "src/transformers/models/vit/feature_extraction_vit.py",
+    "src/transformers/models/vit/image_processing_vit.py",
+    "src/transformers/models/vit/image_processing_vit_fast.py",
+    "src/transformers/models/vit/modeling_vit.py",
+    "src/transformers/models/vit/modeling_tf_vit.py",
+    "src/transformers/models/vit/modeling_flax_vit.py",
+}
+
+WAV2VEC2_MODEL_FILES = {
+    "src/transformers/models/wav2vec2/__init__.py",
+    "src/transformers/models/wav2vec2/configuration_wav2vec2.py",
+    "src/transformers/models/wav2vec2/convert_wav2vec2_original_pytorch_checkpoint_to_pytorch.py",
+    "src/transformers/models/wav2vec2/convert_wav2vec2_original_s3prl_checkpoint_to_pytorch.py",
+    "src/transformers/models/wav2vec2/feature_extraction_wav2vec2.py",
+    "src/transformers/models/wav2vec2/modeling_wav2vec2.py",
+    "src/transformers/models/wav2vec2/modeling_tf_wav2vec2.py",
+    "src/transformers/models/wav2vec2/modeling_flax_wav2vec2.py",
+    "src/transformers/models/wav2vec2/processing_wav2vec2.py",
+    "src/transformers/models/wav2vec2/tokenization_wav2vec2.py",
+}
+
+REPO_PATH = Path(transformers.__path__[0]).parent.parent
+
+
+@require_torch
+@require_tf
+@require_flax
+class TestAddNewModelLike(unittest.TestCase):
+    def init_file(self, file_name, content):
+        with open(file_name, "w", encoding="utf-8") as f:
+            f.write(content)
+
+    def check_result(self, file_name, expected_result):
+        with open(file_name, encoding="utf-8") as f:
+            result = f.read()
+            self.assertEqual(result, expected_result)
+
+    def test_re_class_func(self):
+        self.assertEqual(_re_class_func.search("def my_function(x, y):").groups()[0], "my_function")
+        self.assertEqual(_re_class_func.search("class MyClass:").groups()[0], "MyClass")
+        self.assertEqual(_re_class_func.search("class MyClass(SuperClass):").groups()[0], "MyClass")
+
+    def test_model_patterns_defaults(self):
+        model_patterns = ModelPatterns("GPT-New new", "huggingface/gpt-new-base")
+
+        self.assertEqual(model_patterns.model_type, "gpt-new-new")
+        self.assertEqual(model_patterns.model_lower_cased, "gpt_new_new")
+        self.assertEqual(model_patterns.model_camel_cased, "GPTNewNew")
+        self.assertEqual(model_patterns.model_upper_cased, "GPT_NEW_NEW")
+        self.assertEqual(model_patterns.config_class, "GPTNewNewConfig")
+        self.assertIsNone(model_patterns.tokenizer_class)
+        self.assertIsNone(model_patterns.feature_extractor_class)
+        self.assertIsNone(model_patterns.processor_class)
+
+    def test_parse_module_content(self):
+        test_code = """SOME_CONSTANT = a constant
+
+CONSTANT_DEFINED_ON_SEVERAL_LINES = [
+    first_item,
+    second_item
+]
+
+def function(args):
+    some code
+
+# Copied from transformers.some_module
+class SomeClass:
+    some code
+"""
+
+        expected_parts = [
+            "SOME_CONSTANT = a constant\n",
+            "CONSTANT_DEFINED_ON_SEVERAL_LINES = [\n    first_item,\n    second_item\n]",
+            "",
+            "def function(args):\n    some code\n",
+            "# Copied from transformers.some_module\nclass SomeClass:\n    some code\n",
+        ]
+        self.assertEqual(parse_module_content(test_code), expected_parts)
+
+    def test_add_content_to_text(self):
+        test_text = """all_configs = {
+    "gpt": "GPTConfig",
+    "bert": "BertConfig",
+    "t5": "T5Config",
+}"""
+
+        expected = """all_configs = {
+    "gpt": "GPTConfig",
+    "gpt2": "GPT2Config",
+    "bert": "BertConfig",
+    "t5": "T5Config",
+}"""
+        line = '    "gpt2": "GPT2Config",'
+
+        self.assertEqual(add_content_to_text(test_text, line, add_before="bert"), expected)
+        self.assertEqual(add_content_to_text(test_text, line, add_before="bert", exact_match=True), test_text)
+        self.assertEqual(
+            add_content_to_text(test_text, line, add_before='    "bert": "BertConfig",', exact_match=True), expected
+        )
+        self.assertEqual(add_content_to_text(test_text, line, add_before=re.compile(r'^\s*"bert":')), expected)
+
+        self.assertEqual(add_content_to_text(test_text, line, add_after="gpt"), expected)
+        self.assertEqual(add_content_to_text(test_text, line, add_after="gpt", exact_match=True), test_text)
+        self.assertEqual(
+            add_content_to_text(test_text, line, add_after='    "gpt": "GPTConfig",', exact_match=True), expected
+        )
+        self.assertEqual(add_content_to_text(test_text, line, add_after=re.compile(r'^\s*"gpt":')), expected)
+
+    def test_add_content_to_file(self):
+        test_text = """all_configs = {
+    "gpt": "GPTConfig",
+    "bert": "BertConfig",
+    "t5": "T5Config",
+}"""
+
+        expected = """all_configs = {
+    "gpt": "GPTConfig",
+    "gpt2": "GPT2Config",
+    "bert": "BertConfig",
+    "t5": "T5Config",
+}"""
+        line = '    "gpt2": "GPT2Config",'
+
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            file_name = os.path.join(tmp_dir, "code.py")
+
+            self.init_file(file_name, test_text)
+            add_content_to_file(file_name, line, add_before="bert")
+            self.check_result(file_name, expected)
+
+            self.init_file(file_name, test_text)
+            add_content_to_file(file_name, line, add_before="bert", exact_match=True)
+            self.check_result(file_name, test_text)
+
+            self.init_file(file_name, test_text)
+            add_content_to_file(file_name, line, add_before='    "bert": "BertConfig",', exact_match=True)
+            self.check_result(file_name, expected)
+
+            self.init_file(file_name, test_text)
+            add_content_to_file(file_name, line, add_before=re.compile(r'^\s*"bert":'))
+            self.check_result(file_name, expected)
+
+            self.init_file(file_name, test_text)
+            add_content_to_file(file_name, line, add_after="gpt")
+            self.check_result(file_name, expected)
+
+            self.init_file(file_name, test_text)
+            add_content_to_file(file_name, line, add_after="gpt", exact_match=True)
+            self.check_result(file_name, test_text)
+
+            self.init_file(file_name, test_text)
+            add_content_to_file(file_name, line, add_after='    "gpt": "GPTConfig",', exact_match=True)
+            self.check_result(file_name, expected)
+
+            self.init_file(file_name, test_text)
+            add_content_to_file(file_name, line, add_after=re.compile(r'^\s*"gpt":'))
+            self.check_result(file_name, expected)
+
+    def test_simplify_replacements(self):
+        self.assertEqual(simplify_replacements([("Bert", "NewBert")]), [("Bert", "NewBert")])
+        self.assertEqual(
+            simplify_replacements([("Bert", "NewBert"), ("bert", "new-bert")]),
+            [("Bert", "NewBert"), ("bert", "new-bert")],
+        )
+        self.assertEqual(
+            simplify_replacements([("BertConfig", "NewBertConfig"), ("Bert", "NewBert"), ("bert", "new-bert")]),
+            [("Bert", "NewBert"), ("bert", "new-bert")],
+        )
+
+    def test_replace_model_patterns(self):
+        bert_model_patterns = ModelPatterns("Bert", "google-bert/bert-base-cased")
+        new_bert_model_patterns = ModelPatterns("New Bert", "huggingface/bert-new-base")
+        bert_test = '''class TFBertPreTrainedModel(PreTrainedModel):
+    """
+    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
+    models.
+    """
+
+    config_class = BertConfig
+    load_tf_weights = load_tf_weights_in_bert
+    base_model_prefix = "bert"
+    is_parallelizable = True
+    supports_gradient_checkpointing = True
+    model_type = "bert"
+
+BERT_CONSTANT = "value"
+'''
+        bert_expected = '''class TFNewBertPreTrainedModel(PreTrainedModel):
+    """
+    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
+    models.
+    """
+
+    config_class = NewBertConfig
+    load_tf_weights = load_tf_weights_in_new_bert
+    base_model_prefix = "new_bert"
+    is_parallelizable = True
+    supports_gradient_checkpointing = True
+    model_type = "new-bert"
+
+NEW_BERT_CONSTANT = "value"
+'''
+
+        bert_converted, replacements = replace_model_patterns(bert_test, bert_model_patterns, new_bert_model_patterns)
+        self.assertEqual(bert_converted, bert_expected)
+        # Replacements are empty here since bert as been replaced by bert_new in some instances and bert-new
+        # in others.
+        self.assertEqual(replacements, "")
+
+        # If we remove the model type, we will get replacements
+        bert_test = bert_test.replace('    model_type = "bert"\n', "")
+        bert_expected = bert_expected.replace('    model_type = "new-bert"\n', "")
+        bert_converted, replacements = replace_model_patterns(bert_test, bert_model_patterns, new_bert_model_patterns)
+        self.assertEqual(bert_converted, bert_expected)
+        self.assertEqual(replacements, "BERT->NEW_BERT,Bert->NewBert,bert->new_bert")
+
+        gpt_model_patterns = ModelPatterns("GPT2", "gpt2")
+        new_gpt_model_patterns = ModelPatterns("GPT-New new", "huggingface/gpt-new-base")
+        gpt_test = '''class GPT2PreTrainedModel(PreTrainedModel):
+    """
+    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
+    models.
+    """
+
+    config_class = GPT2Config
+    load_tf_weights = load_tf_weights_in_gpt2
+    base_model_prefix = "transformer"
+    is_parallelizable = True
+    supports_gradient_checkpointing = True
+
+GPT2_CONSTANT = "value"
+'''
+
+        gpt_expected = '''class GPTNewNewPreTrainedModel(PreTrainedModel):
+    """
+    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
+    models.
+    """
+
+    config_class = GPTNewNewConfig
+    load_tf_weights = load_tf_weights_in_gpt_new_new
+    base_model_prefix = "transformer"
+    is_parallelizable = True
+    supports_gradient_checkpointing = True
+
+GPT_NEW_NEW_CONSTANT = "value"
+'''
+
+        gpt_converted, replacements = replace_model_patterns(gpt_test, gpt_model_patterns, new_gpt_model_patterns)
+        self.assertEqual(gpt_converted, gpt_expected)
+        # Replacements are empty here since GPT2 as been replaced by GPTNewNew in some instances and GPT_NEW_NEW
+        # in others.
+        self.assertEqual(replacements, "")
+
+        roberta_model_patterns = ModelPatterns("RoBERTa", "FacebookAI/roberta-base", model_camel_cased="Roberta")
+        new_roberta_model_patterns = ModelPatterns(
+            "RoBERTa-New", "huggingface/roberta-new-base", model_camel_cased="RobertaNew"
+        )
+        roberta_test = '''# Copied from transformers.models.bert.BertModel with Bert->Roberta
+class RobertaModel(RobertaPreTrainedModel):
+    """ The base RoBERTa model. """
+    checkpoint = FacebookAI/roberta-base
+    base_model_prefix = "roberta"
+        '''
+        roberta_expected = '''# Copied from transformers.models.bert.BertModel with Bert->RobertaNew
+class RobertaNewModel(RobertaNewPreTrainedModel):
+    """ The base RoBERTa-New model. """
+    checkpoint = huggingface/roberta-new-base
+    base_model_prefix = "roberta_new"
+        '''
+        roberta_converted, replacements = replace_model_patterns(
+            roberta_test, roberta_model_patterns, new_roberta_model_patterns
+        )
+        self.assertEqual(roberta_converted, roberta_expected)
+
+    def test_get_module_from_file(self):
+        self.assertEqual(
+            get_module_from_file("/git/transformers/src/transformers/models/bert/modeling_tf_bert.py"),
+            "transformers.models.bert.modeling_tf_bert",
+        )
+        self.assertEqual(
+            get_module_from_file("/transformers/models/gpt2/modeling_gpt2.py"),
+            "transformers.models.gpt2.modeling_gpt2",
+        )
+        with self.assertRaises(ValueError):
+            get_module_from_file("/models/gpt2/modeling_gpt2.py")
+
+    def test_duplicate_module(self):
+        bert_model_patterns = ModelPatterns("Bert", "google-bert/bert-base-cased")
+        new_bert_model_patterns = ModelPatterns("New Bert", "huggingface/bert-new-base")
+        bert_test = '''class TFBertPreTrainedModel(PreTrainedModel):
+    """
+    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
+    models.
+    """
+
+    config_class = BertConfig
+    load_tf_weights = load_tf_weights_in_bert
+    base_model_prefix = "bert"
+    is_parallelizable = True
+    supports_gradient_checkpointing = True
+
+BERT_CONSTANT = "value"
+'''
+        bert_expected = '''class TFNewBertPreTrainedModel(PreTrainedModel):
+    """
+    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
+    models.
+    """
+
+    config_class = NewBertConfig
+    load_tf_weights = load_tf_weights_in_new_bert
+    base_model_prefix = "new_bert"
+    is_parallelizable = True
+    supports_gradient_checkpointing = True
+
+NEW_BERT_CONSTANT = "value"
+'''
+        bert_expected_with_copied_from = (
+            "# Copied from transformers.bert_module.TFBertPreTrainedModel with Bert->NewBert,bert->new_bert\n"
+            + bert_expected
+        )
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            work_dir = os.path.join(tmp_dir, "transformers")
+            os.makedirs(work_dir)
+            file_name = os.path.join(work_dir, "bert_module.py")
+            dest_file_name = os.path.join(work_dir, "new_bert_module.py")
+
+            self.init_file(file_name, bert_test)
+            duplicate_module(file_name, bert_model_patterns, new_bert_model_patterns)
+            self.check_result(dest_file_name, bert_expected_with_copied_from)
+
+            self.init_file(file_name, bert_test)
+            duplicate_module(file_name, bert_model_patterns, new_bert_model_patterns, add_copied_from=False)
+            self.check_result(dest_file_name, bert_expected)
+
+    def test_duplicate_module_with_copied_from(self):
+        bert_model_patterns = ModelPatterns("Bert", "google-bert/bert-base-cased")
+        new_bert_model_patterns = ModelPatterns("New Bert", "huggingface/bert-new-base")
+        bert_test = '''# Copied from transformers.models.xxx.XxxModel with Xxx->Bert
+class TFBertPreTrainedModel(PreTrainedModel):
+    """
+    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
+    models.
+    """
+
+    config_class = BertConfig
+    load_tf_weights = load_tf_weights_in_bert
+    base_model_prefix = "bert"
+    is_parallelizable = True
+    supports_gradient_checkpointing = True
+
+BERT_CONSTANT = "value"
+'''
+        bert_expected = '''# Copied from transformers.models.xxx.XxxModel with Xxx->NewBert
+class TFNewBertPreTrainedModel(PreTrainedModel):
+    """
+    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
+    models.
+    """
+
+    config_class = NewBertConfig
+    load_tf_weights = load_tf_weights_in_new_bert
+    base_model_prefix = "new_bert"
+    is_parallelizable = True
+    supports_gradient_checkpointing = True
+
+NEW_BERT_CONSTANT = "value"
+'''
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            work_dir = os.path.join(tmp_dir, "transformers")
+            os.makedirs(work_dir)
+            file_name = os.path.join(work_dir, "bert_module.py")
+            dest_file_name = os.path.join(work_dir, "new_bert_module.py")
+
+            self.init_file(file_name, bert_test)
+            duplicate_module(file_name, bert_model_patterns, new_bert_model_patterns)
+            # There should not be a new Copied from statement, the old one should be adapted.
+            self.check_result(dest_file_name, bert_expected)
+
+            self.init_file(file_name, bert_test)
+            duplicate_module(file_name, bert_model_patterns, new_bert_model_patterns, add_copied_from=False)
+            self.check_result(dest_file_name, bert_expected)
+
+    def test_filter_framework_files(self):
+        files = ["modeling_bert.py", "modeling_tf_bert.py", "modeling_flax_bert.py", "configuration_bert.py"]
+        self.assertEqual(filter_framework_files(files), files)
+        self.assertEqual(set(filter_framework_files(files, ["pt", "tf", "flax"])), set(files))
+
+        self.assertEqual(set(filter_framework_files(files, ["pt"])), {"modeling_bert.py", "configuration_bert.py"})
+        self.assertEqual(set(filter_framework_files(files, ["tf"])), {"modeling_tf_bert.py", "configuration_bert.py"})
+        self.assertEqual(
+            set(filter_framework_files(files, ["flax"])), {"modeling_flax_bert.py", "configuration_bert.py"}
+        )
+
+        self.assertEqual(
+            set(filter_framework_files(files, ["pt", "tf"])),
+            {"modeling_tf_bert.py", "modeling_bert.py", "configuration_bert.py"},
+        )
+        self.assertEqual(
+            set(filter_framework_files(files, ["tf", "flax"])),
+            {"modeling_tf_bert.py", "modeling_flax_bert.py", "configuration_bert.py"},
+        )
+        self.assertEqual(
+            set(filter_framework_files(files, ["pt", "flax"])),
+            {"modeling_bert.py", "modeling_flax_bert.py", "configuration_bert.py"},
+        )
+
+    def test_get_model_files(self):
+        # BERT
+        bert_files = get_model_files("bert")
+
+        doc_file = str(Path(bert_files["doc_file"]).relative_to(REPO_PATH))
+        self.assertEqual(doc_file, "docs/source/en/model_doc/bert.md")
+
+        model_files = {str(Path(f).relative_to(REPO_PATH)) for f in bert_files["model_files"]}
+        self.assertEqual(model_files, BERT_MODEL_FILES)
+
+        self.assertEqual(bert_files["module_name"], "bert")
+
+        test_files = {str(Path(f).relative_to(REPO_PATH)) for f in bert_files["test_files"]}
+        bert_test_files = {
+            "tests/models/bert/test_tokenization_bert.py",
+            "tests/models/bert/test_modeling_bert.py",
+            "tests/models/bert/test_modeling_tf_bert.py",
+            "tests/models/bert/test_modeling_flax_bert.py",
+        }
+        self.assertEqual(test_files, bert_test_files)
+
+        # VIT
+        vit_files = get_model_files("vit")
+        doc_file = str(Path(vit_files["doc_file"]).relative_to(REPO_PATH))
+        self.assertEqual(doc_file, "docs/source/en/model_doc/vit.md")
+
+        model_files = {str(Path(f).relative_to(REPO_PATH)) for f in vit_files["model_files"]}
+        self.assertEqual(model_files, VIT_MODEL_FILES)
+
+        self.assertEqual(vit_files["module_name"], "vit")
+
+        test_files = {str(Path(f).relative_to(REPO_PATH)) for f in vit_files["test_files"]}
+        vit_test_files = {
+            "tests/models/vit/test_image_processing_vit.py",
+            "tests/models/vit/test_modeling_vit.py",
+            "tests/models/vit/test_modeling_tf_vit.py",
+            "tests/models/vit/test_modeling_flax_vit.py",
+        }
+        self.assertEqual(test_files, vit_test_files)
+
+        # Wav2Vec2
+        wav2vec2_files = get_model_files("wav2vec2")
+        doc_file = str(Path(wav2vec2_files["doc_file"]).relative_to(REPO_PATH))
+        self.assertEqual(doc_file, "docs/source/en/model_doc/wav2vec2.md")
+
+        model_files = {str(Path(f).relative_to(REPO_PATH)) for f in wav2vec2_files["model_files"]}
+        self.assertEqual(model_files, WAV2VEC2_MODEL_FILES)
+
+        self.assertEqual(wav2vec2_files["module_name"], "wav2vec2")
+
+        test_files = {str(Path(f).relative_to(REPO_PATH)) for f in wav2vec2_files["test_files"]}
+        wav2vec2_test_files = {
+            "tests/models/wav2vec2/test_feature_extraction_wav2vec2.py",
+            "tests/models/wav2vec2/test_modeling_wav2vec2.py",
+            "tests/models/wav2vec2/test_modeling_tf_wav2vec2.py",
+            "tests/models/wav2vec2/test_modeling_flax_wav2vec2.py",
+            "tests/models/wav2vec2/test_processor_wav2vec2.py",
+            "tests/models/wav2vec2/test_tokenization_wav2vec2.py",
+        }
+        self.assertEqual(test_files, wav2vec2_test_files)
+
+    def test_get_model_files_only_pt(self):
+        # BERT
+        bert_files = get_model_files("bert", frameworks=["pt"])
+
+        doc_file = str(Path(bert_files["doc_file"]).relative_to(REPO_PATH))
+        self.assertEqual(doc_file, "docs/source/en/model_doc/bert.md")
+
+        model_files = {str(Path(f).relative_to(REPO_PATH)) for f in bert_files["model_files"]}
+        bert_model_files = BERT_MODEL_FILES - {
+            "src/transformers/models/bert/modeling_tf_bert.py",
+            "src/transformers/models/bert/modeling_flax_bert.py",
+        }
+        self.assertEqual(model_files, bert_model_files)
+
+        self.assertEqual(bert_files["module_name"], "bert")
+
+        test_files = {str(Path(f).relative_to(REPO_PATH)) for f in bert_files["test_files"]}
+        bert_test_files = {
+            "tests/models/bert/test_tokenization_bert.py",
+            "tests/models/bert/test_modeling_bert.py",
+        }
+        self.assertEqual(test_files, bert_test_files)
+
+        # VIT
+        vit_files = get_model_files("vit", frameworks=["pt"])
+        doc_file = str(Path(vit_files["doc_file"]).relative_to(REPO_PATH))
+        self.assertEqual(doc_file, "docs/source/en/model_doc/vit.md")
+
+        model_files = {str(Path(f).relative_to(REPO_PATH)) for f in vit_files["model_files"]}
+        vit_model_files = VIT_MODEL_FILES - {
+            "src/transformers/models/vit/modeling_tf_vit.py",
+            "src/transformers/models/vit/modeling_flax_vit.py",
+        }
+        self.assertEqual(model_files, vit_model_files)
+
+        self.assertEqual(vit_files["module_name"], "vit")
+
+        test_files = {str(Path(f).relative_to(REPO_PATH)) for f in vit_files["test_files"]}
+        vit_test_files = {
+            "tests/models/vit/test_image_processing_vit.py",
+            "tests/models/vit/test_modeling_vit.py",
+        }
+        self.assertEqual(test_files, vit_test_files)
+
+        # Wav2Vec2
+        wav2vec2_files = get_model_files("wav2vec2", frameworks=["pt"])
+        doc_file = str(Path(wav2vec2_files["doc_file"]).relative_to(REPO_PATH))
+        self.assertEqual(doc_file, "docs/source/en/model_doc/wav2vec2.md")
+
+        model_files = {str(Path(f).relative_to(REPO_PATH)) for f in wav2vec2_files["model_files"]}
+        wav2vec2_model_files = WAV2VEC2_MODEL_FILES - {
+            "src/transformers/models/wav2vec2/modeling_tf_wav2vec2.py",
+            "src/transformers/models/wav2vec2/modeling_flax_wav2vec2.py",
+        }
+        self.assertEqual(model_files, wav2vec2_model_files)
+
+        self.assertEqual(wav2vec2_files["module_name"], "wav2vec2")
+
+        test_files = {str(Path(f).relative_to(REPO_PATH)) for f in wav2vec2_files["test_files"]}
+        wav2vec2_test_files = {
+            "tests/models/wav2vec2/test_feature_extraction_wav2vec2.py",
+            "tests/models/wav2vec2/test_modeling_wav2vec2.py",
+            "tests/models/wav2vec2/test_processor_wav2vec2.py",
+            "tests/models/wav2vec2/test_tokenization_wav2vec2.py",
+        }
+        self.assertEqual(test_files, wav2vec2_test_files)
+
+    def test_get_model_files_tf_and_flax(self):
+        # BERT
+        bert_files = get_model_files("bert", frameworks=["tf", "flax"])
+
+        doc_file = str(Path(bert_files["doc_file"]).relative_to(REPO_PATH))
+        self.assertEqual(doc_file, "docs/source/en/model_doc/bert.md")
+
+        model_files = {str(Path(f).relative_to(REPO_PATH)) for f in bert_files["model_files"]}
+        bert_model_files = BERT_MODEL_FILES - {"src/transformers/models/bert/modeling_bert.py"}
+        self.assertEqual(model_files, bert_model_files)
+
+        self.assertEqual(bert_files["module_name"], "bert")
+
+        test_files = {str(Path(f).relative_to(REPO_PATH)) for f in bert_files["test_files"]}
+        bert_test_files = {
+            "tests/models/bert/test_tokenization_bert.py",
+            "tests/models/bert/test_modeling_tf_bert.py",
+            "tests/models/bert/test_modeling_flax_bert.py",
+        }
+        self.assertEqual(test_files, bert_test_files)
+
+        # VIT
+        vit_files = get_model_files("vit", frameworks=["tf", "flax"])
+        doc_file = str(Path(vit_files["doc_file"]).relative_to(REPO_PATH))
+        self.assertEqual(doc_file, "docs/source/en/model_doc/vit.md")
+
+        model_files = {str(Path(f).relative_to(REPO_PATH)) for f in vit_files["model_files"]}
+        vit_model_files = VIT_MODEL_FILES - {"src/transformers/models/vit/modeling_vit.py"}
+        self.assertEqual(model_files, vit_model_files)
+
+        self.assertEqual(vit_files["module_name"], "vit")
+
+        test_files = {str(Path(f).relative_to(REPO_PATH)) for f in vit_files["test_files"]}
+        vit_test_files = {
+            "tests/models/vit/test_image_processing_vit.py",
+            "tests/models/vit/test_modeling_tf_vit.py",
+            "tests/models/vit/test_modeling_flax_vit.py",
+        }
+        self.assertEqual(test_files, vit_test_files)
+
+        # Wav2Vec2
+        wav2vec2_files = get_model_files("wav2vec2", frameworks=["tf", "flax"])
+        doc_file = str(Path(wav2vec2_files["doc_file"]).relative_to(REPO_PATH))
+        self.assertEqual(doc_file, "docs/source/en/model_doc/wav2vec2.md")
+
+        model_files = {str(Path(f).relative_to(REPO_PATH)) for f in wav2vec2_files["model_files"]}
+        wav2vec2_model_files = WAV2VEC2_MODEL_FILES - {"src/transformers/models/wav2vec2/modeling_wav2vec2.py"}
+        self.assertEqual(model_files, wav2vec2_model_files)
+
+        self.assertEqual(wav2vec2_files["module_name"], "wav2vec2")
+
+        test_files = {str(Path(f).relative_to(REPO_PATH)) for f in wav2vec2_files["test_files"]}
+        wav2vec2_test_files = {
+            "tests/models/wav2vec2/test_feature_extraction_wav2vec2.py",
+            "tests/models/wav2vec2/test_modeling_tf_wav2vec2.py",
+            "tests/models/wav2vec2/test_modeling_flax_wav2vec2.py",
+            "tests/models/wav2vec2/test_processor_wav2vec2.py",
+            "tests/models/wav2vec2/test_tokenization_wav2vec2.py",
+        }
+        self.assertEqual(test_files, wav2vec2_test_files)
+
+    def test_find_base_model_checkpoint(self):
+        self.assertEqual(find_base_model_checkpoint("bert"), "google-bert/bert-base-uncased")
+        self.assertEqual(find_base_model_checkpoint("gpt2"), "openai-community/gpt2")
+
+    def test_retrieve_model_classes(self):
+        gpt_classes = {k: set(v) for k, v in retrieve_model_classes("gpt2").items()}
+        expected_gpt_classes = {
+            "pt": {
+                "GPT2ForTokenClassification",
+                "GPT2Model",
+                "GPT2LMHeadModel",
+                "GPT2ForSequenceClassification",
+                "GPT2ForQuestionAnswering",
+            },
+            "tf": {"TFGPT2Model", "TFGPT2ForSequenceClassification", "TFGPT2LMHeadModel"},
+            "flax": {"FlaxGPT2Model", "FlaxGPT2LMHeadModel"},
+        }
+        self.assertEqual(gpt_classes, expected_gpt_classes)
+
+        del expected_gpt_classes["flax"]
+        gpt_classes = {k: set(v) for k, v in retrieve_model_classes("gpt2", frameworks=["pt", "tf"]).items()}
+        self.assertEqual(gpt_classes, expected_gpt_classes)
+
+        del expected_gpt_classes["pt"]
+        gpt_classes = {k: set(v) for k, v in retrieve_model_classes("gpt2", frameworks=["tf"]).items()}
+        self.assertEqual(gpt_classes, expected_gpt_classes)
+
+    def test_retrieve_info_for_model_with_bert(self):
+        bert_info = retrieve_info_for_model("bert")
+        bert_classes = [
+            "BertForTokenClassification",
+            "BertForQuestionAnswering",
+            "BertForNextSentencePrediction",
+            "BertForSequenceClassification",
+            "BertForMaskedLM",
+            "BertForMultipleChoice",
+            "BertModel",
+            "BertForPreTraining",
+            "BertLMHeadModel",
+        ]
+        expected_model_classes = {
+            "pt": set(bert_classes),
+            "tf": {f"TF{m}" for m in bert_classes},
+            "flax": {f"Flax{m}" for m in bert_classes[:-1] + ["BertForCausalLM"]},
+        }
+
+        self.assertEqual(set(bert_info["frameworks"]), {"pt", "tf", "flax"})
+        model_classes = {k: set(v) for k, v in bert_info["model_classes"].items()}
+        self.assertEqual(model_classes, expected_model_classes)
+
+        all_bert_files = bert_info["model_files"]
+        model_files = {str(Path(f).relative_to(REPO_PATH)) for f in all_bert_files["model_files"]}
+        self.assertEqual(model_files, BERT_MODEL_FILES)
+
+        test_files = {str(Path(f).relative_to(REPO_PATH)) for f in all_bert_files["test_files"]}
+        bert_test_files = {
+            "tests/models/bert/test_tokenization_bert.py",
+            "tests/models/bert/test_modeling_bert.py",
+            "tests/models/bert/test_modeling_tf_bert.py",
+            "tests/models/bert/test_modeling_flax_bert.py",
+        }
+        self.assertEqual(test_files, bert_test_files)
+
+        doc_file = str(Path(all_bert_files["doc_file"]).relative_to(REPO_PATH))
+        self.assertEqual(doc_file, "docs/source/en/model_doc/bert.md")
+
+        self.assertEqual(all_bert_files["module_name"], "bert")
+
+        bert_model_patterns = bert_info["model_patterns"]
+        self.assertEqual(bert_model_patterns.model_name, "BERT")
+        self.assertEqual(bert_model_patterns.checkpoint, "google-bert/bert-base-uncased")
+        self.assertEqual(bert_model_patterns.model_type, "bert")
+        self.assertEqual(bert_model_patterns.model_lower_cased, "bert")
+        self.assertEqual(bert_model_patterns.model_camel_cased, "Bert")
+        self.assertEqual(bert_model_patterns.model_upper_cased, "BERT")
+        self.assertEqual(bert_model_patterns.config_class, "BertConfig")
+        self.assertEqual(bert_model_patterns.tokenizer_class, "BertTokenizer")
+        self.assertIsNone(bert_model_patterns.feature_extractor_class)
+        self.assertIsNone(bert_model_patterns.processor_class)
+
+    def test_retrieve_info_for_model_with_vit(self):
+        vit_info = retrieve_info_for_model("vit")
+        vit_classes = ["ViTForImageClassification", "ViTModel"]
+        pt_only_classes = ["ViTForMaskedImageModeling"]
+        expected_model_classes = {
+            "pt": set(vit_classes + pt_only_classes),
+            "tf": {f"TF{m}" for m in vit_classes},
+            "flax": {f"Flax{m}" for m in vit_classes},
+        }
+
+        self.assertEqual(set(vit_info["frameworks"]), {"pt", "tf", "flax"})
+        model_classes = {k: set(v) for k, v in vit_info["model_classes"].items()}
+        self.assertEqual(model_classes, expected_model_classes)
+
+        all_vit_files = vit_info["model_files"]
+        model_files = {str(Path(f).relative_to(REPO_PATH)) for f in all_vit_files["model_files"]}
+        self.assertEqual(model_files, VIT_MODEL_FILES)
+
+        test_files = {str(Path(f).relative_to(REPO_PATH)) for f in all_vit_files["test_files"]}
+        vit_test_files = {
+            "tests/models/vit/test_image_processing_vit.py",
+            "tests/models/vit/test_modeling_vit.py",
+            "tests/models/vit/test_modeling_tf_vit.py",
+            "tests/models/vit/test_modeling_flax_vit.py",
+        }
+        self.assertEqual(test_files, vit_test_files)
+
+        doc_file = str(Path(all_vit_files["doc_file"]).relative_to(REPO_PATH))
+        self.assertEqual(doc_file, "docs/source/en/model_doc/vit.md")
+
+        self.assertEqual(all_vit_files["module_name"], "vit")
+
+        vit_model_patterns = vit_info["model_patterns"]
+        self.assertEqual(vit_model_patterns.model_name, "ViT")
+        self.assertEqual(vit_model_patterns.checkpoint, "google/vit-base-patch16-224-in21k")
+        self.assertEqual(vit_model_patterns.model_type, "vit")
+        self.assertEqual(vit_model_patterns.model_lower_cased, "vit")
+        self.assertEqual(vit_model_patterns.model_camel_cased, "ViT")
+        self.assertEqual(vit_model_patterns.model_upper_cased, "VIT")
+        self.assertEqual(vit_model_patterns.config_class, "ViTConfig")
+        self.assertEqual(vit_model_patterns.feature_extractor_class, "ViTFeatureExtractor")
+        self.assertEqual(vit_model_patterns.image_processor_class, "ViTImageProcessor")
+        self.assertIsNone(vit_model_patterns.tokenizer_class)
+        self.assertIsNone(vit_model_patterns.processor_class)
+
+    def test_retrieve_info_for_model_with_wav2vec2(self):
+        wav2vec2_info = retrieve_info_for_model("wav2vec2")
+        wav2vec2_classes = [
+            "Wav2Vec2Model",
+            "Wav2Vec2ForPreTraining",
+            "Wav2Vec2ForAudioFrameClassification",
+            "Wav2Vec2ForCTC",
+            "Wav2Vec2ForMaskedLM",
+            "Wav2Vec2ForSequenceClassification",
+            "Wav2Vec2ForXVector",
+        ]
+        expected_model_classes = {
+            "pt": set(wav2vec2_classes),
+            "tf": {f"TF{m}" for m in [wav2vec2_classes[0], wav2vec2_classes[-2]]},
+            "flax": {f"Flax{m}" for m in wav2vec2_classes[:2]},
+        }
+
+        self.assertEqual(set(wav2vec2_info["frameworks"]), {"pt", "tf", "flax"})
+        model_classes = {k: set(v) for k, v in wav2vec2_info["model_classes"].items()}
+        self.assertEqual(model_classes, expected_model_classes)
+
+        all_wav2vec2_files = wav2vec2_info["model_files"]
+        model_files = {str(Path(f).relative_to(REPO_PATH)) for f in all_wav2vec2_files["model_files"]}
+        self.assertEqual(model_files, WAV2VEC2_MODEL_FILES)
+
+        test_files = {str(Path(f).relative_to(REPO_PATH)) for f in all_wav2vec2_files["test_files"]}
+        wav2vec2_test_files = {
+            "tests/models/wav2vec2/test_feature_extraction_wav2vec2.py",
+            "tests/models/wav2vec2/test_modeling_wav2vec2.py",
+            "tests/models/wav2vec2/test_modeling_tf_wav2vec2.py",
+            "tests/models/wav2vec2/test_modeling_flax_wav2vec2.py",
+            "tests/models/wav2vec2/test_processor_wav2vec2.py",
+            "tests/models/wav2vec2/test_tokenization_wav2vec2.py",
+        }
+        self.assertEqual(test_files, wav2vec2_test_files)
+
+        doc_file = str(Path(all_wav2vec2_files["doc_file"]).relative_to(REPO_PATH))
+        self.assertEqual(doc_file, "docs/source/en/model_doc/wav2vec2.md")
+
+        self.assertEqual(all_wav2vec2_files["module_name"], "wav2vec2")
+
+        wav2vec2_model_patterns = wav2vec2_info["model_patterns"]
+        self.assertEqual(wav2vec2_model_patterns.model_name, "Wav2Vec2")
+        self.assertEqual(wav2vec2_model_patterns.checkpoint, "facebook/wav2vec2-base-960h")
+        self.assertEqual(wav2vec2_model_patterns.model_type, "wav2vec2")
+        self.assertEqual(wav2vec2_model_patterns.model_lower_cased, "wav2vec2")
+        self.assertEqual(wav2vec2_model_patterns.model_camel_cased, "Wav2Vec2")
+        self.assertEqual(wav2vec2_model_patterns.model_upper_cased, "WAV2VEC2")
+        self.assertEqual(wav2vec2_model_patterns.config_class, "Wav2Vec2Config")
+        self.assertEqual(wav2vec2_model_patterns.feature_extractor_class, "Wav2Vec2FeatureExtractor")
+        self.assertEqual(wav2vec2_model_patterns.processor_class, "Wav2Vec2Processor")
+        self.assertEqual(wav2vec2_model_patterns.tokenizer_class, "Wav2Vec2CTCTokenizer")
+
+    def test_clean_frameworks_in_init_with_gpt(self):
+        test_init = """
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule, is_flax_available, is_tf_available, is_tokenizers_available, is_torch_available
+
+_import_structure = {
+    "configuration_gpt2": ["GPT2Config", "GPT2OnnxConfig"],
+    "tokenization_gpt2": ["GPT2Tokenizer"],
+}
+
+try:
+    if not is_tokenizers_available():
+        raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+    pass
+else:
+    _import_structure["tokenization_gpt2_fast"] = ["GPT2TokenizerFast"]
+
+try:
+    if not is_torch_available():
+        raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+    pass
+else:
+    _import_structure["modeling_gpt2"] = ["GPT2Model"]
+
+try:
+    if not is_tf_available():
+        raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+    pass
+else:
+    _import_structure["modeling_tf_gpt2"] = ["TFGPT2Model"]
+
+try:
+    if not is_flax_available():
+        raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+    pass
+else:
+    _import_structure["modeling_flax_gpt2"] = ["FlaxGPT2Model"]
+
+if TYPE_CHECKING:
+    from .configuration_gpt2 import GPT2Config, GPT2OnnxConfig
+    from .tokenization_gpt2 import GPT2Tokenizer
+
+    try:
+        if not is_tokenizers_available():
+            raise OptionalDependencyNotAvailable()
+    except OptionalDependencyNotAvailable:
+        pass
+    else:
+        from .tokenization_gpt2_fast import GPT2TokenizerFast
+
+    try:
+        if not is_torch_available():
+            raise OptionalDependencyNotAvailable()
+    except OptionalDependencyNotAvailable:
+        pass
+    else:
+        from .modeling_gpt2 import GPT2Model
+
+    try:
+        if not is_tf_available():
+            raise OptionalDependencyNotAvailable()
+    except OptionalDependencyNotAvailable:
+        pass
+    else:
+        from .modeling_tf_gpt2 import TFGPT2Model
+
+    try:
+        if not is_flax_available():
+            raise OptionalDependencyNotAvailable()
+    except OptionalDependencyNotAvailable:
+        pass
+    else:
+        from .modeling_flax_gpt2 import FlaxGPT2Model
+
+else:
+    import sys
+
+    sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure)
+"""
+
+        init_no_tokenizer = """
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule, is_flax_available, is_tf_available, is_torch_available
+
+_import_structure = {
+    "configuration_gpt2": ["GPT2Config", "GPT2OnnxConfig"],
+}
+
+try:
+    if not is_torch_available():
+        raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+    pass
+else:
+    _import_structure["modeling_gpt2"] = ["GPT2Model"]
+
+try:
+    if not is_tf_available():
+        raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+    pass
+else:
+    _import_structure["modeling_tf_gpt2"] = ["TFGPT2Model"]
+
+try:
+    if not is_flax_available():
+        raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+    pass
+else:
+    _import_structure["modeling_flax_gpt2"] = ["FlaxGPT2Model"]
+
+if TYPE_CHECKING:
+    from .configuration_gpt2 import GPT2Config, GPT2OnnxConfig
+
+    try:
+        if not is_torch_available():
+            raise OptionalDependencyNotAvailable()
+    except OptionalDependencyNotAvailable:
+        pass
+    else:
+        from .modeling_gpt2 import GPT2Model
+
+    try:
+        if not is_tf_available():
+            raise OptionalDependencyNotAvailable()
+    except OptionalDependencyNotAvailable:
+        pass
+    else:
+        from .modeling_tf_gpt2 import TFGPT2Model
+
+    try:
+        if not is_flax_available():
+            raise OptionalDependencyNotAvailable()
+    except OptionalDependencyNotAvailable:
+        pass
+    else:
+        from .modeling_flax_gpt2 import FlaxGPT2Model
+
+else:
+    import sys
+
+    sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure)
+"""
+
+        init_pt_only = """
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule, is_tokenizers_available, is_torch_available
+
+_import_structure = {
+    "configuration_gpt2": ["GPT2Config", "GPT2OnnxConfig"],
+    "tokenization_gpt2": ["GPT2Tokenizer"],
+}
+
+try:
+    if not is_tokenizers_available():
+        raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+    pass
+else:
+    _import_structure["tokenization_gpt2_fast"] = ["GPT2TokenizerFast"]
+
+try:
+    if not is_torch_available():
+        raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+    pass
+else:
+    _import_structure["modeling_gpt2"] = ["GPT2Model"]
+
+if TYPE_CHECKING:
+    from .configuration_gpt2 import GPT2Config, GPT2OnnxConfig
+    from .tokenization_gpt2 import GPT2Tokenizer
+
+    try:
+        if not is_tokenizers_available():
+            raise OptionalDependencyNotAvailable()
+    except OptionalDependencyNotAvailable:
+        pass
+    else:
+        from .tokenization_gpt2_fast import GPT2TokenizerFast
+
+    try:
+        if not is_torch_available():
+            raise OptionalDependencyNotAvailable()
+    except OptionalDependencyNotAvailable:
+        pass
+    else:
+        from .modeling_gpt2 import GPT2Model
+
+else:
+    import sys
+
+    sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure)
+"""
+
+        init_pt_only_no_tokenizer = """
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule, is_torch_available
+
+_import_structure = {
+    "configuration_gpt2": ["GPT2Config", "GPT2OnnxConfig"],
+}
+
+try:
+    if not is_torch_available():
+        raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+    pass
+else:
+    _import_structure["modeling_gpt2"] = ["GPT2Model"]
+
+if TYPE_CHECKING:
+    from .configuration_gpt2 import GPT2Config, GPT2OnnxConfig
+
+    try:
+        if not is_torch_available():
+            raise OptionalDependencyNotAvailable()
+    except OptionalDependencyNotAvailable:
+        pass
+    else:
+        from .modeling_gpt2 import GPT2Model
+
+else:
+    import sys
+
+    sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure)
+"""
+
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            file_name = os.path.join(tmp_dir, "../__init__.py")
+
+            self.init_file(file_name, test_init)
+            clean_frameworks_in_init(file_name, keep_processing=False)
+            self.check_result(file_name, init_no_tokenizer)
+
+            self.init_file(file_name, test_init)
+            clean_frameworks_in_init(file_name, frameworks=["pt"])
+            self.check_result(file_name, init_pt_only)
+
+            self.init_file(file_name, test_init)
+            clean_frameworks_in_init(file_name, frameworks=["pt"], keep_processing=False)
+            self.check_result(file_name, init_pt_only_no_tokenizer)
+
+    def test_clean_frameworks_in_init_with_vit(self):
+        test_init = """
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule, is_flax_available, is_tf_available, is_torch_available, is_vision_available
+
+_import_structure = {
+    "configuration_vit": ["ViTConfig"],
+}
+
+try:
+    if not is_vision_available():
+        raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+    pass
+else:
+    _import_structure["image_processing_vit"] = ["ViTImageProcessor"]
+
+try:
+    if not is_torch_available():
+        raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+    pass
+else:
+    _import_structure["modeling_vit"] = ["ViTModel"]
+
+try:
+    if not is_tf_available():
+        raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+    pass
+else:
+    _import_structure["modeling_tf_vit"] = ["TFViTModel"]
+
+try:
+    if not is_flax_available():
+        raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+    pass
+else:
+    _import_structure["modeling_flax_vit"] = ["FlaxViTModel"]
+
+if TYPE_CHECKING:
+    from .configuration_vit import ViTConfig
+
+    try:
+        if not is_vision_available():
+            raise OptionalDependencyNotAvailable()
+    except OptionalDependencyNotAvailable:
+        pass
+    else:
+        from .image_processing_vit import ViTImageProcessor
+
+    try:
+        if not is_torch_available():
+            raise OptionalDependencyNotAvailable()
+    except OptionalDependencyNotAvailable:
+        pass
+    else:
+        from .modeling_vit import ViTModel
+
+    try:
+        if not is_tf_available():
+            raise OptionalDependencyNotAvailable()
+    except OptionalDependencyNotAvailable:
+        pass
+    else:
+        from .modeling_tf_vit import TFViTModel
+
+    try:
+        if not is_flax_available():
+            raise OptionalDependencyNotAvailable()
+    except OptionalDependencyNotAvailable:
+        pass
+    else:
+        from .modeling_flax_vit import FlaxViTModel
+
+else:
+    import sys
+
+    sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure)
+"""
+
+        init_no_feature_extractor = """
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule, is_flax_available, is_tf_available, is_torch_available
+
+_import_structure = {
+    "configuration_vit": ["ViTConfig"],
+}
+
+try:
+    if not is_torch_available():
+        raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+    pass
+else:
+    _import_structure["modeling_vit"] = ["ViTModel"]
+
+try:
+    if not is_tf_available():
+        raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+    pass
+else:
+    _import_structure["modeling_tf_vit"] = ["TFViTModel"]
+
+try:
+    if not is_flax_available():
+        raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+    pass
+else:
+    _import_structure["modeling_flax_vit"] = ["FlaxViTModel"]
+
+if TYPE_CHECKING:
+    from .configuration_vit import ViTConfig
+
+    try:
+        if not is_torch_available():
+            raise OptionalDependencyNotAvailable()
+    except OptionalDependencyNotAvailable:
+        pass
+    else:
+        from .modeling_vit import ViTModel
+
+    try:
+        if not is_tf_available():
+            raise OptionalDependencyNotAvailable()
+    except OptionalDependencyNotAvailable:
+        pass
+    else:
+        from .modeling_tf_vit import TFViTModel
+
+    try:
+        if not is_flax_available():
+            raise OptionalDependencyNotAvailable()
+    except OptionalDependencyNotAvailable:
+        pass
+    else:
+        from .modeling_flax_vit import FlaxViTModel
+
+else:
+    import sys
+
+    sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure)
+"""
+
+        init_pt_only = """
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule, is_torch_available, is_vision_available
+
+_import_structure = {
+    "configuration_vit": ["ViTConfig"],
+}
+
+try:
+    if not is_vision_available():
+        raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+    pass
+else:
+    _import_structure["image_processing_vit"] = ["ViTImageProcessor"]
+
+try:
+    if not is_torch_available():
+        raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+    pass
+else:
+    _import_structure["modeling_vit"] = ["ViTModel"]
+
+if TYPE_CHECKING:
+    from .configuration_vit import ViTConfig
+
+    try:
+        if not is_vision_available():
+            raise OptionalDependencyNotAvailable()
+    except OptionalDependencyNotAvailable:
+        pass
+    else:
+        from .image_processing_vit import ViTImageProcessor
+
+    try:
+        if not is_torch_available():
+            raise OptionalDependencyNotAvailable()
+    except OptionalDependencyNotAvailable:
+        pass
+    else:
+        from .modeling_vit import ViTModel
+
+else:
+    import sys
+
+    sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure)
+"""
+
+        init_pt_only_no_feature_extractor = """
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule, is_torch_available
+
+_import_structure = {
+    "configuration_vit": ["ViTConfig"],
+}
+
+try:
+    if not is_torch_available():
+        raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+    pass
+else:
+    _import_structure["modeling_vit"] = ["ViTModel"]
+
+if TYPE_CHECKING:
+    from .configuration_vit import ViTConfig
+
+    try:
+        if not is_torch_available():
+            raise OptionalDependencyNotAvailable()
+    except OptionalDependencyNotAvailable:
+        pass
+    else:
+        from .modeling_vit import ViTModel
+
+else:
+    import sys
+
+    sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure)
+"""
+
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            file_name = os.path.join(tmp_dir, "../__init__.py")
+
+            self.init_file(file_name, test_init)
+            clean_frameworks_in_init(file_name, keep_processing=False)
+            self.check_result(file_name, init_no_feature_extractor)
+
+            self.init_file(file_name, test_init)
+            clean_frameworks_in_init(file_name, frameworks=["pt"])
+            self.check_result(file_name, init_pt_only)
+
+            self.init_file(file_name, test_init)
+            clean_frameworks_in_init(file_name, frameworks=["pt"], keep_processing=False)
+            self.check_result(file_name, init_pt_only_no_feature_extractor)
+
+    def test_duplicate_doc_file(self):
+        test_doc = """
+# GPT2
+
+## Overview
+
+Overview of the model.
+
+## GPT2Config
+
+[[autodoc]] GPT2Config
+
+## GPT2Tokenizer
+
+[[autodoc]] GPT2Tokenizer
+    - save_vocabulary
+
+## GPT2TokenizerFast
+
+[[autodoc]] GPT2TokenizerFast
+
+## GPT2 specific outputs
+
+[[autodoc]] models.gpt2.modeling_gpt2.GPT2DoubleHeadsModelOutput
+
+[[autodoc]] models.gpt2.modeling_tf_gpt2.TFGPT2DoubleHeadsModelOutput
+
+## GPT2Model
+
+[[autodoc]] GPT2Model
+    - forward
+
+## TFGPT2Model
+
+[[autodoc]] TFGPT2Model
+    - call
+
+## FlaxGPT2Model
+
+[[autodoc]] FlaxGPT2Model
+    - __call__
+
+"""
+        test_new_doc = """
+# GPT-New New
+
+## Overview
+
+The GPT-New New model was proposed in [<INSERT PAPER NAME HERE>](<INSERT PAPER LINK HERE>) by <INSERT AUTHORS HERE>.
+<INSERT SHORT SUMMARY HERE>
+
+The abstract from the paper is the following:
+
+*<INSERT PAPER ABSTRACT HERE>*
+
+Tips:
+
+<INSERT TIPS ABOUT MODEL HERE>
+
+This model was contributed by [INSERT YOUR HF USERNAME HERE](https://huggingface.co/<INSERT YOUR HF USERNAME HERE>).
+The original code can be found [here](<INSERT LINK TO GITHUB REPO HERE>).
+
+
+## GPTNewNewConfig
+
+[[autodoc]] GPTNewNewConfig
+
+## GPTNewNewTokenizer
+
+[[autodoc]] GPTNewNewTokenizer
+    - save_vocabulary
+
+## GPTNewNewTokenizerFast
+
+[[autodoc]] GPTNewNewTokenizerFast
+
+## GPTNewNew specific outputs
+
+[[autodoc]] models.gpt_new_new.modeling_gpt_new_new.GPTNewNewDoubleHeadsModelOutput
+
+[[autodoc]] models.gpt_new_new.modeling_tf_gpt_new_new.TFGPTNewNewDoubleHeadsModelOutput
+
+## GPTNewNewModel
+
+[[autodoc]] GPTNewNewModel
+    - forward
+
+## TFGPTNewNewModel
+
+[[autodoc]] TFGPTNewNewModel
+    - call
+
+## FlaxGPTNewNewModel
+
+[[autodoc]] FlaxGPTNewNewModel
+    - __call__
+
+"""
+
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            doc_file = os.path.join(tmp_dir, "gpt2.md")
+            new_doc_file = os.path.join(tmp_dir, "gpt-new-new.md")
+
+            gpt2_model_patterns = ModelPatterns("GPT2", "gpt2", tokenizer_class="GPT2Tokenizer")
+            new_model_patterns = ModelPatterns(
+                "GPT-New New", "huggingface/gpt-new-new", tokenizer_class="GPTNewNewTokenizer"
+            )
+
+            self.init_file(doc_file, test_doc)
+            duplicate_doc_file(doc_file, gpt2_model_patterns, new_model_patterns)
+            self.check_result(new_doc_file, test_new_doc)
+
+            test_new_doc_pt_only = test_new_doc.replace(
+                """
+## TFGPTNewNewModel
+
+[[autodoc]] TFGPTNewNewModel
+    - call
+
+## FlaxGPTNewNewModel
+
+[[autodoc]] FlaxGPTNewNewModel
+    - __call__
+
+""",
+                "",
+            )
+            self.init_file(doc_file, test_doc)
+            duplicate_doc_file(doc_file, gpt2_model_patterns, new_model_patterns, frameworks=["pt"])
+            self.check_result(new_doc_file, test_new_doc_pt_only)
+
+            test_new_doc_no_tok = test_new_doc.replace(
+                """
+## GPTNewNewTokenizer
+
+[[autodoc]] GPTNewNewTokenizer
+    - save_vocabulary
+
+## GPTNewNewTokenizerFast
+
+[[autodoc]] GPTNewNewTokenizerFast
+""",
+                "",
+            )
+            new_model_patterns = ModelPatterns(
+                "GPT-New New", "huggingface/gpt-new-new", tokenizer_class="GPT2Tokenizer"
+            )
+            self.init_file(doc_file, test_doc)
+            duplicate_doc_file(doc_file, gpt2_model_patterns, new_model_patterns)
+            print(test_new_doc_no_tok)
+            self.check_result(new_doc_file, test_new_doc_no_tok)
+
+            test_new_doc_pt_only_no_tok = test_new_doc_no_tok.replace(
+                """
+## TFGPTNewNewModel
+
+[[autodoc]] TFGPTNewNewModel
+    - call
+
+## FlaxGPTNewNewModel
+
+[[autodoc]] FlaxGPTNewNewModel
+    - __call__
+
+""",
+                "",
+            )
+            self.init_file(doc_file, test_doc)
+            duplicate_doc_file(doc_file, gpt2_model_patterns, new_model_patterns, frameworks=["pt"])
+            self.check_result(new_doc_file, test_new_doc_pt_only_no_tok)

docs/transformers/tests/utils/test_audio_utils.py

@@ -0,0 +1,1751 @@
+# Copyright 2023 HuggingFace Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+import pytest
+
+from transformers.audio_utils import (
+    amplitude_to_db,
+    amplitude_to_db_batch,
+    chroma_filter_bank,
+    hertz_to_mel,
+    mel_filter_bank,
+    mel_to_hertz,
+    power_to_db,
+    power_to_db_batch,
+    spectrogram,
+    spectrogram_batch,
+    window_function,
+)
+from transformers.testing_utils import is_librosa_available, require_librosa
+
+
+if is_librosa_available():
+    from librosa.filters import chroma
+
+
+class AudioUtilsFunctionTester(unittest.TestCase):
+    # will be set in `def _load_datasamples`
+    _dataset = None
+
+    def test_hertz_to_mel(self):
+        self.assertEqual(hertz_to_mel(0.0), 0.0)
+        self.assertAlmostEqual(hertz_to_mel(100), 150.48910241)
+
+        inputs = np.array([100, 200])
+        expected = np.array([150.48910241, 283.22989816])
+        self.assertTrue(np.allclose(hertz_to_mel(inputs), expected))
+
+        self.assertEqual(hertz_to_mel(0.0, "slaney"), 0.0)
+        self.assertEqual(hertz_to_mel(100, "slaney"), 1.5)
+
+        inputs = np.array([60, 100, 200, 1000, 1001, 2000])
+        expected = np.array([0.9, 1.5, 3.0, 15.0, 15.01453781, 25.08188016])
+        self.assertTrue(np.allclose(hertz_to_mel(inputs, "slaney"), expected))
+
+        inputs = np.array([60, 100, 200, 1000, 1001, 2000])
+        expected = np.array([92.6824, 150.4899, 283.2313, 999.9907, 1000.6534, 1521.3674])
+        self.assertTrue(np.allclose(hertz_to_mel(inputs, "kaldi"), expected))
+
+        with pytest.raises(ValueError):
+            hertz_to_mel(100, mel_scale=None)
+
+    def test_mel_to_hertz(self):
+        self.assertEqual(mel_to_hertz(0.0), 0.0)
+        self.assertAlmostEqual(mel_to_hertz(150.48910241), 100)
+
+        inputs = np.array([150.48910241, 283.22989816])
+        expected = np.array([100, 200])
+        self.assertTrue(np.allclose(mel_to_hertz(inputs), expected))
+
+        self.assertEqual(mel_to_hertz(0.0, "slaney"), 0.0)
+        self.assertEqual(mel_to_hertz(1.5, "slaney"), 100)
+
+        inputs = np.array([0.9, 1.5, 3.0, 15.0, 15.01453781, 25.08188016])
+        expected = np.array([60, 100, 200, 1000, 1001, 2000])
+        self.assertTrue(np.allclose(mel_to_hertz(inputs, "slaney"), expected))
+
+        inputs = np.array([92.6824, 150.4899, 283.2313, 999.9907, 1000.6534, 1521.3674])
+        expected = np.array([60, 100, 200, 1000, 1001, 2000])
+        self.assertTrue(np.allclose(mel_to_hertz(inputs, "kaldi"), expected))
+
+        with pytest.raises(ValueError):
+            mel_to_hertz(100, mel_scale=None)
+
+    def test_mel_filter_bank_shape(self):
+        mel_filters = mel_filter_bank(
+            num_frequency_bins=513,
+            num_mel_filters=13,
+            min_frequency=100,
+            max_frequency=4000,
+            sampling_rate=16000,
+            norm=None,
+            mel_scale="htk",
+        )
+        self.assertEqual(mel_filters.shape, (513, 13))
+
+        mel_filters = mel_filter_bank(
+            num_frequency_bins=513,
+            num_mel_filters=13,
+            min_frequency=100,
+            max_frequency=4000,
+            sampling_rate=16000,
+            norm="slaney",
+            mel_scale="slaney",
+        )
+        self.assertEqual(mel_filters.shape, (513, 13))
+
+        mel_filters = mel_filter_bank(
+            num_frequency_bins=513,
+            num_mel_filters=13,
+            min_frequency=100,
+            max_frequency=4000,
+            sampling_rate=16000,
+            norm="slaney",
+            mel_scale="slaney",
+            triangularize_in_mel_space=True,
+        )
+        self.assertEqual(mel_filters.shape, (513, 13))
+
+    def test_mel_filter_bank_htk(self):
+        mel_filters = mel_filter_bank(
+            num_frequency_bins=16,
+            num_mel_filters=4,
+            min_frequency=0,
+            max_frequency=2000,
+            sampling_rate=4000,
+            norm=None,
+            mel_scale="htk",
+        )
+        # fmt: off
+        expected = np.array([
+            [0.0       , 0.0       , 0.0       , 0.0       ],
+            [0.61454786, 0.0       , 0.0       , 0.0       ],
+            [0.82511046, 0.17488954, 0.0       , 0.0       ],
+            [0.35597035, 0.64402965, 0.0       , 0.0       ],
+            [0.0       , 0.91360726, 0.08639274, 0.0       ],
+            [0.0       , 0.55547007, 0.44452993, 0.0       ],
+            [0.0       , 0.19733289, 0.80266711, 0.0       ],
+            [0.0       , 0.0       , 0.87724349, 0.12275651],
+            [0.0       , 0.0       , 0.6038449 , 0.3961551 ],
+            [0.0       , 0.0       , 0.33044631, 0.66955369],
+            [0.0       , 0.0       , 0.05704771, 0.94295229],
+            [0.0       , 0.0       , 0.0       , 0.83483975],
+            [0.0       , 0.0       , 0.0       , 0.62612982],
+            [0.0       , 0.0       , 0.0       , 0.41741988],
+            [0.0       , 0.0       , 0.0       , 0.20870994],
+            [0.0       , 0.0       , 0.0       , 0.0       ]
+        ])
+        # fmt: on
+        self.assertTrue(np.allclose(mel_filters, expected))
+
+    def test_mel_filter_bank_slaney(self):
+        mel_filters = mel_filter_bank(
+            num_frequency_bins=16,
+            num_mel_filters=4,
+            min_frequency=0,
+            max_frequency=2000,
+            sampling_rate=4000,
+            norm=None,
+            mel_scale="slaney",
+        )
+        # fmt: off
+        expected = np.array([
+            [0.0       , 0.0       , 0.0       , 0.0       ],
+            [0.39869419, 0.0       , 0.0       , 0.0       ],
+            [0.79738839, 0.0       , 0.0       , 0.0       ],
+            [0.80391742, 0.19608258, 0.0       , 0.0       ],
+            [0.40522322, 0.59477678, 0.0       , 0.0       ],
+            [0.00652903, 0.99347097, 0.0       , 0.0       ],
+            [0.0       , 0.60796161, 0.39203839, 0.0       ],
+            [0.0       , 0.20939631, 0.79060369, 0.0       ],
+            [0.0       , 0.0       , 0.84685344, 0.15314656],
+            [0.0       , 0.0       , 0.52418477, 0.47581523],
+            [0.0       , 0.0       , 0.2015161 , 0.7984839 ],
+            [0.0       , 0.0       , 0.0       , 0.9141874 ],
+            [0.0       , 0.0       , 0.0       , 0.68564055],
+            [0.0       , 0.0       , 0.0       , 0.4570937 ],
+            [0.0       , 0.0       , 0.0       , 0.22854685],
+            [0.0       , 0.0       , 0.0       , 0.0       ]
+        ])
+        # fmt: on
+        self.assertTrue(np.allclose(mel_filters, expected))
+
+    def test_mel_filter_bank_kaldi(self):
+        mel_filters = mel_filter_bank(
+            num_frequency_bins=16,
+            num_mel_filters=4,
+            min_frequency=0,
+            max_frequency=2000,
+            sampling_rate=4000,
+            norm=None,
+            mel_scale="kaldi",
+            triangularize_in_mel_space=True,
+        )
+        # fmt: off
+        # here the expected values from torchaudio.compliance.kaldi.get_mel_banks
+        # note that we compute values in float64 while they do it in float32
+        expected = np.array(
+            [
+                [0.0000000000000000, 0.0000000000000000, 0.0000000000000000, 0.0000000000000000],
+                [0.6457883715629578, 0.0000000000000000, 0.0000000000000000, 0.0000000000000000],
+                [0.8044781088829041, 0.1955219060182571, 0.0000000000000000, 0.0000000000000000],
+                [0.3258901536464691, 0.6741098165512085, 0.0000000000000000, 0.0000000000000000],
+                [0.0000000000000000, 0.9021250009536743, 0.0978749766945839, 0.0000000000000000],
+                [0.0000000000000000, 0.5219038724899292, 0.4780961275100708, 0.0000000000000000],
+                [0.0000000000000000, 0.1771058291196823, 0.8228941559791565, 0.0000000000000000],
+                [0.0000000000000000, 0.0000000000000000, 0.8616894483566284, 0.1383105516433716],
+                [0.0000000000000000, 0.0000000000000000, 0.5710380673408508, 0.4289619624614716],
+                [0.0000000000000000, 0.0000000000000000, 0.3015440106391907, 0.6984559893608093],
+                [0.0000000000000000, 0.0000000000000000, 0.0503356307744980, 0.9496643543243408],
+                [0.0000000000000000, 0.0000000000000000, 0.0000000000000000, 0.8150880336761475],
+                [0.0000000000000000, 0.0000000000000000, 0.0000000000000000, 0.5938932299613953],
+                [0.0000000000000000, 0.0000000000000000, 0.0000000000000000, 0.3851676583290100],
+                [0.0000000000000000, 0.0000000000000000, 0.0000000000000000, 0.1875794380903244],
+            ],
+            dtype=np.float64,
+        )
+        # fmt: on
+
+        # kaldi implementation does not compute values for last fft bin
+        # indeed, they enforce max_frequency <= sampling_rate / 2 and
+        # therefore they know that last fft bin filter bank values will be all 0
+        # and pad after with zeros
+        # to comply with our API for `mel_filter_bank`, we need to also pad here
+        expected = np.pad(expected, ((0, 1), (0, 0)))
+
+        self.assertTrue(np.allclose(mel_filters, expected))
+
+    def test_mel_filter_bank_slaney_norm(self):
+        mel_filters = mel_filter_bank(
+            num_frequency_bins=16,
+            num_mel_filters=4,
+            min_frequency=0,
+            max_frequency=2000,
+            sampling_rate=4000,
+            norm="slaney",
+            mel_scale="slaney",
+        )
+        # fmt: off
+        expected = np.array([
+            [0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0.00000000e+00],
+            [1.19217795e-03, 0.00000000e+00, 0.00000000e+00, 0.00000000e+00],
+            [2.38435591e-03, 0.00000000e+00, 0.00000000e+00, 0.00000000e+00],
+            [2.40387905e-03, 5.86232616e-04, 0.00000000e+00, 0.00000000e+00],
+            [1.21170110e-03, 1.77821783e-03, 0.00000000e+00, 0.00000000e+00],
+            [1.95231437e-05, 2.97020305e-03, 0.00000000e+00, 0.00000000e+00],
+            [0.00000000e+00, 1.81763684e-03, 1.04857612e-03, 0.00000000e+00],
+            [0.00000000e+00, 6.26036972e-04, 2.11460963e-03, 0.00000000e+00],
+            [0.00000000e+00, 0.00000000e+00, 2.26505954e-03, 3.07332945e-04],
+            [0.00000000e+00, 0.00000000e+00, 1.40202503e-03, 9.54861093e-04],
+            [0.00000000e+00, 0.00000000e+00, 5.38990521e-04, 1.60238924e-03],
+            [0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 1.83458185e-03],
+            [0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 1.37593638e-03],
+            [0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 9.17290923e-04],
+            [0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 4.58645462e-04],
+            [0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0.00000000e+00]
+        ])
+        # fmt: on
+        self.assertTrue(np.allclose(mel_filters, expected))
+
+    def test_window_function(self):
+        window = window_function(16, "hann")
+        self.assertEqual(len(window), 16)
+
+        # fmt: off
+        expected = np.array([
+            0.0, 0.03806023, 0.14644661, 0.30865828, 0.5, 0.69134172, 0.85355339, 0.96193977,
+            1.0, 0.96193977, 0.85355339, 0.69134172, 0.5, 0.30865828, 0.14644661, 0.03806023,
+        ])
+        # fmt: on
+        self.assertTrue(np.allclose(window, expected))
+
+    def _load_datasamples(self, num_samples):
+        from datasets import load_dataset
+
+        if self._dataset is None:
+            self._dataset = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")
+        speech_samples = self._dataset.sort("id").select(range(num_samples))[:num_samples]["audio"]
+        return [x["array"] for x in speech_samples]
+
+    def test_spectrogram_impulse(self):
+        waveform = np.zeros(40)
+        waveform[9] = 1.0  # impulse shifted in time
+
+        spec = spectrogram(
+            waveform,
+            window_function(12, "hann", frame_length=16),
+            frame_length=16,
+            hop_length=4,
+            power=1.0,
+            center=True,
+            pad_mode="reflect",
+            onesided=True,
+        )
+        self.assertEqual(spec.shape, (9, 11))
+
+        expected = np.array([[0.0, 0.0669873, 0.9330127, 0.5, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]])
+        self.assertTrue(np.allclose(spec, expected))
+
+    def test_spectrogram_batch_impulse(self):
+        waveform1 = np.zeros(40)
+        waveform1[9] = 1.0
+
+        waveform2 = np.zeros(28)
+        waveform2[12] = 3.0
+
+        waveform3 = np.zeros(51)
+        waveform3[26] = 4.5
+
+        waveform_list = [waveform1, waveform2, waveform3]
+
+        spec_list = spectrogram_batch(
+            waveform_list,
+            window_function(12, "hann", frame_length=16),
+            frame_length=16,
+            hop_length=4,
+            power=1.0,
+            center=True,
+            pad_mode="reflect",
+            onesided=True,
+        )
+
+        self.assertEqual(spec_list[0].shape, (9, 11))
+        self.assertEqual(spec_list[1].shape, (9, 8))
+        self.assertEqual(spec_list[2].shape, (9, 13))
+
+        expected1 = np.array([[0.0, 0.0669873, 0.9330127, 0.5, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]])
+        expected2 = np.array([[0.0, 0.0, 0.75, 3.0, 0.75, 0.0, 0.0, 0.0]])
+        expected3 = np.array([[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 3.375, 3.375, 0.0, 0.0, 0.0, 0.0, 0.0]])
+
+        self.assertTrue(np.allclose(spec_list[0], expected1))
+        self.assertTrue(np.allclose(spec_list[1], expected2))
+        self.assertTrue(np.allclose(spec_list[2], expected3))
+
+    def test_spectrogram_integration_test(self):
+        waveform = self._load_datasamples(1)[0]
+
+        spec = spectrogram(
+            waveform,
+            window_function(400, "hann", frame_length=512),
+            frame_length=512,
+            hop_length=128,
+            power=1.0,
+            center=True,
+            pad_mode="reflect",
+            onesided=True,
+        )
+        self.assertEqual(spec.shape, (257, 732))
+
+        # fmt: off
+        expected = np.array([
+            0.02464888, 0.04648664, 0.05872392, 0.02311783, 0.0327175 ,
+            0.02433643, 0.01198814, 0.02055709, 0.01559287, 0.01394357,
+            0.01299037, 0.01728045, 0.0254554 , 0.02486533, 0.02011792,
+            0.01755333, 0.02100457, 0.02337024, 0.01436963, 0.01464558,
+            0.0211017 , 0.0193489 , 0.01272165, 0.01858462, 0.03722598,
+            0.0456542 , 0.03281558, 0.00620586, 0.02226466, 0.03618042,
+            0.03508182, 0.02271432, 0.01051649, 0.01225771, 0.02315293,
+            0.02331886, 0.01417785, 0.0106844 , 0.01791214, 0.017177  ,
+            0.02125114, 0.05028201, 0.06830665, 0.05216664, 0.01963666,
+            0.06941418, 0.11513043, 0.12257859, 0.10948435, 0.08568069,
+            0.05509328, 0.05047818, 0.047112  , 0.05060737, 0.02982424,
+            0.02803827, 0.02933729, 0.01760491, 0.00587815, 0.02117637,
+            0.0293578 , 0.03452379, 0.02194803, 0.01676056,
+        ])
+        # fmt: on
+        self.assertTrue(np.allclose(spec[:64, 400], expected))
+
+        spec = spectrogram(
+            waveform,
+            window_function(400, "hann"),
+            frame_length=400,
+            hop_length=128,
+            fft_length=512,
+            power=1.0,
+            center=True,
+            pad_mode="reflect",
+            onesided=True,
+        )
+        self.assertEqual(spec.shape, (257, 732))
+        self.assertTrue(np.allclose(spec[:64, 400], expected))
+
+        mel_filters = mel_filter_bank(
+            num_frequency_bins=257,
+            num_mel_filters=400,
+            min_frequency=20,
+            max_frequency=8000,
+            sampling_rate=16000,
+            norm=None,
+            mel_scale="kaldi",
+            triangularize_in_mel_space=True,
+        )
+
+        spec = spectrogram(
+            waveform,
+            window_function(400, "povey", periodic=False),
+            frame_length=400,
+            hop_length=160,
+            fft_length=512,
+            power=2.0,
+            center=False,
+            pad_mode="reflect",
+            onesided=True,
+            preemphasis=0.97,
+            mel_filters=mel_filters,
+            log_mel="log",
+            mel_floor=1.1920928955078125e-07,
+            remove_dc_offset=True,
+        )
+        self.assertEqual(spec.shape, (400, 584))
+
+        # fmt: off
+        expected = np.array([-15.94238515,  -8.20712299,  -8.22704352, -15.94238515,
+       -15.94238515, -15.94238515, -15.94238515, -15.94238515,
+        -6.52463769,  -7.73677889, -15.94238515, -15.94238515,
+       -15.94238515, -15.94238515,  -4.18650018,  -3.37195286,
+       -15.94238515, -15.94238515, -15.94238515, -15.94238515,
+        -4.70190154,  -2.4217066 , -15.94238515, -15.94238515,
+       -15.94238515, -15.94238515,  -5.62755239,  -3.53385194,
+       -15.94238515, -15.94238515, -15.94238515, -15.94238515,
+        -9.43303023,  -8.77480925, -15.94238515, -15.94238515,
+       -15.94238515, -15.94238515,  -4.2951092 ,  -5.51585994,
+       -15.94238515, -15.94238515, -15.94238515,  -4.40151721,
+        -3.95228878, -15.94238515, -15.94238515, -15.94238515,
+        -6.10365415,  -4.59494697, -15.94238515, -15.94238515,
+       -15.94238515,  -8.10727767,  -6.2585298 , -15.94238515,
+       -15.94238515, -15.94238515,  -5.60161702,  -4.47217004,
+       -15.94238515, -15.94238515, -15.94238515,  -5.91641988]
+        )
+        # fmt: on
+        self.assertTrue(np.allclose(spec[:64, 400], expected, atol=1e-5))
+
+    def test_spectrogram_batch_integration_test(self):
+        waveform_list = self._load_datasamples(3)
+
+        spec_list = spectrogram_batch(
+            waveform_list,
+            window_function(400, "hann", frame_length=512),
+            frame_length=512,
+            hop_length=128,
+            power=1.0,
+            center=True,
+            pad_mode="reflect",
+            onesided=True,
+        )
+        self.assertEqual(spec_list[0].shape, (257, 732))
+        self.assertEqual(spec_list[1].shape, (257, 602))
+        self.assertEqual(spec_list[2].shape, (257, 1561))
+
+        # fmt: off
+        expected1 = np.array([
+            0.02464888, 0.04648664, 0.05872392, 0.02311783, 0.0327175 ,
+            0.02433643, 0.01198814, 0.02055709, 0.01559287, 0.01394357,
+            0.01299037, 0.01728045, 0.0254554 , 0.02486533, 0.02011792,
+            0.01755333, 0.02100457, 0.02337024, 0.01436963, 0.01464558,
+            0.0211017 , 0.0193489 , 0.01272165, 0.01858462, 0.03722598,
+            0.0456542 , 0.03281558, 0.00620586, 0.02226466, 0.03618042,
+            0.03508182, 0.02271432, 0.01051649, 0.01225771, 0.02315293,
+            0.02331886, 0.01417785, 0.0106844 , 0.01791214, 0.017177  ,
+            0.02125114, 0.05028201, 0.06830665, 0.05216664, 0.01963666,
+            0.06941418, 0.11513043, 0.12257859, 0.10948435, 0.08568069,
+            0.05509328, 0.05047818, 0.047112  , 0.05060737, 0.02982424,
+            0.02803827, 0.02933729, 0.01760491, 0.00587815, 0.02117637,
+            0.0293578 , 0.03452379, 0.02194803, 0.01676056,
+        ])
+        expected2 = np.array([
+            7.61983171e-02, 1.45338190e-01, 2.63903728e+00, 7.74429535e+00,
+            9.61932980e+00, 5.40767686e+00, 1.08924884e+00, 3.40908262e+00,
+            3.59484250e+00, 1.68451077e+00, 5.88405873e-01, 1.17042530e+00,
+            9.94803324e-01, 3.53757065e-01, 5.47699239e-01, 9.48368581e-01,
+            7.17770457e-01, 2.09396633e-01, 1.77574463e-01, 2.35644731e-01,
+            1.31535991e-01, 1.53539552e-02, 4.34416305e-02, 5.32897267e-02,
+            4.03567305e-02, 1.41842226e-02, 2.90514538e-02, 3.36549485e-02,
+            1.53516624e-02, 2.37464225e-02, 4.60092464e-02, 4.05769324e-02,
+            4.82633401e-03, 4.12675364e-02, 7.13859796e-02, 6.16866566e-02,
+            2.55657822e-02, 1.68923281e-02, 1.91299946e-02, 1.60033798e-02,
+            1.33405095e-02, 1.52065457e-02, 1.21833352e-02, 2.25786382e-03,
+            6.15358376e-03, 1.07647616e-02, 1.23051018e-02, 6.75289378e-03,
+            2.71127435e-03, 1.06515263e-02, 1.18463583e-02, 7.14347935e-03,
+            1.87912782e-03, 4.44236027e-03, 5.19630243e-03, 2.46666998e-03,
+            1.01598645e-03, 1.21589237e-03, 1.29095500e-03, 1.07447628e-03,
+            1.40218156e-03, 3.65402623e-03, 4.00592755e-03, 4.20001841e-03
+        ])
+        expected3 = np.array([
+            0.07805249, 0.34305022, 0.55617084, 1.22475182, 1.17040678,
+            0.51540532, 0.23570016, 0.06630775, 0.09017777, 0.07693192,
+            0.0333643 , 0.04873054, 0.04668559, 0.02384041, 0.02780435,
+            0.0289717 , 0.01704903, 0.0201644 , 0.01700376, 0.02176975,
+            0.02042491, 0.00732129, 0.00326042, 0.00245065, 0.00510645,
+            0.00681892, 0.00739329, 0.00551437, 0.0070674 , 0.00630015,
+            0.00379566, 0.0060098 , 0.00311543, 0.00902284, 0.01171038,
+            0.01202166, 0.01759194, 0.01652899, 0.01201872, 0.01295351,
+            0.00756432, 0.01415318, 0.02349972, 0.02296833, 0.02429341,
+            0.02447459, 0.01835044, 0.01437871, 0.02262246, 0.02972324,
+            0.03392252, 0.03037546, 0.01116927, 0.01555062, 0.02833379,
+            0.02294212, 0.02069847, 0.02496927, 0.02273526, 0.01341643,
+            0.00805407, 0.00624943, 0.01076262, 0.01876003
+        ])
+        # fmt: on
+        self.assertTrue(np.allclose(spec_list[0][:64, 400], expected1))
+        self.assertTrue(np.allclose(spec_list[1][:64, 400], expected2))
+        self.assertTrue(np.allclose(spec_list[2][:64, 400], expected3))
+
+        spec_list = spectrogram_batch(
+            waveform_list,
+            window_function(400, "hann"),
+            frame_length=400,
+            hop_length=128,
+            fft_length=512,
+            power=1.0,
+            center=True,
+            pad_mode="reflect",
+            onesided=True,
+        )
+        self.assertEqual(spec_list[0].shape, (257, 732))
+        self.assertEqual(spec_list[1].shape, (257, 602))
+        self.assertEqual(spec_list[2].shape, (257, 1561))
+        self.assertTrue(np.allclose(spec_list[0][:64, 400], expected1))
+        self.assertTrue(np.allclose(spec_list[1][:64, 400], expected2))
+        self.assertTrue(np.allclose(spec_list[2][:64, 400], expected3))
+
+        mel_filters = mel_filter_bank(
+            num_frequency_bins=257,
+            num_mel_filters=400,
+            min_frequency=20,
+            max_frequency=8000,
+            sampling_rate=16000,
+            norm=None,
+            mel_scale="kaldi",
+            triangularize_in_mel_space=True,
+        )
+
+        spec_list = spectrogram_batch(
+            waveform_list,
+            window_function(400, "povey", periodic=False),
+            frame_length=400,
+            hop_length=160,
+            fft_length=512,
+            power=2.0,
+            center=False,
+            pad_mode="reflect",
+            onesided=True,
+            preemphasis=0.97,
+            mel_filters=mel_filters,
+            log_mel="log",
+            mel_floor=1.1920928955078125e-07,
+            remove_dc_offset=True,
+        )
+        self.assertEqual(spec_list[0].shape, (400, 584))
+        self.assertEqual(spec_list[1].shape, (400, 480))
+        self.assertEqual(spec_list[2].shape, (400, 1247))
+
+        # fmt: off
+        expected1 = np.array([-15.94238515,  -8.20712299,  -8.22704352, -15.94238515,
+       -15.94238515, -15.94238515, -15.94238515, -15.94238515,
+        -6.52463769,  -7.73677889, -15.94238515, -15.94238515,
+       -15.94238515, -15.94238515,  -4.18650018,  -3.37195286,
+       -15.94238515, -15.94238515, -15.94238515, -15.94238515,
+        -4.70190154,  -2.4217066 , -15.94238515, -15.94238515,
+       -15.94238515, -15.94238515,  -5.62755239,  -3.53385194,
+       -15.94238515, -15.94238515, -15.94238515, -15.94238515,
+        -9.43303023,  -8.77480925, -15.94238515, -15.94238515,
+       -15.94238515, -15.94238515,  -4.2951092 ,  -5.51585994,
+       -15.94238515, -15.94238515, -15.94238515,  -4.40151721,
+        -3.95228878, -15.94238515, -15.94238515, -15.94238515,
+        -6.10365415,  -4.59494697, -15.94238515, -15.94238515,
+       -15.94238515,  -8.10727767,  -6.2585298 , -15.94238515,
+       -15.94238515, -15.94238515,  -5.60161702,  -4.47217004,
+       -15.94238515, -15.94238515, -15.94238515,  -5.91641988]
+        )
+        expected2 = np.array([-15.942385, -8.531508, -8.551396, -15.942385, -15.942385,
+        -15.942385, -15.942385, -15.942385, -5.626043, -6.8381968,
+        -15.942385, -15.942385, -15.942385, -15.942385, -3.3122184,
+        -2.49764, -15.942385, -15.942385, -15.942385, -15.942385,
+        -3.625868, -1.3457257, -15.942385, -15.942385, -15.942385,
+        -15.942385, -4.2223063, -2.1285915, -15.942385, -15.942385,
+        -15.942385, -15.942385, -8.611152, -7.952894, -15.942385,
+        -15.942385, -15.942385, -15.942385, -2.7585578, -3.9793255,
+        -15.942385, -15.942385, -15.942385, -2.5377562, -2.0885658,
+        -15.942385, -15.942385, -15.942385, -3.8310733, -2.322393,
+        -15.942385, -15.942385, -15.942385, -7.674944, -5.8261633,
+        -15.942385, -15.942385, -15.942385, -3.5960004, -2.4665844,
+        -15.942385, -15.942385, -15.942385, -1.7905309]
+        )
+        expected3 = np.array([-15.942385, -13.406995, -13.426883, -15.942385, -15.942385,
+        -15.942385, -15.942385, -15.942385, -15.942385, -15.942385,
+        -15.942385, -15.942385, -15.942385, -15.942385, -13.493383,
+        -12.678805, -15.942385, -15.942385, -15.942385, -15.942385,
+        -14.809377, -12.529235, -15.942385, -15.942385, -15.942385,
+        -15.942385, -13.838827, -11.745112, -15.942385, -15.942385,
+        -15.942385, -15.942385, -13.9336405, -13.275384, -15.942385,
+        -15.942385, -15.942385, -15.942385, -13.043786, -14.264554,
+        -15.942385, -15.942385, -15.942385, -13.060181, -12.610991,
+        -15.942385, -15.942385, -15.942385, -14.152064, -12.643384,
+        -15.942385, -15.942385, -15.942385, -14.48317, -12.634389,
+        -15.942385, -15.942385, -15.942385, -14.627316, -13.4979,
+        -15.942385, -15.942385, -15.942385, -12.6279955]
+        )
+        # fmt: on
+        self.assertTrue(np.allclose(spec_list[0][:64, 400], expected1, atol=1e-5))
+        self.assertTrue(np.allclose(spec_list[1][:64, 400], expected2, atol=1e-5))
+        self.assertTrue(np.allclose(spec_list[2][:64, 400], expected3, atol=1e-5))
+
+    def test_spectrogram_center_padding(self):
+        waveform = self._load_datasamples(1)[0]
+
+        spec = spectrogram(
+            waveform,
+            window_function(512, "hann"),
+            frame_length=512,
+            hop_length=128,
+            center=True,
+            pad_mode="reflect",
+        )
+        self.assertEqual(spec.shape, (257, 732))
+
+        # fmt: off
+        expected = np.array([
+            0.1287945 , 0.12792738, 0.08311573, 0.03155122, 0.02470202,
+            0.00727857, 0.00910694, 0.00686163, 0.01238981, 0.01473668,
+            0.00336144, 0.00370314, 0.00600871, 0.01120164, 0.01942998,
+            0.03132008, 0.0232842 , 0.01124642, 0.02754783, 0.02423725,
+            0.00147893, 0.00038027, 0.00112299, 0.00596233, 0.00571529,
+            0.02084235, 0.0231855 , 0.00810006, 0.01837943, 0.00651339,
+            0.00093931, 0.00067426, 0.01058399, 0.01270507, 0.00151734,
+            0.00331913, 0.00302416, 0.01081792, 0.00754549, 0.00148963,
+            0.00111943, 0.00152573, 0.00608017, 0.01749986, 0.01205949,
+            0.0143082 , 0.01910573, 0.00413786, 0.03916619, 0.09873404,
+            0.08302026, 0.02673891, 0.00401255, 0.01397392, 0.00751862,
+            0.01024884, 0.01544606, 0.00638907, 0.00623633, 0.0085103 ,
+            0.00217659, 0.00276204, 0.00260835, 0.00299299,
+        ])
+        # fmt: on
+        self.assertTrue(np.allclose(spec[:64, 0], expected))
+
+        spec = spectrogram(
+            waveform,
+            window_function(512, "hann"),
+            frame_length=512,
+            hop_length=128,
+            center=True,
+            pad_mode="constant",
+        )
+        self.assertEqual(spec.shape, (257, 732))
+
+        # fmt: off
+        expected = np.array([
+            0.06558744, 0.06889656, 0.06263352, 0.04264418, 0.03404115,
+            0.03244197, 0.02279134, 0.01646339, 0.01452216, 0.00826055,
+            0.00062093, 0.0031821 , 0.00419456, 0.00689327, 0.01106367,
+            0.01712119, 0.01721762, 0.00977533, 0.01606626, 0.02275621,
+            0.01727687, 0.00992739, 0.01217688, 0.01049927, 0.01022947,
+            0.01302475, 0.01166873, 0.01081812, 0.01057327, 0.00767912,
+            0.00429567, 0.00089625, 0.00654583, 0.00912084, 0.00700984,
+            0.00225026, 0.00290545, 0.00667712, 0.00730663, 0.00410813,
+            0.00073102, 0.00219296, 0.00527618, 0.00996585, 0.01123781,
+            0.00872816, 0.01165121, 0.02047945, 0.03681747, 0.0514379 ,
+            0.05137928, 0.03960042, 0.02821562, 0.01813349, 0.01201322,
+            0.01260964, 0.00900654, 0.00207905, 0.00456714, 0.00850599,
+            0.00788239, 0.00664407, 0.00824227, 0.00628301,
+        ])
+        # fmt: on
+        self.assertTrue(np.allclose(spec[:64, 0], expected))
+
+        spec = spectrogram(
+            waveform,
+            window_function(512, "hann"),
+            frame_length=512,
+            hop_length=128,
+            center=False,
+        )
+        self.assertEqual(spec.shape, (257, 728))
+
+        # fmt: off
+        expected = np.array([
+            0.00250445, 0.02161521, 0.06232229, 0.04339567, 0.00937727,
+            0.01080616, 0.00248685, 0.0095264 , 0.00727476, 0.0079152 ,
+            0.00839946, 0.00254932, 0.00716622, 0.005559  , 0.00272623,
+            0.00581774, 0.01896395, 0.01829788, 0.01020514, 0.01632692,
+            0.00870888, 0.02065827, 0.0136022 , 0.0132382 , 0.011827  ,
+            0.00194505, 0.0189979 , 0.026874  , 0.02194014, 0.01923883,
+            0.01621437, 0.00661967, 0.00289517, 0.00470257, 0.00957801,
+            0.00191455, 0.00431664, 0.00544359, 0.01126213, 0.00785778,
+            0.00423469, 0.01322504, 0.02226548, 0.02318576, 0.03428908,
+            0.03648811, 0.0202938 , 0.011902  , 0.03226198, 0.06347476,
+            0.01306318, 0.05308729, 0.05474771, 0.03127991, 0.00998512,
+            0.01449977, 0.01272741, 0.00868176, 0.00850386, 0.00313876,
+            0.00811857, 0.00538216, 0.00685749, 0.00535275,
+        ])
+        # fmt: on
+        self.assertTrue(np.allclose(spec[:64, 0], expected))
+
+    def test_spectrogram_batch_center_padding(self):
+        waveform_list = self._load_datasamples(3)
+
+        spec_list = spectrogram_batch(
+            waveform_list,
+            window_function(512, "hann"),
+            frame_length=512,
+            hop_length=128,
+            center=True,
+            pad_mode="reflect",
+        )
+        self.assertEqual(spec_list[0].shape, (257, 732))
+        self.assertEqual(spec_list[1].shape, (257, 602))
+        self.assertEqual(spec_list[2].shape, (257, 1561))
+
+        # fmt: off
+        expected1 = np.array([
+            0.1287945 , 0.12792738, 0.08311573, 0.03155122, 0.02470202,
+            0.00727857, 0.00910694, 0.00686163, 0.01238981, 0.01473668,
+            0.00336144, 0.00370314, 0.00600871, 0.01120164, 0.01942998,
+            0.03132008, 0.0232842 , 0.01124642, 0.02754783, 0.02423725,
+            0.00147893, 0.00038027, 0.00112299, 0.00596233, 0.00571529,
+            0.02084235, 0.0231855 , 0.00810006, 0.01837943, 0.00651339,
+            0.00093931, 0.00067426, 0.01058399, 0.01270507, 0.00151734,
+            0.00331913, 0.00302416, 0.01081792, 0.00754549, 0.00148963,
+            0.00111943, 0.00152573, 0.00608017, 0.01749986, 0.01205949,
+            0.0143082 , 0.01910573, 0.00413786, 0.03916619, 0.09873404,
+            0.08302026, 0.02673891, 0.00401255, 0.01397392, 0.00751862,
+            0.01024884, 0.01544606, 0.00638907, 0.00623633, 0.0085103 ,
+            0.00217659, 0.00276204, 0.00260835, 0.00299299,
+        ])
+        expected2 = np.array([
+            1.89624839e-02, 1.23274978e-02, 3.69160250e-02, 4.76267971e-02,
+            1.39258439e-02, 2.98370440e-02, 2.74845166e-03, 3.01934010e-03,
+            1.18722776e-02, 9.70834121e-03, 2.06300567e-04, 6.32975250e-04,
+            8.20603687e-03, 1.21864351e-02, 3.28791840e-03, 3.36801982e-04,
+            2.79373326e-03, 5.00530424e-03, 8.46884679e-03, 1.14089288e-02,
+            8.59052036e-03, 2.88538425e-03, 9.95071139e-03, 6.80431770e-03,
+            2.95809377e-03, 1.46285209e-04, 3.36268265e-03, 4.80051298e-04,
+            2.84506916e-03, 9.34222655e-04, 3.42161348e-03, 2.79612141e-03,
+            3.38875921e-03, 2.85030343e-03, 5.39513239e-05, 2.72908504e-03,
+            2.09591188e-03, 5.00271388e-04, 8.31917219e-04, 2.37967237e-03,
+            1.75001193e-03, 1.31826295e-04, 8.83622793e-04, 1.54303256e-04,
+            3.09544569e-03, 4.08527814e-03, 2.73566321e-03, 1.78805250e-03,
+            9.53314066e-06, 1.74316950e-03, 1.51099428e-03, 8.65990878e-04,
+            8.44859460e-04, 5.35220199e-04, 5.36562002e-04, 8.33181897e-04,
+            8.22705682e-04, 1.81083288e-03, 9.75003233e-04, 6.73114730e-04,
+            6.81665202e-04, 2.05180887e-03, 1.10151991e-03, 4.75923851e-04,
+        ])
+        expected3 = np.array([
+            0.07079848, 0.04237922, 0.0220724, 0.04446052, 0.03598337,
+            0.03327273, 0.02545774, 0.01319528, 0.00919659, 0.01376867,
+            0.00361992, 0.00608425, 0.01105873, 0.0105565, 0.00744286,
+            0.00244849, 0.00257317, 0.00749989, 0.01061386, 0.01525312,
+            0.00656914, 0.01199581, 0.00487319, 0.00830956, 0.0046706,
+            0.00588962, 0.00544486, 0.00565179, 0.00050112, 0.01108059,
+            0.00217417, 0.00453234, 0.00537306, 0.00269329, 0.00342333,
+            0.00095484, 0.00708934, 0.00660373, 0.00543686, 0.00217186,
+            0.00431519, 0.00457764, 0.00503529, 0.01166454, 0.01375581,
+            0.01467224, 0.00873404, 0.00534086, 0.00476848, 0.0226163,
+            0.0314, 0.00151021, 0.01975221, 0.01637519, 0.00046068,
+            0.0460544, 0.06285986, 0.03151625, 0.0013598, 0.004804,
+            0.0073824, 0.02312599, 0.02613977, 0.01056851
+        ])
+        # fmt: on
+        self.assertTrue(np.allclose(spec_list[0][:64, 0], expected1))
+        self.assertTrue(np.allclose(spec_list[1][:64, 0], expected2))
+        self.assertTrue(np.allclose(spec_list[2][:64, 0], expected3))
+
+        spec_list = spectrogram_batch(
+            waveform_list,
+            window_function(512, "hann"),
+            frame_length=512,
+            hop_length=128,
+            center=True,
+            pad_mode="constant",
+        )
+        self.assertEqual(spec_list[0].shape, (257, 732))
+        self.assertEqual(spec_list[1].shape, (257, 602))
+        self.assertEqual(spec_list[2].shape, (257, 1561))
+
+        # fmt: off
+        expected1 = np.array([
+            0.06558744, 0.06889656, 0.06263352, 0.04264418, 0.03404115,
+            0.03244197, 0.02279134, 0.01646339, 0.01452216, 0.00826055,
+            0.00062093, 0.0031821 , 0.00419456, 0.00689327, 0.01106367,
+            0.01712119, 0.01721762, 0.00977533, 0.01606626, 0.02275621,
+            0.01727687, 0.00992739, 0.01217688, 0.01049927, 0.01022947,
+            0.01302475, 0.01166873, 0.01081812, 0.01057327, 0.00767912,
+            0.00429567, 0.00089625, 0.00654583, 0.00912084, 0.00700984,
+            0.00225026, 0.00290545, 0.00667712, 0.00730663, 0.00410813,
+            0.00073102, 0.00219296, 0.00527618, 0.00996585, 0.01123781,
+            0.00872816, 0.01165121, 0.02047945, 0.03681747, 0.0514379 ,
+            0.05137928, 0.03960042, 0.02821562, 0.01813349, 0.01201322,
+            0.01260964, 0.00900654, 0.00207905, 0.00456714, 0.00850599,
+            0.00788239, 0.00664407, 0.00824227, 0.00628301,
+        ])
+        expected2 = np.array([
+            0.00955754, 0.01445548, 0.02393902, 0.02903068, 0.02512844,
+            0.01508297, 0.00474784, 0.00440362, 0.0073898, 0.00546519,
+            0.00126077, 0.00240507, 0.00523254, 0.00632742, 0.00415215,
+            0.00056628, 0.00161288, 0.0026956, 0.00431587, 0.00621471,
+            0.00791291, 0.0079454, 0.00594525, 0.00334581, 0.00180047,
+            0.00144485, 0.00175764, 0.00188037, 0.00134889, 0.00150253,
+            0.00178821, 0.00158875, 0.00204339, 0.00266497, 0.00280556,
+            0.00221949, 0.00108956, 0.000532, 0.00108454, 0.00129254,
+            0.00089315, 0.00022803, 0.00038176, 0.0011302, 0.00189306,
+            0.0021964, 0.00203576, 0.00207306, 0.00217727, 0.00174297,
+            0.00103331, 0.00076695, 0.0007422, 0.00061986, 0.00081204,
+            0.00079615, 0.00089417, 0.00105452, 0.00042615, 0.00066372,
+            0.00132765, 0.00122087, 0.00054903, 0.00107945,
+        ])
+        expected3 = np.array([
+            0.03573493, 0.03625983, 0.03341755, 0.02431477, 0.01770546,
+            0.0169356 , 0.01579034, 0.01600499, 0.01329064, 0.00747957,
+            0.00367372, 0.00403853, 0.00519597, 0.00551022, 0.00532757,
+            0.00367569, 0.00130341, 0.00345149, 0.00520744, 0.00872308,
+            0.01172503, 0.00948154, 0.00344236, 0.00387997, 0.00425455,
+            0.00394357, 0.00711733, 0.00615654, 0.00055756, 0.00656414,
+            0.00852001, 0.00666252, 0.00509767, 0.00246784, 0.00376049,
+            0.00682879, 0.00641118, 0.00469685, 0.00358701, 0.0015552 ,
+            0.00261458, 0.00701979, 0.00929578, 0.00894536, 0.00828491,
+            0.00773528, 0.00552091, 0.00259871, 0.00933179, 0.01588626,
+            0.01697887, 0.01268552, 0.00957255, 0.01204092, 0.02123362,
+            0.03062669, 0.03215763, 0.02629963, 0.01769568, 0.01088869,
+            0.01151334, 0.01378197, 0.01319263, 0.01066859,
+       ])
+        # fmt: on
+        self.assertTrue(np.allclose(spec_list[0][:64, 0], expected1))
+        self.assertTrue(np.allclose(spec_list[1][:64, 0], expected2))
+        self.assertTrue(np.allclose(spec_list[2][:64, 0], expected3))
+
+        spec_list = spectrogram_batch(
+            waveform_list,
+            window_function(512, "hann"),
+            frame_length=512,
+            hop_length=128,
+            center=False,
+        )
+        self.assertEqual(spec_list[0].shape, (257, 728))
+        self.assertEqual(spec_list[1].shape, (257, 598))
+        self.assertEqual(spec_list[2].shape, (257, 1557))
+
+        # fmt: off
+        expected1 = np.array([
+            0.00250445, 0.02161521, 0.06232229, 0.04339567, 0.00937727,
+            0.01080616, 0.00248685, 0.0095264 , 0.00727476, 0.0079152 ,
+            0.00839946, 0.00254932, 0.00716622, 0.005559  , 0.00272623,
+            0.00581774, 0.01896395, 0.01829788, 0.01020514, 0.01632692,
+            0.00870888, 0.02065827, 0.0136022 , 0.0132382 , 0.011827  ,
+            0.00194505, 0.0189979 , 0.026874  , 0.02194014, 0.01923883,
+            0.01621437, 0.00661967, 0.00289517, 0.00470257, 0.00957801,
+            0.00191455, 0.00431664, 0.00544359, 0.01126213, 0.00785778,
+            0.00423469, 0.01322504, 0.02226548, 0.02318576, 0.03428908,
+            0.03648811, 0.0202938 , 0.011902  , 0.03226198, 0.06347476,
+            0.01306318, 0.05308729, 0.05474771, 0.03127991, 0.00998512,
+            0.01449977, 0.01272741, 0.00868176, 0.00850386, 0.00313876,
+            0.00811857, 0.00538216, 0.00685749, 0.00535275,
+        ])
+        expected2 = np.array([
+            0.01232908, 0.05980514, 0.08285419, 0.01850723, 0.02823627,
+            0.00204369, 0.01372626, 0.00956435, 0.02267217, 0.00947112,
+            0.00355174, 0.00418008, 0.00843608, 0.01559252, 0.01125505,
+            0.00183573, 0.00765051, 0.0109983 , 0.00890545, 0.00583453,
+            0.00115901, 0.00579039, 0.00151353, 0.00395812, 0.00231413,
+            0.00384272, 0.00313914, 0.00072331, 0.00338935, 0.00383328,
+            0.00218129, 0.00284516, 0.00228538, 0.00083603, 0.00111663,
+            0.00235799, 0.00142748, 0.00092908, 0.0012966 , 0.0011403 ,
+            0.0010619 , 0.00158732, 0.00289866, 0.00216709, 0.00313325,
+            0.00361277, 0.00202507, 0.0009948 , 0.00114428, 0.00200851,
+            0.0009234 , 0.00063468, 0.00018746, 0.00100463, 0.00053799,
+            0.00080009, 0.00158291, 0.00172077, 0.00173586, 0.00197127,
+            0.00107058, 0.00043486, 0.0009859 , 0.00215484,
+        ])
+        expected3 = np.array([
+            0.01864123, 0.06131337, 0.08346292, 0.04936386, 0.02792609,
+            0.01005205, 0.00884826, 0.02198604, 0.02421535, 0.00957573,
+            0.00503561, 0.00241331, 0.00175652, 0.00195889, 0.00453299,
+            0.0020317 , 0.00249264, 0.00517483, 0.01111943, 0.0150079 ,
+            0.01977743, 0.01253825, 0.00517561, 0.01031712, 0.00579466,
+            0.00783679, 0.0071415 , 0.00591847, 0.01510728, 0.01194921,
+            0.00518072, 0.00125978, 0.00577552, 0.01050614, 0.0077644 ,
+            0.0042905 , 0.00278469, 0.00166695, 0.00255013, 0.00578153,
+            0.00586451, 0.00929514, 0.01501226, 0.00741419, 0.00310625,
+            0.00086757, 0.00595618, 0.0053882 , 0.0116266 , 0.02504773,
+            0.02889692, 0.03739442, 0.04730207, 0.03856638, 0.05700104,
+            0.04299267, 0.02153366, 0.03740607, 0.03811468, 0.01575022,
+            0.00676344, 0.01359865, 0.01769319, 0.00907966,
+        ])
+        # fmt: on
+        self.assertTrue(np.allclose(spec_list[0][:64, 0], expected1))
+        self.assertTrue(np.allclose(spec_list[1][:64, 0], expected2))
+        self.assertTrue(np.allclose(spec_list[2][:64, 0], expected3))
+
+    def test_spectrogram_shapes(self):
+        waveform = self._load_datasamples(1)[0]
+
+        spec = spectrogram(
+            waveform,
+            window_function(400, "hann"),
+            frame_length=400,
+            hop_length=128,
+            power=1.0,
+            center=True,
+            pad_mode="reflect",
+            onesided=True,
+        )
+        self.assertEqual(spec.shape, (201, 732))
+
+        spec = spectrogram(
+            waveform,
+            window_function(400, "hann"),
+            frame_length=400,
+            hop_length=128,
+            power=1.0,
+            center=False,
+            pad_mode="reflect",
+            onesided=True,
+        )
+        self.assertEqual(spec.shape, (201, 729))
+
+        spec = spectrogram(
+            waveform,
+            window_function(400, "hann"),
+            frame_length=400,
+            hop_length=128,
+            fft_length=512,
+            power=1.0,
+            center=True,
+            pad_mode="reflect",
+            onesided=True,
+        )
+        self.assertEqual(spec.shape, (257, 732))
+
+        spec = spectrogram(
+            waveform,
+            window_function(400, "hann", frame_length=512),
+            frame_length=512,
+            hop_length=64,
+            power=1.0,
+            center=True,
+            pad_mode="reflect",
+            onesided=False,
+        )
+        self.assertEqual(spec.shape, (512, 1464))
+
+        spec = spectrogram(
+            waveform,
+            window_function(512, "hann"),
+            frame_length=512,
+            hop_length=64,
+            power=1.0,
+            center=True,
+            pad_mode="reflect",
+            onesided=False,
+        )
+        self.assertEqual(spec.shape, (512, 1464))
+
+        spec = spectrogram(
+            waveform,
+            window_function(512, "hann"),
+            frame_length=512,
+            hop_length=512,
+            power=1.0,
+            center=True,
+            pad_mode="reflect",
+            onesided=False,
+        )
+        self.assertEqual(spec.shape, (512, 183))
+
+    def test_spectrogram_batch_shapes(self):
+        waveform_list = self._load_datasamples(3)
+
+        spec_list = spectrogram_batch(
+            waveform_list,
+            window_function(400, "hann"),
+            frame_length=400,
+            hop_length=128,
+            power=1.0,
+            center=True,
+            pad_mode="reflect",
+            onesided=True,
+        )
+        self.assertEqual(spec_list[0].shape, (201, 732))
+        self.assertEqual(spec_list[1].shape, (201, 602))
+        self.assertEqual(spec_list[2].shape, (201, 1561))
+
+        spec_list = spectrogram_batch(
+            waveform_list,
+            window_function(400, "hann"),
+            frame_length=400,
+            hop_length=128,
+            power=1.0,
+            center=False,
+            pad_mode="reflect",
+            onesided=True,
+        )
+        self.assertEqual(spec_list[0].shape, (201, 729))
+        self.assertEqual(spec_list[1].shape, (201, 599))
+        self.assertEqual(spec_list[2].shape, (201, 1558))
+
+        spec_list = spectrogram_batch(
+            waveform_list,
+            window_function(400, "hann"),
+            frame_length=400,
+            hop_length=128,
+            fft_length=512,
+            power=1.0,
+            center=True,
+            pad_mode="reflect",
+            onesided=True,
+        )
+        self.assertEqual(spec_list[0].shape, (257, 732))
+        self.assertEqual(spec_list[1].shape, (257, 602))
+        self.assertEqual(spec_list[2].shape, (257, 1561))
+
+        spec_list = spectrogram_batch(
+            waveform_list,
+            window_function(400, "hann", frame_length=512),
+            frame_length=512,
+            hop_length=64,
+            power=1.0,
+            center=True,
+            pad_mode="reflect",
+            onesided=False,
+        )
+        self.assertEqual(spec_list[0].shape, (512, 1464))
+        self.assertEqual(spec_list[1].shape, (512, 1204))
+        self.assertEqual(spec_list[2].shape, (512, 3122))
+
+        spec_list = spectrogram_batch(
+            waveform_list,
+            window_function(512, "hann"),
+            frame_length=512,
+            hop_length=64,
+            power=1.0,
+            center=True,
+            pad_mode="reflect",
+            onesided=False,
+        )
+        self.assertEqual(spec_list[0].shape, (512, 1464))
+        self.assertEqual(spec_list[1].shape, (512, 1204))
+        self.assertEqual(spec_list[2].shape, (512, 3122))
+
+        spec_list = spectrogram_batch(
+            waveform_list,
+            window_function(512, "hann"),
+            frame_length=512,
+            hop_length=512,
+            power=1.0,
+            center=True,
+            pad_mode="reflect",
+            onesided=False,
+        )
+        self.assertEqual(spec_list[0].shape, (512, 183))
+        self.assertEqual(spec_list[1].shape, (512, 151))
+        self.assertEqual(spec_list[2].shape, (512, 391))
+
+    def test_mel_spectrogram(self):
+        waveform = self._load_datasamples(1)[0]
+
+        mel_filters = mel_filter_bank(
+            num_frequency_bins=513,
+            num_mel_filters=13,
+            min_frequency=100,
+            max_frequency=4000,
+            sampling_rate=16000,
+            norm=None,
+            mel_scale="htk",
+        )
+        self.assertEqual(mel_filters.shape, (513, 13))
+
+        spec = spectrogram(
+            waveform,
+            window_function(800, "hann", frame_length=1024),
+            frame_length=1024,
+            hop_length=128,
+            power=2.0,
+        )
+        self.assertEqual(spec.shape, (513, 732))
+
+        spec = spectrogram(
+            waveform,
+            window_function(800, "hann", frame_length=1024),
+            frame_length=1024,
+            hop_length=128,
+            power=2.0,
+            mel_filters=mel_filters,
+        )
+        self.assertEqual(spec.shape, (13, 732))
+
+        # fmt: off
+        expected = np.array([
+            1.08027889e+02, 1.48080673e+01, 7.70758213e+00, 9.57676639e-01,
+            8.81639061e-02, 5.26073833e-02, 1.52736155e-02, 9.95350117e-03,
+            7.95364356e-03, 1.01148004e-02, 4.29241020e-03, 9.90708797e-03,
+            9.44153646e-04
+        ])
+        # fmt: on
+        self.assertTrue(np.allclose(spec[:, 300], expected))
+
+    def test_mel_spectrogram_batch(self):
+        waveform_list = self._load_datasamples(3)
+
+        mel_filters = mel_filter_bank(
+            num_frequency_bins=513,
+            num_mel_filters=13,
+            min_frequency=100,
+            max_frequency=4000,
+            sampling_rate=16000,
+            norm=None,
+            mel_scale="htk",
+        )
+        self.assertEqual(mel_filters.shape, (513, 13))
+
+        spec_list = spectrogram_batch(
+            waveform_list,
+            window_function(800, "hann", frame_length=1024),
+            frame_length=1024,
+            hop_length=128,
+            power=2.0,
+        )
+        self.assertEqual(spec_list[0].shape, (513, 732))
+        self.assertEqual(spec_list[1].shape, (513, 602))
+        self.assertEqual(spec_list[2].shape, (513, 1561))
+
+        spec_list = spectrogram_batch(
+            waveform_list,
+            window_function(800, "hann", frame_length=1024),
+            frame_length=1024,
+            hop_length=128,
+            power=2.0,
+            mel_filters=mel_filters,
+        )
+        self.assertEqual(spec_list[0].shape, (13, 732))
+        self.assertEqual(spec_list[1].shape, (13, 602))
+        self.assertEqual(spec_list[2].shape, (13, 1561))
+
+        # fmt: off
+        expected1 = np.array([
+            1.08027889e+02, 1.48080673e+01, 7.70758213e+00, 9.57676639e-01,
+            8.81639061e-02, 5.26073833e-02, 1.52736155e-02, 9.95350117e-03,
+            7.95364356e-03, 1.01148004e-02, 4.29241020e-03, 9.90708797e-03,
+            9.44153646e-04
+        ])
+        expected2 = np.array([
+            71.82577165, 109.44693334, 272.4834194, 164.90450355,
+            16.54056349, 11.60810547, 24.87525946, 21.07317022,
+            1.26736284, 1.4583074, 1.36659061, 1.76305768,
+            2.03703503
+        ])
+        expected3 = np.array([
+            5.22246749e+02, 6.92660728e+02, 2.65895922e+02, 2.06526565e+01,
+            2.28692104e+00, 1.19473622e+00, 8.43228216e-01, 3.20760592e+00,
+            1.33654151e+00, 1.51050684e-01, 2.78282477e-01, 9.25020981e-01,
+            2.29908841e-01
+        ])
+        # fmt: on
+        self.assertTrue(np.allclose(spec_list[0][:, 300], expected1))
+        self.assertTrue(np.allclose(spec_list[1][:, 300], expected2))
+        self.assertTrue(np.allclose(spec_list[2][:, 300], expected3))
+
+    def test_spectrogram_power(self):
+        waveform = self._load_datasamples(1)[0]
+
+        spec = spectrogram(
+            waveform,
+            window_function(400, "hann", frame_length=512),
+            frame_length=512,
+            hop_length=128,
+            power=None,
+        )
+        self.assertEqual(spec.shape, (257, 732))
+        self.assertEqual(spec.dtype, np.complex64)
+
+        # fmt: off
+        expected = np.array([
+             0.01452305+0.01820039j, -0.01737362-0.01641946j,
+             0.0121028 +0.01565081j, -0.02794554-0.03021514j,
+             0.04719803+0.04086519j, -0.04391563-0.02779365j,
+             0.05682834+0.01571325j, -0.08604821-0.02023657j,
+             0.07497991+0.0186641j , -0.06366091-0.00922475j,
+             0.11003416+0.0114788j , -0.13677941-0.01523552j,
+             0.10934535-0.00117226j, -0.11635598+0.02551187j,
+             0.14708674-0.03469823j, -0.1328196 +0.06034218j,
+             0.12667368-0.13973421j, -0.14764774+0.18912019j,
+             0.10235471-0.12181523j, -0.00773012+0.04730498j,
+            -0.01487191-0.07312611j, -0.02739162+0.09619419j,
+             0.02895459-0.05398273j,  0.01198589+0.05276592j,
+            -0.02117299-0.10123465j,  0.00666388+0.09526499j,
+            -0.01672773-0.05649684j,  0.02723125+0.05939891j,
+            -0.01879361-0.062954j  ,  0.03686557+0.04568823j,
+            -0.07394181-0.07949649j,  0.06238583+0.13905765j,
+        ])
+        # fmt: on
+        self.assertTrue(np.allclose(spec[64:96, 321], expected))
+
+        spec = spectrogram(
+            waveform,
+            window_function(400, "hann", frame_length=512),
+            frame_length=512,
+            hop_length=128,
+            power=1.0,
+        )
+        self.assertEqual(spec.shape, (257, 732))
+        self.assertEqual(spec.dtype, np.float64)
+
+        # fmt: off
+        expected = np.array([
+            0.02328461, 0.02390484, 0.01978448, 0.04115711, 0.0624309 ,
+            0.05197181, 0.05896072, 0.08839577, 0.07726794, 0.06432579,
+            0.11063128, 0.13762532, 0.10935163, 0.11911998, 0.15112405,
+            0.14588428, 0.18860507, 0.23992978, 0.15910825, 0.04793241,
+            0.07462307, 0.10001811, 0.06125769, 0.05411011, 0.10342509,
+            0.09549777, 0.05892122, 0.06534349, 0.06569936, 0.05870678,
+            0.10856833, 0.1524107 , 0.11463385, 0.05766969, 0.12385171,
+            0.14472842, 0.11978184, 0.10353675, 0.07244056, 0.03461861,
+            0.02624896, 0.02227475, 0.01238363, 0.00885281, 0.0110049 ,
+            0.00807005, 0.01033663, 0.01703181, 0.01445856, 0.00585615,
+            0.0132431 , 0.02754132, 0.01524478, 0.0204908 , 0.07453328,
+            0.10716327, 0.07195779, 0.08816078, 0.18340898, 0.16449876,
+            0.12322842, 0.1621659 , 0.12334293, 0.06033659,
+        ])
+        # fmt: on
+        self.assertTrue(np.allclose(spec[64:128, 321], expected))
+
+        spec = spectrogram(
+            waveform,
+            window_function(400, "hann", frame_length=512),
+            frame_length=512,
+            hop_length=128,
+            power=2.0,
+        )
+        self.assertEqual(spec.shape, (257, 732))
+        self.assertEqual(spec.dtype, np.float64)
+
+        # fmt: off
+        expected = np.array([
+            5.42173162e-04, 5.71441371e-04, 3.91425507e-04, 1.69390778e-03,
+            3.89761780e-03, 2.70106923e-03, 3.47636663e-03, 7.81381316e-03,
+            5.97033510e-03, 4.13780799e-03, 1.22392802e-02, 1.89407300e-02,
+            1.19577805e-02, 1.41895693e-02, 2.28384770e-02, 2.12822221e-02,
+            3.55718732e-02, 5.75663000e-02, 2.53154356e-02, 2.29751552e-03,
+            5.56860259e-03, 1.00036217e-02, 3.75250424e-03, 2.92790355e-03,
+            1.06967501e-02, 9.11982451e-03, 3.47171025e-03, 4.26977174e-03,
+            4.31640586e-03, 3.44648538e-03, 1.17870830e-02, 2.32290216e-02,
+            1.31409196e-02, 3.32579296e-03, 1.53392460e-02, 2.09463164e-02,
+            1.43476883e-02, 1.07198600e-02, 5.24763530e-03, 1.19844836e-03,
+            6.89007982e-04, 4.96164430e-04, 1.53354369e-04, 7.83722571e-05,
+            1.21107812e-04, 6.51257360e-05, 1.06845939e-04, 2.90082477e-04,
+            2.09049831e-04, 3.42945241e-05, 1.75379610e-04, 7.58524227e-04,
+            2.32403356e-04, 4.19872697e-04, 5.55520924e-03, 1.14839673e-02,
+            5.17792348e-03, 7.77232368e-03, 3.36388536e-02, 2.70598419e-02,
+            1.51852425e-02, 2.62977779e-02, 1.52134784e-02, 3.64050455e-03,
+        ])
+        # fmt: on
+        self.assertTrue(np.allclose(spec[64:128, 321], expected))
+
+    def test_spectrogram_batch_power(self):
+        waveform_list = self._load_datasamples(3)
+
+        spec_list = spectrogram_batch(
+            waveform_list,
+            window_function(400, "hann", frame_length=512),
+            frame_length=512,
+            hop_length=128,
+            power=None,
+        )
+        self.assertEqual(spec_list[0].shape, (257, 732))
+        self.assertEqual(spec_list[0].dtype, np.complex64)
+        self.assertEqual(spec_list[1].shape, (257, 602))
+        self.assertEqual(spec_list[1].dtype, np.complex64)
+        self.assertEqual(spec_list[2].shape, (257, 1561))
+        self.assertEqual(spec_list[2].dtype, np.complex64)
+
+        # fmt: off
+        expected1 = np.array([
+             0.01452305+0.01820039j, -0.01737362-0.01641946j,
+             0.0121028 +0.01565081j, -0.02794554-0.03021514j,
+             0.04719803+0.04086519j, -0.04391563-0.02779365j,
+             0.05682834+0.01571325j, -0.08604821-0.02023657j,
+             0.07497991+0.0186641j , -0.06366091-0.00922475j,
+             0.11003416+0.0114788j , -0.13677941-0.01523552j,
+             0.10934535-0.00117226j, -0.11635598+0.02551187j,
+             0.14708674-0.03469823j, -0.1328196 +0.06034218j,
+             0.12667368-0.13973421j, -0.14764774+0.18912019j,
+             0.10235471-0.12181523j, -0.00773012+0.04730498j,
+            -0.01487191-0.07312611j, -0.02739162+0.09619419j,
+             0.02895459-0.05398273j,  0.01198589+0.05276592j,
+            -0.02117299-0.10123465j,  0.00666388+0.09526499j,
+            -0.01672773-0.05649684j,  0.02723125+0.05939891j,
+            -0.01879361-0.062954j  ,  0.03686557+0.04568823j,
+            -0.07394181-0.07949649j,  0.06238583+0.13905765j,
+        ])
+        expected2 = np.array([
+            -0.01634146-7.0067253e-03j, -0.00068403+9.2661660e-03j,
+            0.00571721-3.9035487e-03j, -0.00915086+1.5033451e-03j,
+            0.01138636+5.4256055e-03j, -0.00294282-1.2016168e-02j,
+            -0.00428711+7.3687937e-03j, -0.001002  -1.3972387e-03j,
+            0.00622582+3.7551194e-03j, -0.00137886-7.0342086e-03j,
+            -0.00824075+3.8430823e-03j,  0.0107349 +7.1450039e-03j,
+            0.00363763-1.4242286e-02j, -0.01499857+1.7917662e-05j,
+            -0.0046242 +1.2500680e-02j,  0.02180984+7.2047939e-03j,
+            -0.00273568-1.6844695e-02j, -0.00178986-7.5209686e-03j,
+            -0.01661806+1.2662713e-03j, -0.01045276+2.0611197e-02j,
+            0.03252975+2.5592113e-02j,  0.03945662-6.7136563e-02j,
+            -0.10622615+4.9393820e-03j,  0.06684612+6.4607985e-02j,
+            -0.00753762-5.1637031e-02j, -0.00220644+1.8002450e-02j,
+            -0.00357443-4.1291970e-03j,  0.01463647-1.4063751e-03j,
+            -0.02252573-1.1189026e-02j,  0.00276293+1.9019062e-02j,
+            0.01216721+1.2095908e-03j,  0.00034753-7.4386634e-03j
+        ])
+        expected3 = np.array([
+            2.3276670e-02+0.0406534j, -2.4413882e-02-0.07868771j,
+            1.0993068e-02+0.05550544j, -1.5825305e-02+0.00480187j,
+            4.7617555e-02-0.04421869j, -7.1669750e-02+0.06317082j,
+            5.9706111e-02-0.08369736j, -2.2317577e-02+0.08915959j,
+            -2.3291381e-02-0.06601578j,  5.9362967e-02+0.03185856j,
+            -6.5269925e-02+0.0030586j,  5.0898481e-02-0.04319243j,
+            -4.0413942e-02+0.08051146j,  3.0059000e-02-0.09730332j,
+            -1.2479190e-02+0.09703682j, -6.1806822e-03-0.09617531j,
+            2.6907364e-02+0.08084074j, -4.1639723e-02-0.03391053j,
+            3.1113219e-02-0.01497662j,  3.4023849e-03+0.03632669j,
+            -4.9804080e-02-0.039231j,  8.9777440e-02+0.02577243j,
+            -9.2947647e-02+0.01514865j,  6.2368069e-02-0.05954866j,
+            -2.9966677e-02+0.06520324j, -8.2365885e-05-0.0440613j ,
+            2.0203773e-02+0.04350767j, -8.9924788e-04-0.05406843j,
+            -3.5951469e-02+0.03055602j,  3.3790238e-02+0.02182594j,
+            1.0919777e-03-0.06437822j, -1.8534327e-02+0.07866792j
+        ])
+        # fmt: on
+        self.assertTrue(np.allclose(spec_list[0][64:96, 321], expected1))
+        self.assertTrue(np.allclose(spec_list[1][64:96, 321], expected2))
+        self.assertTrue(np.allclose(spec_list[2][64:96, 321], expected3))
+
+        spec_list = spectrogram_batch(
+            waveform_list,
+            window_function(400, "hann", frame_length=512),
+            frame_length=512,
+            hop_length=128,
+            power=1.0,
+        )
+        self.assertEqual(spec_list[0].shape, (257, 732))
+        self.assertEqual(spec_list[0].dtype, np.float64)
+        self.assertEqual(spec_list[1].shape, (257, 602))
+        self.assertEqual(spec_list[1].dtype, np.float64)
+        self.assertEqual(spec_list[2].shape, (257, 1561))
+        self.assertEqual(spec_list[2].dtype, np.float64)
+
+        # fmt: off
+        expected1 = np.array([
+            0.02328461, 0.02390484, 0.01978448, 0.04115711, 0.0624309 ,
+            0.05197181, 0.05896072, 0.08839577, 0.07726794, 0.06432579,
+            0.11063128, 0.13762532, 0.10935163, 0.11911998, 0.15112405,
+            0.14588428, 0.18860507, 0.23992978, 0.15910825, 0.04793241,
+            0.07462307, 0.10001811, 0.06125769, 0.05411011, 0.10342509,
+            0.09549777, 0.05892122, 0.06534349, 0.06569936, 0.05870678,
+            0.10856833, 0.1524107 , 0.11463385, 0.05766969, 0.12385171,
+            0.14472842, 0.11978184, 0.10353675, 0.07244056, 0.03461861,
+            0.02624896, 0.02227475, 0.01238363, 0.00885281, 0.0110049 ,
+            0.00807005, 0.01033663, 0.01703181, 0.01445856, 0.00585615,
+            0.0132431 , 0.02754132, 0.01524478, 0.0204908 , 0.07453328,
+            0.10716327, 0.07195779, 0.08816078, 0.18340898, 0.16449876,
+            0.12322842, 0.1621659 , 0.12334293, 0.06033659,
+        ])
+        expected2 = np.array([
+            0.01778026, 0.00929138, 0.00692273, 0.00927352, 0.01261294,
+            0.01237128, 0.00852516, 0.00171938, 0.00727061, 0.00716808,
+            0.00909281, 0.01289532, 0.01469949, 0.01499858, 0.01332855,
+            0.02296907, 0.01706539, 0.00773101, 0.01666623, 0.02311021,
+            0.0413901, 0.07787261, 0.10634092, 0.09296556, 0.05218428,
+            0.01813716, 0.00546139, 0.01470388, 0.02515159, 0.0192187,
+            0.01222719, 0.00744678, 0.01045674, 0.01923522, 0.01990819,
+            0.01174323, 0.01535391, 0.02786647, 0.02904595, 0.0313408 ,
+            0.0340503, 0.03118268, 0.02915136, 0.04200513, 0.05563153,
+            0.05429446, 0.05021769, 0.05882667, 0.06668596, 0.06555867,
+            0.04523559, 0.01489498, 0.01031892, 0.02134155, 0.01736669,
+            0.0195216, 0.03971575, 0.03938636, 0.02052712, 0.03104931,
+            0.0902727, 0.09022622, 0.03275532, 0.0172633,
+        ])
+        expected3 = np.array([
+            0.04684551, 0.08238806, 0.05658358, 0.01653778, 0.06498249,
+            0.09553589, 0.10281084, 0.09191031, 0.07000408, 0.06737158,
+            0.06534155, 0.06675509, 0.09008541, 0.10184046, 0.09783596,
+            0.0963737, 0.08520112, 0.05370093, 0.03453015, 0.03648568,
+            0.06339967, 0.09340346, 0.09417402, 0.08623119, 0.07175977,
+            0.04406138, 0.04796988, 0.05407591, 0.0471824 , 0.04022626,
+            0.06438748, 0.0808218, 0.0745263, 0.06191467, 0.03116328,
+            0.03206497, 0.05867718, 0.04424652, 0.04448404, 0.07032498,
+            0.08300796, 0.07895744, 0.0816894, 0.09392357, 0.07571699,
+            0.03967651, 0.07703795, 0.06464871, 0.08704693, 0.14085226,
+            0.1350321, 0.18794712, 0.27043005, 0.26596246, 0.19948336,
+            0.06545141, 0.13204652, 0.08554521, 0.2262849, 0.33900721,
+            0.3970475, 0.3482436, 0.17134947, 0.46249565,
+        ])
+        # fmt: on
+        self.assertTrue(np.allclose(spec_list[0][64:128, 321], expected1))
+        self.assertTrue(np.allclose(spec_list[1][64:128, 321], expected2))
+        self.assertTrue(np.allclose(spec_list[2][64:128, 321], expected3))
+
+        spec_list = spectrogram_batch(
+            waveform_list,
+            window_function(400, "hann", frame_length=512),
+            frame_length=512,
+            hop_length=128,
+            power=2.0,
+        )
+        self.assertEqual(spec_list[0].shape, (257, 732))
+        self.assertEqual(spec_list[0].dtype, np.float64)
+        self.assertEqual(spec_list[1].shape, (257, 602))
+        self.assertEqual(spec_list[1].dtype, np.float64)
+        self.assertEqual(spec_list[2].shape, (257, 1561))
+        self.assertEqual(spec_list[2].dtype, np.float64)
+
+        # fmt: off
+        expected1 = np.array([
+            5.42173162e-04, 5.71441371e-04, 3.91425507e-04, 1.69390778e-03,
+            3.89761780e-03, 2.70106923e-03, 3.47636663e-03, 7.81381316e-03,
+            5.97033510e-03, 4.13780799e-03, 1.22392802e-02, 1.89407300e-02,
+            1.19577805e-02, 1.41895693e-02, 2.28384770e-02, 2.12822221e-02,
+            3.55718732e-02, 5.75663000e-02, 2.53154356e-02, 2.29751552e-03,
+            5.56860259e-03, 1.00036217e-02, 3.75250424e-03, 2.92790355e-03,
+            1.06967501e-02, 9.11982451e-03, 3.47171025e-03, 4.26977174e-03,
+            4.31640586e-03, 3.44648538e-03, 1.17870830e-02, 2.32290216e-02,
+            1.31409196e-02, 3.32579296e-03, 1.53392460e-02, 2.09463164e-02,
+            1.43476883e-02, 1.07198600e-02, 5.24763530e-03, 1.19844836e-03,
+            6.89007982e-04, 4.96164430e-04, 1.53354369e-04, 7.83722571e-05,
+            1.21107812e-04, 6.51257360e-05, 1.06845939e-04, 2.90082477e-04,
+            2.09049831e-04, 3.42945241e-05, 1.75379610e-04, 7.58524227e-04,
+            2.32403356e-04, 4.19872697e-04, 5.55520924e-03, 1.14839673e-02,
+            5.17792348e-03, 7.77232368e-03, 3.36388536e-02, 2.70598419e-02,
+            1.51852425e-02, 2.62977779e-02, 1.52134784e-02, 3.64050455e-03,
+        ])
+        expected2 = np.array([
+            3.16137604e-04, 8.63297362e-05, 4.79241720e-05, 8.59982493e-05,
+            1.59086326e-04, 1.53048476e-04, 7.26783945e-05, 2.95627100e-06,
+            5.28617352e-05, 5.13813355e-05, 8.26792588e-05, 1.66289156e-04,
+            2.16075069e-04, 2.24957314e-04, 1.77650211e-04, 5.27578282e-04,
+            2.91227688e-04, 5.97685493e-05, 2.77763360e-04, 5.34081651e-04,
+            1.71314057e-03, 6.06414277e-03, 1.13083916e-02, 8.64259617e-03,
+            2.72319867e-03, 3.28956593e-04, 2.98268126e-05, 2.16204145e-04,
+            6.32602626e-04, 3.69358508e-04, 1.49504171e-04, 5.54544917e-05,
+            1.09343371e-04, 3.69993847e-04, 3.96335839e-04, 1.37903521e-04,
+            2.35742483e-04, 7.76540114e-04, 8.43667068e-04, 9.82245923e-04,
+            1.15942286e-03, 9.72359636e-04, 8.49801853e-04, 1.76443092e-03,
+            3.09486753e-03, 2.94788822e-03, 2.52181630e-03, 3.46057723e-03,
+            4.44701769e-03, 4.29793858e-03, 2.04625858e-03, 2.21860290e-04,
+            1.06480179e-04, 4.55461892e-04, 3.01601836e-04, 3.81092892e-04,
+            1.57734053e-03, 1.55128531e-03, 4.21362677e-04, 9.64059883e-04,
+            8.14916019e-03, 8.14077014e-03, 1.07291131e-03, 2.98021545e-04,
+        ])
+        expected3 = np.array([
+            0.0021945 , 0.00678779, 0.0032017 , 0.0002735 , 0.00422272,
+            0.00912711, 0.01057007, 0.00844751, 0.00490057, 0.00453893,
+            0.00426952, 0.00445624, 0.00811538, 0.01037148, 0.00957188,
+            0.00928789, 0.00725923, 0.00288379, 0.00119233, 0.0013312 ,
+            0.00401952, 0.00872421, 0.00886875, 0.00743582, 0.00514946,
+            0.00194141, 0.00230111, 0.0029242 , 0.00222618, 0.00161815,
+            0.00414575, 0.00653216, 0.00555417, 0.00383343, 0.00097115,
+            0.00102816, 0.00344301, 0.00195775, 0.00197883, 0.0049456 ,
+            0.00689032, 0.00623428, 0.00667316, 0.00882164, 0.00573306,
+            0.00157423, 0.00593485, 0.00417946, 0.00757717, 0.01983936,
+            0.01823367, 0.03532412, 0.07313241, 0.07073603, 0.03979361,
+            0.00428389, 0.01743628, 0.00731798, 0.05120486, 0.11492589,
+            0.15764671, 0.1212736 , 0.02936064, 0.21390222
+        ])
+        # fmt: on
+        self.assertTrue(np.allclose(spec_list[0][64:128, 321], expected1))
+        self.assertTrue(np.allclose(spec_list[1][64:128, 321], expected2))
+        self.assertTrue(np.allclose(spec_list[2][64:128, 321], expected3))
+
+    def test_power_to_db(self):
+        spectrogram = np.zeros((2, 3))
+        spectrogram[0, 0] = 2.0
+        spectrogram[0, 1] = 0.5
+        spectrogram[0, 2] = 0.707
+        spectrogram[1, 1] = 1.0
+
+        output = power_to_db(spectrogram, reference=1.0)
+        expected = np.array([[3.01029996, -3.01029996, -1.50580586], [-100.0, 0.0, -100.0]])
+        self.assertTrue(np.allclose(output, expected))
+
+        output = power_to_db(spectrogram, reference=2.0)
+        expected = np.array([[0.0, -6.02059991, -4.51610582], [-103.01029996, -3.01029996, -103.01029996]])
+        self.assertTrue(np.allclose(output, expected))
+
+        output = power_to_db(spectrogram, min_value=1e-6)
+        expected = np.array([[3.01029996, -3.01029996, -1.50580586], [-60.0, 0.0, -60.0]])
+        self.assertTrue(np.allclose(output, expected))
+
+        output = power_to_db(spectrogram, db_range=80)
+        expected = np.array([[3.01029996, -3.01029996, -1.50580586], [-76.98970004, 0.0, -76.98970004]])
+        self.assertTrue(np.allclose(output, expected))
+
+        output = power_to_db(spectrogram, reference=2.0, db_range=80)
+        expected = np.array([[0.0, -6.02059991, -4.51610582], [-80.0, -3.01029996, -80.0]])
+        self.assertTrue(np.allclose(output, expected))
+
+        output = power_to_db(spectrogram, reference=2.0, min_value=1e-6, db_range=80)
+        expected = np.array([[0.0, -6.02059991, -4.51610582], [-63.01029996, -3.01029996, -63.01029996]])
+        self.assertTrue(np.allclose(output, expected))
+
+        with pytest.raises(ValueError):
+            power_to_db(spectrogram, reference=0.0)
+        with pytest.raises(ValueError):
+            power_to_db(spectrogram, min_value=0.0)
+        with pytest.raises(ValueError):
+            power_to_db(spectrogram, db_range=-80)
+
+    def test_power_to_db_batch(self):
+        # Setup a batch of spectrograms with varying values and lengths
+        batch_spectrogram = np.zeros((3, 2, 3))
+        batch_spectrogram[0, 0, 0] = 2.0
+        batch_spectrogram[0, 0, 1] = 0.5
+        batch_spectrogram[0, 0, 2] = 0.707
+        batch_spectrogram[0, 1, 1] = 1.0
+        batch_spectrogram[1, :, :2] = batch_spectrogram[0, :, :2] * 1.5
+        batch_spectrogram[2, :, :1] = batch_spectrogram[0, :, :1] * 0.5
+
+        # Expected values computed by applying `power_to_db` iteratively
+        output = power_to_db_batch(batch_spectrogram, reference=1.0)
+        expected = np.array(
+            [
+                [[3.01029996, -3.01029996, -1.50580586], [-100, 0, -100]],
+                [[4.77121255, -1.24938737, -100], [-100, 1.76091259, -100]],
+                [[0, -100, -100], [-100, -100, -100]],
+            ]
+        )
+        self.assertTrue(np.allclose(output, expected))
+
+        output = power_to_db_batch(batch_spectrogram, reference=2.0)
+        expected = np.array(
+            [
+                [[0, -6.02059991, -4.51610582], [-103.01029996, -3.01029996, -103.01029996]],
+                [[1.76091259, -4.25968732, -103.01029996], [-103.01029996, -1.24938737, -103.01029996]],
+                [[-3.01029996, -103.01029996, -103.01029996], [-103.01029996, -103.01029996, -103.01029996]],
+            ]
+        )
+        self.assertTrue(np.allclose(output, expected))
+
+        output = power_to_db_batch(batch_spectrogram, min_value=1e-6)
+        expected = np.array(
+            [
+                [[3.01029996, -3.01029996, -1.50580586], [-60, 0, -60]],
+                [[4.77121255, -1.24938737, -60], [-60, 1.76091259, -60]],
+                [[0, -60, -60], [-60, -60, -60]],
+            ]
+        )
+        self.assertTrue(np.allclose(output, expected))
+
+        output = power_to_db_batch(batch_spectrogram, db_range=80)
+        expected = np.array(
+            [
+                [[3.01029996, -3.01029996, -1.50580586], [-76.98970004, 0, -76.98970004]],
+                [[4.77121255, -1.24938737, -75.22878745], [-75.22878745, 1.76091259, -75.22878745]],
+                [[0, -80, -80], [-80, -80, -80]],
+            ]
+        )
+        self.assertTrue(np.allclose(output, expected))
+
+        output = power_to_db_batch(batch_spectrogram, reference=2.0, db_range=80)
+        expected = np.array(
+            [
+                [[0, -6.02059991, -4.51610582], [-80, -3.01029996, -80]],
+                [[1.76091259, -4.25968732, -78.23908741], [-78.23908741, -1.24938737, -78.23908741]],
+                [[-3.01029996, -83.01029996, -83.01029996], [-83.01029996, -83.01029996, -83.01029996]],
+            ]
+        )
+        self.assertTrue(np.allclose(output, expected))
+
+        output = power_to_db_batch(batch_spectrogram, reference=2.0, min_value=1e-6, db_range=80)
+        expected = np.array(
+            [
+                [[0, -6.02059991, -4.51610582], [-63.01029996, -3.01029996, -63.01029996]],
+                [[1.76091259, -4.25968732, -63.01029996], [-63.01029996, -1.24938737, -63.01029996]],
+                [[-3.01029996, -63.01029996, -63.01029996], [-63.01029996, -63.01029996, -63.01029996]],
+            ]
+        )
+        self.assertTrue(np.allclose(output, expected))
+
+        with pytest.raises(ValueError):
+            power_to_db_batch(batch_spectrogram, reference=0.0)
+        with pytest.raises(ValueError):
+            power_to_db_batch(batch_spectrogram, min_value=0.0)
+        with pytest.raises(ValueError):
+            power_to_db_batch(batch_spectrogram, db_range=-80)
+
+    def test_amplitude_to_db(self):
+        spectrogram = np.zeros((2, 3))
+        spectrogram[0, 0] = 2.0
+        spectrogram[0, 1] = 0.5
+        spectrogram[0, 2] = 0.707
+        spectrogram[1, 1] = 1.0
+
+        output = amplitude_to_db(spectrogram, reference=1.0)
+        expected = np.array([[6.02059991, -6.02059991, -3.01161172], [-100.0, 0.0, -100.0]])
+        self.assertTrue(np.allclose(output, expected))
+
+        output = amplitude_to_db(spectrogram, reference=2.0)
+        expected = np.array([[0.0, -12.04119983, -9.03221164], [-106.02059991, -6.02059991, -106.02059991]])
+        self.assertTrue(np.allclose(output, expected))
+
+        output = amplitude_to_db(spectrogram, min_value=1e-3)
+        expected = np.array([[6.02059991, -6.02059991, -3.01161172], [-60.0, 0.0, -60.0]])
+        self.assertTrue(np.allclose(output, expected))
+
+        output = amplitude_to_db(spectrogram, db_range=80)
+        expected = np.array([[6.02059991, -6.02059991, -3.01161172], [-73.97940009, 0.0, -73.97940009]])
+        self.assertTrue(np.allclose(output, expected))
+
+        output = amplitude_to_db(spectrogram, reference=2.0, db_range=80)
+        expected = np.array([[0.0, -12.04119983, -9.03221164], [-80.0, -6.02059991, -80.0]])
+        self.assertTrue(np.allclose(output, expected))
+
+        output = amplitude_to_db(spectrogram, reference=2.0, min_value=1e-3, db_range=80)
+        expected = np.array([[0.0, -12.04119983, -9.03221164], [-66.02059991, -6.02059991, -66.02059991]])
+        self.assertTrue(np.allclose(output, expected))
+
+        with pytest.raises(ValueError):
+            amplitude_to_db(spectrogram, reference=0.0)
+        with pytest.raises(ValueError):
+            amplitude_to_db(spectrogram, min_value=0.0)
+        with pytest.raises(ValueError):
+            amplitude_to_db(spectrogram, db_range=-80)
+
+    def test_amplitude_to_db_batch(self):
+        # Setup a batch of spectrograms with varying values and lengths
+        batch_spectrogram = np.zeros((3, 2, 3))
+        batch_spectrogram[0, 0, 0] = 2.0
+        batch_spectrogram[0, 0, 1] = 0.5
+        batch_spectrogram[0, 0, 2] = 0.707
+        batch_spectrogram[0, 1, 1] = 1.0
+        batch_spectrogram[1, :, :2] = batch_spectrogram[0, :, :2] * 1.5
+        batch_spectrogram[2, :, :1] = batch_spectrogram[0, :, :1] * 0.5
+
+        # Expected values computed by applying `amplitude_to_db` iteratively
+        output = amplitude_to_db_batch(batch_spectrogram, reference=1.0)
+        expected = np.array(
+            [
+                [[6.02059991, -6.02059991, -3.01161172], [-100, 0, -100]],
+                [[9.54242509, -2.49877473, -100], [-100, 3.52182518, -100]],
+                [[0, -100, -100], [-100, -100, -100]],
+            ]
+        )
+        self.assertTrue(np.allclose(output, expected))
+
+        output = amplitude_to_db_batch(batch_spectrogram, reference=2.0)
+        expected = np.array(
+            [
+                [[0, -12.04119983, -9.03221164], [-106.02059991, -6.02059991, -106.02059991]],
+                [[3.52182518, -8.51937465, -106.02059991], [-106.02059991, -2.49877473, -106.02059991]],
+                [[-6.02059991, -106.02059991, -106.02059991], [-106.02059991, -106.02059991, -106.02059991]],
+            ]
+        )
+        self.assertTrue(np.allclose(output, expected))
+
+        output = amplitude_to_db_batch(batch_spectrogram, min_value=1e-3)
+        expected = np.array(
+            [
+                [[6.02059991, -6.02059991, -3.01161172], [-60, 0, -60]],
+                [[9.54242509, -2.49877473, -60], [-60, 3.52182518, -60]],
+                [[0, -60, -60], [-60, -60, -60]],
+            ]
+        )
+        self.assertTrue(np.allclose(output, expected))
+
+        output = amplitude_to_db_batch(batch_spectrogram, db_range=80)
+        expected = np.array(
+            [
+                [[6.02059991, -6.02059991, -3.01161172], [-73.97940009, 0, -73.97940009]],
+                [[9.54242509, -2.49877473, -70.45757491], [-70.45757491, 3.52182518, -70.45757491]],
+                [[0, -80, -80], [-80, -80, -80]],
+            ]
+        )
+        self.assertTrue(np.allclose(output, expected))
+
+        output = amplitude_to_db_batch(batch_spectrogram, reference=2.0, db_range=80)
+        expected = np.array(
+            [
+                [[0, -12.04119983, -9.03221164], [-80, -6.02059991, -80]],
+                [[3.52182518, -8.51937465, -76.47817482], [-76.47817482, -2.49877473, -76.47817482]],
+                [[-6.02059991, -86.02059991, -86.02059991], [-86.02059991, -86.02059991, -86.02059991]],
+            ]
+        )
+        self.assertTrue(np.allclose(output, expected))
+
+        output = amplitude_to_db_batch(batch_spectrogram, reference=2.0, min_value=1e-3, db_range=80)
+        expected = np.array(
+            [
+                [[0, -12.04119983, -9.03221164], [-66.02059991, -6.02059991, -66.02059991]],
+                [[3.52182518, -8.51937465, -66.02059991], [-66.02059991, -2.49877473, -66.02059991]],
+                [[-6.02059991, -66.02059991, -66.02059991], [-66.02059991, -66.02059991, -66.02059991]],
+            ]
+        )
+        self.assertTrue(np.allclose(output, expected))
+
+        with pytest.raises(ValueError):
+            amplitude_to_db_batch(batch_spectrogram, reference=0.0)
+        with pytest.raises(ValueError):
+            amplitude_to_db_batch(batch_spectrogram, min_value=0.0)
+        with pytest.raises(ValueError):
+            amplitude_to_db_batch(batch_spectrogram, db_range=-80)
+
+    @require_librosa
+    def test_chroma_equivalence(self):
+        num_frequency_bins = 25
+        num_chroma = 6
+        sampling_rate = 24000
+
+        # test default parameters
+        original_chroma = chroma(sr=sampling_rate, n_chroma=num_chroma, n_fft=num_frequency_bins)
+        utils_chroma = chroma_filter_bank(
+            num_frequency_bins=num_frequency_bins, num_chroma=num_chroma, sampling_rate=sampling_rate
+        )
+
+        self.assertTrue(np.allclose(original_chroma, utils_chroma))
+
+        # test no weighting_parameters
+        original_chroma = chroma(sr=sampling_rate, n_chroma=num_chroma, n_fft=num_frequency_bins, octwidth=None)
+        utils_chroma = chroma_filter_bank(
+            num_frequency_bins=num_frequency_bins,
+            num_chroma=num_chroma,
+            sampling_rate=sampling_rate,
+            weighting_parameters=None,
+        )
+
+        self.assertTrue(np.allclose(original_chroma, utils_chroma))
+
+        # test with L1 norm
+        original_chroma = chroma(sr=sampling_rate, n_chroma=num_chroma, n_fft=num_frequency_bins, norm=1.0)
+        utils_chroma = chroma_filter_bank(
+            num_frequency_bins=num_frequency_bins, num_chroma=num_chroma, sampling_rate=sampling_rate, power=1.0
+        )
+
+        self.assertTrue(np.allclose(original_chroma, utils_chroma))
+
+        # test starting at 'A' chroma, power = None, tuning = 0, different weighting_parameters
+        original_chroma = chroma(
+            sr=sampling_rate,
+            n_chroma=num_chroma,
+            n_fft=num_frequency_bins,
+            norm=None,
+            base_c=None,
+            octwidth=1.0,
+            ctroct=4.0,
+        )
+        utils_chroma = chroma_filter_bank(
+            num_frequency_bins=num_frequency_bins,
+            num_chroma=num_chroma,
+            sampling_rate=sampling_rate,
+            power=None,
+            start_at_c_chroma=False,
+            weighting_parameters=(4.0, 1.0),
+        )
+
+        self.assertTrue(np.allclose(original_chroma, utils_chroma))

docs/transformers/tests/utils/test_backbone_utils.py

@@ -0,0 +1,272 @@
+# Copyright 2023 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import pytest
+
+from transformers import DetrConfig, MaskFormerConfig, ResNetBackbone, ResNetConfig, TimmBackbone
+from transformers.testing_utils import require_torch, slow
+from transformers.utils.backbone_utils import (
+    BackboneMixin,
+    get_aligned_output_features_output_indices,
+    load_backbone,
+    verify_out_features_out_indices,
+)
+from transformers.utils.import_utils import is_torch_available
+
+
+if is_torch_available():
+    import torch
+
+    from transformers import BertPreTrainedModel
+
+
+class BackboneUtilsTester(unittest.TestCase):
+    def test_get_aligned_output_features_output_indices(self):
+        stage_names = ["a", "b", "c"]
+
+        # Defaults to last layer if both are None
+        out_features, out_indices = get_aligned_output_features_output_indices(None, None, stage_names)
+        self.assertEqual(out_features, ["c"])
+        self.assertEqual(out_indices, [2])
+
+        # Out indices set to match out features
+        out_features, out_indices = get_aligned_output_features_output_indices(["a", "c"], None, stage_names)
+        self.assertEqual(out_features, ["a", "c"])
+        self.assertEqual(out_indices, [0, 2])
+
+        # Out features set to match out indices
+        out_features, out_indices = get_aligned_output_features_output_indices(None, [0, 2], stage_names)
+        self.assertEqual(out_features, ["a", "c"])
+        self.assertEqual(out_indices, [0, 2])
+
+        # Out features selected from negative indices
+        out_features, out_indices = get_aligned_output_features_output_indices(None, [-3, -1], stage_names)
+        self.assertEqual(out_features, ["a", "c"])
+        self.assertEqual(out_indices, [-3, -1])
+
+    def test_verify_out_features_out_indices(self):
+        # Stage names must be set
+        with pytest.raises(ValueError, match="Stage_names must be set for transformers backbones"):
+            verify_out_features_out_indices(["a", "b"], (0, 1), None)
+
+        # Out features must be a list
+        with pytest.raises(ValueError, match="out_features must be a list got <class 'tuple'>"):
+            verify_out_features_out_indices(("a", "b"), (0, 1), ["a", "b"])
+
+        # Out features must be a subset of stage names
+        with pytest.raises(
+            ValueError, match=r"out_features must be a subset of stage_names: \['a'\] got \['a', 'b'\]"
+        ):
+            verify_out_features_out_indices(["a", "b"], [0, 1], ["a"])
+
+        # Out features must contain no duplicates
+        with pytest.raises(ValueError, match=r"out_features must not contain any duplicates, got \['a', 'a'\]"):
+            verify_out_features_out_indices(["a", "a"], None, ["a"])
+
+        # Out indices must be a list
+        with pytest.raises(ValueError, match="out_indices must be a list, got <class 'int'>"):
+            verify_out_features_out_indices(None, 0, ["a", "b"])
+
+        with pytest.raises(ValueError, match="out_indices must be a list, got <class 'tuple'>"):
+            verify_out_features_out_indices(None, (0, 1), ["a", "b"])
+
+        # Out indices must be a subset of stage names
+        with pytest.raises(
+            ValueError, match=r"out_indices must be valid indices for stage_names \['a'\], got \[0, 1\]"
+        ):
+            verify_out_features_out_indices(None, [0, 1], ["a"])
+
+        # Out indices must contain no duplicates
+        with pytest.raises(ValueError, match=r"out_indices must not contain any duplicates, got \[0, 0\]"):
+            verify_out_features_out_indices(None, [0, 0], ["a"])
+
+        # Out features and out indices must be the same length
+        with pytest.raises(
+            ValueError, match="out_features and out_indices should have the same length if both are set"
+        ):
+            verify_out_features_out_indices(["a", "b"], [0], ["a", "b", "c"])
+
+        # Out features should match out indices
+        with pytest.raises(
+            ValueError, match="out_features and out_indices should correspond to the same stages if both are set"
+        ):
+            verify_out_features_out_indices(["a", "b"], [0, 2], ["a", "b", "c"])
+
+        # Out features and out indices should be in order
+        with pytest.raises(
+            ValueError,
+            match=r"out_features must be in the same order as stage_names, expected \['a', 'b'\] got \['b', 'a'\]",
+        ):
+            verify_out_features_out_indices(["b", "a"], [0, 1], ["a", "b"])
+
+        with pytest.raises(
+            ValueError, match=r"out_indices must be in the same order as stage_names, expected \[-2, 1\] got \[1, -2\]"
+        ):
+            verify_out_features_out_indices(["a", "b"], [1, -2], ["a", "b"])
+
+        # Check passes with valid inputs
+        verify_out_features_out_indices(["a", "b", "d"], [0, 1, -1], ["a", "b", "c", "d"])
+
+    def test_backbone_mixin(self):
+        backbone = BackboneMixin()
+
+        backbone.stage_names = ["a", "b", "c"]
+        backbone._out_features = ["a", "c"]
+        backbone._out_indices = [0, 2]
+
+        # Check that the output features and indices are set correctly
+        self.assertEqual(backbone.out_features, ["a", "c"])
+        self.assertEqual(backbone.out_indices, [0, 2])
+
+        # Check out features and indices are updated correctly
+        backbone.out_features = ["a", "b"]
+        self.assertEqual(backbone.out_features, ["a", "b"])
+        self.assertEqual(backbone.out_indices, [0, 1])
+
+        backbone.out_indices = [-3, -1]
+        self.assertEqual(backbone.out_features, ["a", "c"])
+        self.assertEqual(backbone.out_indices, [-3, -1])
+
+    @slow
+    @require_torch
+    def test_load_backbone_from_config(self):
+        """
+        Test that load_backbone correctly loads a backbone from a backbone config.
+        """
+        config = MaskFormerConfig(backbone_config=ResNetConfig(out_indices=(0, 2)))
+        backbone = load_backbone(config)
+        self.assertEqual(backbone.out_features, ["stem", "stage2"])
+        self.assertEqual(backbone.out_indices, (0, 2))
+        self.assertIsInstance(backbone, ResNetBackbone)
+
+    @slow
+    @require_torch
+    def test_load_backbone_from_checkpoint(self):
+        """
+        Test that load_backbone correctly loads a backbone from a checkpoint.
+        """
+        config = MaskFormerConfig(backbone="microsoft/resnet-18", backbone_config=None)
+        backbone = load_backbone(config)
+        self.assertEqual(backbone.out_indices, [4])
+        self.assertEqual(backbone.out_features, ["stage4"])
+        self.assertIsInstance(backbone, ResNetBackbone)
+
+        config = MaskFormerConfig(
+            backbone="resnet18",
+            use_timm_backbone=True,
+        )
+        backbone = load_backbone(config)
+        # We can't know ahead of time the exact output features and indices, or the layer names before
+        # creating the timm model, so it defaults to the last layer (-1,) and has a different layer name
+        self.assertEqual(backbone.out_indices, (-1,))
+        self.assertEqual(backbone.out_features, ["layer4"])
+        self.assertIsInstance(backbone, TimmBackbone)
+
+    @slow
+    @require_torch
+    def test_load_backbone_backbone_kwargs(self):
+        """
+        Test that load_backbone correctly configures the loaded backbone with the provided kwargs.
+        """
+        config = MaskFormerConfig(backbone="resnet18", use_timm_backbone=True, backbone_kwargs={"out_indices": (0, 1)})
+        backbone = load_backbone(config)
+        self.assertEqual(backbone.out_indices, (0, 1))
+        self.assertIsInstance(backbone, TimmBackbone)
+
+        config = MaskFormerConfig(backbone="microsoft/resnet-18", backbone_kwargs={"out_indices": (0, 2)})
+        backbone = load_backbone(config)
+        self.assertEqual(backbone.out_indices, (0, 2))
+        self.assertIsInstance(backbone, ResNetBackbone)
+
+        # Check can't be passed with a backone config
+        with pytest.raises(ValueError):
+            config = MaskFormerConfig(
+                backbone="microsoft/resnet-18",
+                backbone_config=ResNetConfig(out_indices=(0, 2)),
+                backbone_kwargs={"out_indices": (0, 1)},
+            )
+
+    @slow
+    @require_torch
+    def test_load_backbone_in_new_model(self):
+        """
+        Tests that new model can be created, with its weights instantiated and pretrained backbone weights loaded.
+        """
+
+        # Inherit from PreTrainedModel to ensure that the weights are initialized
+        class NewModel(BertPreTrainedModel):
+            def __init__(self, config):
+                super().__init__(config)
+                self.backbone = load_backbone(config)
+                self.layer_0 = torch.nn.Linear(config.hidden_size, config.hidden_size)
+                self.layer_1 = torch.nn.Linear(config.hidden_size, config.hidden_size)
+
+        def get_equal_not_equal_weights(model_0, model_1):
+            equal_weights = []
+            not_equal_weights = []
+            for (k0, v0), (k1, v1) in zip(model_0.named_parameters(), model_1.named_parameters()):
+                self.assertEqual(k0, k1)
+                weights_are_equal = torch.allclose(v0, v1)
+                if weights_are_equal:
+                    equal_weights.append(k0)
+                else:
+                    not_equal_weights.append(k0)
+            return equal_weights, not_equal_weights
+
+        config = MaskFormerConfig(use_pretrained_backbone=False, backbone="microsoft/resnet-18")
+        model_0 = NewModel(config)
+        model_1 = NewModel(config)
+        equal_weights, not_equal_weights = get_equal_not_equal_weights(model_0, model_1)
+
+        # Norm layers are always initialized with the same weights
+        equal_weights = [w for w in equal_weights if "normalization" not in w]
+        self.assertEqual(len(equal_weights), 0)
+        self.assertEqual(len(not_equal_weights), 24)
+
+        # Now we create a new model with backbone weights that are pretrained
+        config.use_pretrained_backbone = True
+        model_0 = NewModel(config)
+        model_1 = NewModel(config)
+        equal_weights, not_equal_weights = get_equal_not_equal_weights(model_0, model_1)
+
+        # Norm layers are always initialized with the same weights
+        equal_weights = [w for w in equal_weights if "normalization" not in w]
+        self.assertEqual(len(equal_weights), 20)
+        # Linear layers are still initialized randomly
+        self.assertEqual(len(not_equal_weights), 4)
+
+        # Check loading in timm backbone
+        config = DetrConfig(use_pretrained_backbone=False, backbone="resnet18", use_timm_backbone=True)
+        model_0 = NewModel(config)
+        model_1 = NewModel(config)
+        equal_weights, not_equal_weights = get_equal_not_equal_weights(model_0, model_1)
+
+        # Norm layers are always initialized with the same weights
+        equal_weights = [w for w in equal_weights if "bn" not in w and "downsample.1" not in w]
+        self.assertEqual(len(equal_weights), 0)
+        self.assertEqual(len(not_equal_weights), 24)
+
+        # Now we create a new model with backbone weights that are pretrained
+        config.use_pretrained_backbone = True
+        model_0 = NewModel(config)
+        model_1 = NewModel(config)
+        equal_weights, not_equal_weights = get_equal_not_equal_weights(model_0, model_1)
+
+        # Norm layers are always initialized with the same weights
+        equal_weights = [w for w in equal_weights if "bn" not in w and "downsample.1" not in w]
+        self.assertEqual(len(equal_weights), 20)
+        # Linear layers are still initialized randomly
+        self.assertEqual(len(not_equal_weights), 4)

docs/transformers/tests/utils/test_cache_utils.py

@@ -0,0 +1,766 @@
+# Copyright 2023 HuggingFace Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import copy
+import unittest
+
+from parameterized import parameterized
+
+from transformers import set_seed
+from transformers.testing_utils import (
+    CaptureStderr,
+    get_gpu_count,
+    is_torch_available,
+    require_gptq,
+    require_non_xpu,
+    require_read_token,
+    require_torch,
+    require_torch_accelerator,
+    require_torch_gpu,
+    require_torch_multi_gpu,
+    slow,
+    torch_device,
+)
+
+
+if is_torch_available():
+    import torch
+
+    from transformers import (
+        AutoModelForCausalLM,
+        AutoTokenizer,
+        ClvpForCausalLM,
+        DynamicCache,
+        GenerationConfig,
+        LlamaConfig,
+        SinkCache,
+        StaticCache,
+        convert_and_export_with_cache,
+    )
+    from transformers.utils import is_torch_greater_or_equal
+
+
+@require_torch
+class CacheTest(unittest.TestCase):
+    def test_dynamic_cache_retrocompatibility(self):
+        """Tests that we can convert back and forth between the legacy cache format and DynamicCache"""
+        legacy_cache = ()
+        new_cache = DynamicCache()
+
+        # Creates a new cache with 10 layers in both formats
+        for layer_idx in range(10):
+            new_key = torch.rand((2, 4, 8, 16))
+            new_value = torch.rand((2, 4, 8, 16))
+            new_cache.update(new_key, new_value, layer_idx)
+            legacy_cache += ((new_key, new_value),)
+
+        # Sanity check 1: they must have the same shapes
+        self.assertTrue(len(legacy_cache), len(new_cache))
+        for layer_idx in range(10):
+            self.assertTrue(len(legacy_cache[layer_idx]), len(legacy_cache[layer_idx]))
+            for key_value_idx in range(2):
+                self.assertTrue(
+                    legacy_cache[layer_idx][key_value_idx].shape == new_cache[layer_idx][key_value_idx].shape
+                )
+
+        # Sanity check 2: we can get the sequence length in multiple ways with DynamicCache, and they return the
+        # expected value
+        self.assertTrue(legacy_cache[0][0].shape[-2] == new_cache[0][0].shape[-2] == new_cache.get_seq_length() == 8)
+
+        # Sanity check 3: they must be equal, and both support indexing
+        for layer_idx in range(10):
+            for key_value_idx in range(2):
+                self.assertTrue(
+                    torch.allclose(new_cache[layer_idx][key_value_idx], legacy_cache[layer_idx][key_value_idx])
+                )
+
+        # Test 1: We can convert from legacy to new with no changes
+        from_legacy = DynamicCache.from_legacy_cache(legacy_cache)
+        for layer_idx in range(10):
+            for key_value_idx in range(2):
+                self.assertTrue(
+                    torch.allclose(from_legacy[layer_idx][key_value_idx], legacy_cache[layer_idx][key_value_idx])
+                )
+
+        # Test 2: We can convert from new to legacy with no changes
+        to_legacy = new_cache.to_legacy_cache()
+        for layer_idx in range(10):
+            for key_value_idx in range(2):
+                self.assertTrue(
+                    torch.allclose(to_legacy[layer_idx][key_value_idx], new_cache[layer_idx][key_value_idx])
+                )
+
+    def test_reorder_cache_retrocompatibility(self):
+        """Tests that Cache.reorder_cache is retrocompatible with the legacy code path"""
+        legacy_reorder_fn = ClvpForCausalLM._reorder_cache  # An example of a legacy `_reorder_cache` function
+
+        legacy_cache = ()
+        new_cache = DynamicCache()
+
+        # Creates a new cache with 10 layers in both formats
+        for layer_idx in range(10):
+            new_key = torch.rand((4, 4, 8, 16))
+            new_value = torch.rand((4, 4, 8, 16))
+            new_cache.update(new_key, new_value, layer_idx)
+            legacy_cache += ((new_key, new_value),)
+
+        # Let's create some dummy beam indices. From the shape above, it is equivalent to the case where num_beams=4
+        # and batch_size=1
+        beam_idx = torch.randint(low=0, high=4, size=(4,))
+
+        legacy_cache_reordered = legacy_reorder_fn(legacy_cache, beam_idx)
+        new_cache.reorder_cache(beam_idx)
+
+        # Let's check that the results are the same
+        for layer_idx in range(10):
+            for key_value_idx in range(2):
+                self.assertTrue(
+                    torch.allclose(
+                        new_cache[layer_idx][key_value_idx], legacy_cache_reordered[layer_idx][key_value_idx]
+                    )
+                )
+
+    def test_static_cache_mha_mqa_gqa(self):
+        """
+        Tests that static cache works with multi-head attention (MHA), grouped query attention (GQA), and multi-query
+        attention (MQA)
+        """
+
+        def _random_kvs(config):
+            # shape for key and values: (batch_size, num_heads, seq_len, head_dim)
+            random_keys = torch.rand(
+                (1, config.num_key_value_heads, 1, config.hidden_size // config.num_attention_heads),
+                device=torch_device,
+            )
+            random_values = torch.rand(
+                (1, config.num_key_value_heads, 1, config.hidden_size // config.num_attention_heads),
+                device=torch_device,
+            )
+            return random_keys, random_values
+
+        mha_config = LlamaConfig(num_attention_heads=32)
+        mha_static_cache = StaticCache(config=mha_config, max_batch_size=1, max_cache_len=10, device=torch_device)
+        cached_keys, cached_values = mha_static_cache.update(
+            *_random_kvs(mha_config), 0, cache_kwargs={"cache_position": torch.arange(1).to(torch_device)}
+        )
+        self.assertTrue(cached_keys.shape == (1, 32, 10, 128))
+        self.assertTrue(cached_values.shape == (1, 32, 10, 128))
+
+        gqa_config = LlamaConfig(num_attention_heads=32, num_key_value_heads=4)
+        gqa_static_cache = StaticCache(config=gqa_config, max_batch_size=1, max_cache_len=10, device=torch_device)
+        cached_keys, cached_values = gqa_static_cache.update(
+            *_random_kvs(gqa_config), 0, cache_kwargs={"cache_position": torch.arange(1).to(torch_device)}
+        )
+        self.assertTrue(cached_keys.shape == (1, 4, 10, 128))
+        self.assertTrue(cached_values.shape == (1, 4, 10, 128))
+
+        mqa_config = LlamaConfig(num_attention_heads=32, num_key_value_heads=1)
+        mqa_static_cache = StaticCache(config=mqa_config, max_batch_size=1, max_cache_len=10, device=torch_device)
+        cached_keys, cached_values = mqa_static_cache.update(
+            *_random_kvs(mqa_config), 0, cache_kwargs={"cache_position": torch.arange(1).to(torch_device)}
+        )
+        self.assertTrue(cached_keys.shape == (1, 1, 10, 128))
+        self.assertTrue(cached_values.shape == (1, 1, 10, 128))
+
+    def test_dynamic_cache_exportability(self):
+        model = AutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-MistralForCausalLM")
+        model = model.eval()
+        tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-MistralForCausalLM")
+        prompt = "What is the best way to debug python script?"
+        inputs = tokenizer(prompt, return_tensors="pt")
+        attention_mask = inputs.attention_mask
+        input_ids = inputs.input_ids
+
+        past_key_values = DynamicCache()
+        ep = torch.export.export(
+            model,
+            (),
+            {
+                "input_ids": input_ids,
+                "attention_mask": attention_mask,
+                "past_key_values": past_key_values,
+                "use_cache": True,
+            },
+            strict=False,
+        )
+        res = ep.module()(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            past_key_values=past_key_values,
+            use_cache=True,
+        )
+        self.assertTrue(len(res.past_key_values.key_cache) == model.config.num_hidden_layers)
+        self.assertEqual(2 * model.config.num_hidden_layers + 1, len(ep.graph_signature.output_specs))
+        self.assertEqual(
+            3,
+            len(
+                [
+                    x
+                    for x in ep.graph_signature.input_specs
+                    if x.kind == torch.export.graph_signature.InputKind.USER_INPUT
+                ]
+            ),
+        )
+
+        past_key_values_eager = DynamicCache()
+        res_eager = model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            past_key_values=past_key_values_eager,
+            use_cache=True,
+        )
+        self.assertTrue(torch.allclose(res.logits, res_eager.logits))
+        for k1, k2 in zip(res.past_key_values.key_cache, res_eager.past_key_values.key_cache):
+            self.assertTrue(torch.allclose(k1, k2))
+
+        for v1, v2 in zip(res.past_key_values.value_cache, res_eager.past_key_values.value_cache):
+            self.assertTrue(torch.allclose(v1, v2))
+
+    @slow
+    @require_read_token
+    def test_static_cache_exportability(self):
+        """
+        Tests that static cache works with `torch.export()`
+        """
+        if not is_torch_greater_or_equal("2.3"):
+            self.skipTest(reason="This test requires torch >= 2.3 to run.")
+
+        set_seed(0)
+        device = "cpu"
+        dtype = "bfloat16"
+        cache_implementation = "static"
+        attn_implementation = "sdpa"  # Export and ExecuTorch only works for SdpaAttention
+        batch_size = 1
+        max_cache_len = 1234
+        model = AutoModelForCausalLM.from_pretrained(
+            "google/gemma-2b",
+            device_map=device,
+            torch_dtype=dtype,
+            attn_implementation=attn_implementation,
+            generation_config=GenerationConfig(
+                use_cache=True,
+                cache_implementation=cache_implementation,
+                max_length=max_cache_len,
+                cache_config={
+                    "batch_size": batch_size,
+                    "max_cache_len": max_cache_len,
+                    "device": device,
+                },
+            ),
+        )
+        # Check if cache config is passed through correctly
+        self.assertEqual(model.generation_config.use_cache, True)
+        self.assertEqual(model.generation_config.cache_implementation, cache_implementation)
+        self.assertEqual(model.generation_config.max_length, max_cache_len)
+        self.assertTrue(model.generation_config.cache_config is not None)
+        self.assertEqual(model.generation_config.cache_config.batch_size, batch_size)
+        self.assertEqual(model.generation_config.cache_config.max_cache_len, max_cache_len)
+
+        exported_program = convert_and_export_with_cache(model)
+
+        # Check if the exported model is configured with the `StaticCache` correctly
+        n_static_key_caches = n_static_value_caches = 0
+        for buffer_name, buffer in exported_program.named_buffers():
+            if buffer_name.startswith("key_cache"):
+                self.assertTrue(buffer.shape[0] == batch_size)
+                self.assertTrue(buffer.shape[2] == max_cache_len)
+                n_static_key_caches = n_static_key_caches + 1
+            if buffer_name.startswith("value_cache"):
+                self.assertTrue(buffer.shape[0] == batch_size)
+                self.assertTrue(buffer.shape[2] == max_cache_len)
+                n_static_value_caches = n_static_value_caches + 1
+        self.assertEqual(n_static_key_caches, model.config.num_hidden_layers)
+        self.assertEqual(n_static_value_caches, model.config.num_hidden_layers)
+
+
+@require_torch_accelerator
+@slow
+class CacheIntegrationTest(unittest.TestCase):
+    def test_dynamic_cache_hard(self):
+        tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf", padding_side="left")
+        model = AutoModelForCausalLM.from_pretrained(
+            "meta-llama/Llama-2-7b-hf", device_map="auto", torch_dtype=torch.float16
+        )
+        inputs = tokenizer(["Here's everything I know about cats. Cats"], return_tensors="pt").to(model.device)
+
+        # DynamicCache and the legacy cache format should be equivalent
+        set_seed(0)
+        gen_out_legacy = model.generate(**inputs, do_sample=True, max_new_tokens=256)
+        set_seed(0)
+        gen_out = model.generate(**inputs, do_sample=True, max_new_tokens=256, past_key_values=DynamicCache())
+        self.assertListEqual(gen_out_legacy.tolist(), gen_out.tolist())
+
+        decoded = tokenizer.batch_decode(gen_out, skip_special_tokens=True)
+        expected_text = (
+            "Here's everything I know about cats. Cats are mysterious creatures. They can't talk, and they don't like "
+            "to be held. They don't play fetch, and they don't like to be hugged. But they do like to be petted.\n"
+            "Cats are also very independent. They don't like to be told what to do, and they don't like to be told "
+            "what to eat. They are also very territorial. They don't like to share their food or their toys.\nCats "
+            "are also very curious. They like to explore, and they like to play. They are also very fast. They can "
+            "run very fast, and they can jump very high.\nCats are also very smart. They can learn tricks, and they "
+            "can solve problems. They are also very playful. They like to play with toys, and they like to play with "
+            "other cats.\nCats are also very affectionate. They like to be petted, and they like to be held. They "
+            "also like to be scratched.\nCats are also very clean. They like to groom themselves, and they like to "
+            "clean their litter box.\nCats are also very independent. They don't"
+        )
+        self.assertEqual(decoded[0], expected_text)
+
+    def test_dynamic_cache_batched(self):
+        tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf", padding_side="left")
+        tokenizer.pad_token = tokenizer.eos_token
+        model = AutoModelForCausalLM.from_pretrained(
+            "meta-llama/Llama-2-7b-hf", device_map="auto", torch_dtype=torch.float16
+        )
+        inputs = tokenizer(["A sequence: 1, 2, 3, 4, 5", "A sequence: A, B, C"], padding=True, return_tensors="pt").to(
+            model.device
+        )
+
+        gen_out = model.generate(**inputs, do_sample=False, max_new_tokens=10, past_key_values=DynamicCache())
+        decoded = tokenizer.batch_decode(gen_out, skip_special_tokens=True)
+        expected_text = ["A sequence: 1, 2, 3, 4, 5, 6, 7, 8,", "A sequence: A, B, C, D, E, F, G, H"]
+        self.assertListEqual(decoded, expected_text)
+
+    def test_dynamic_cache_beam_search(self):
+        tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf", padding_side="left")
+        model = AutoModelForCausalLM.from_pretrained(
+            "meta-llama/Llama-2-7b-hf", device_map="auto", torch_dtype=torch.float16
+        )
+
+        inputs = tokenizer(["The best color is"], return_tensors="pt").to(model.device)
+        gen_out = model.generate(
+            **inputs,
+            do_sample=False,
+            max_new_tokens=20,
+            num_beams=2,
+            num_return_sequences=2,
+        )
+        decoded = tokenizer.batch_decode(gen_out, skip_special_tokens=True)
+        expected_text = [
+            "The best color is the one that makes you feel good.\nThe best color is the one that makes you feel good",
+            "The best color is the one that suits you.\nThe best color is the one that suits you. The",
+        ]
+        self.assertListEqual(decoded, expected_text)
+
+    def test_hybrid_cache_n_sequences(self):
+        tokenizer = AutoTokenizer.from_pretrained("google/gemma-2-9b")
+        model = AutoModelForCausalLM.from_pretrained(
+            "google/gemma-2-9b",
+            device_map="auto",
+            torch_dtype=torch.bfloat16,
+            attn_implementation="eager",
+        )
+
+        inputs = tokenizer(["Hello I am doing"], return_tensors="pt").to(model.device)
+
+        gen_out = model.generate(
+            **inputs,
+            do_sample=False,
+            max_new_tokens=20,
+            num_return_sequences=2,
+            num_beams=2,
+        )
+        decoded = tokenizer.batch_decode(gen_out, skip_special_tokens=True)
+        expected_text = [
+            "Hello I am doing a project for my school and I am trying to make a program that will allow me to input a",
+            "Hello I am doing a project for my school and I am trying to make a program that will allow me to use a",
+        ]
+        self.assertListEqual(decoded, expected_text)
+
+    @require_non_xpu
+    @require_gptq
+    def test_sink_cache_hard(self):
+        tokenizer = AutoTokenizer.from_pretrained("TheBloke/LLaMa-7B-GPTQ")
+        model = AutoModelForCausalLM.from_pretrained("TheBloke/LLaMa-7B-GPTQ", device_map="auto")
+
+        inputs = tokenizer(["Vaswani et al. (2017) introduced the Transformers"], return_tensors="pt").to(model.device)
+
+        # Set up the SinkCache. Using a small window length to contain computational complexity. If this example is run
+        # without a SinkCache, the last few tokens are gibberish (ends in "of the of the of a of a of")
+        cache = SinkCache(window_length=508, num_sink_tokens=4)
+        gen_out = model.generate(**inputs, do_sample=False, max_new_tokens=3000, past_key_values=cache)
+        decoded = tokenizer.batch_decode(gen_out, skip_special_tokens=True)
+        self.assertTrue(decoded[0].endswith("to perform a variety of tasks. The Transformer is a neural network"))
+
+    def test_sink_cache_iterative_prompts(self):
+        """Tests that SinkCache supports more than one new token at once, when shifting the cache"""
+        tokenizer = AutoTokenizer.from_pretrained("HuggingFaceH4/zephyr-7b-beta")
+        model = AutoModelForCausalLM.from_pretrained(
+            "HuggingFaceH4/zephyr-7b-beta", device_map="auto", torch_dtype=torch.float16
+        )
+        prompt = (
+            "Compose an engaging travel blog post about a recent trip to Hawaii, highlighting cultural experiences "
+            "and must-see attractions."
+        )
+
+        # Prepare generation settings
+        cache = SinkCache(window_length=256, num_sink_tokens=4)
+        input_ids = torch.tensor([], device=model.device, dtype=torch.int)
+        for _ in range(3):
+            # Tokenize the prompt with the correct chat template
+            chat = [{"role": "user", "content": prompt}]
+            tokenized_chat = tokenizer.apply_chat_template(chat, return_tensors="pt", add_generation_prompt=True).to(
+                model.device
+            )
+            input_ids = torch.cat((input_ids, tokenized_chat), dim=1)
+
+            # Perform the generation
+            gen_out = model.generate(
+                input_ids, do_sample=False, max_new_tokens=100, past_key_values=cache, use_cache=True
+            )
+            input_ids = gen_out
+
+        # We went well beyond the cache length
+        self.assertTrue(input_ids.shape[1] > cache.get_max_cache_shape() * 1.5)
+
+        # And it still produces a coherent english
+        decoded = tokenizer.batch_decode(input_ids, skip_special_tokens=True)
+        last_output = (
+            "<|assistant|>\nAs the sun began to set over the Pacific Ocean, I found myself standing on the shores of "
+            "Waikiki Beach, my heart filled with awe and wonder. I had just returned from a two-week journey to the "
+            "beautiful island of Hawaii, and it had been an unforgettable experience filled with cultural experiences "
+            "and must-see attractions that left me breathless.\n\nOne of the most memorable experiences of my trip "
+            "was visiting the historic district of Honolulu. Here,"
+        )
+        self.assertTrue(decoded[0].endswith(last_output))
+
+    @require_torch_gpu
+    @parameterized.expand(
+        [
+            ("eager", "static"),
+            ("sdpa", "static"),
+        ]
+    )
+    def test_static_cache_greedy_decoding_pad_left(self, attn_implementation, cache_implementation):
+        EXPECTED_GENERATION = [
+            "The best color is the one that complements the skin tone of the",
+            "We should not undermind the issues at hand.\nWe should not undermind the issues",
+        ]
+
+        tokenizer = AutoTokenizer.from_pretrained(
+            "NousResearch/Llama-2-7b-chat-hf", padding_side="left", pad_token="<s>"
+        )
+        model = AutoModelForCausalLM.from_pretrained(
+            "NousResearch/Llama-2-7b-chat-hf",
+            torch_dtype=torch.bfloat16,
+            attn_implementation=attn_implementation,
+        ).to(torch_device)
+        inputs = tokenizer(
+            ["The best color is", "We should not undermind the issues at hand"], padding=True, return_tensors="pt"
+        ).to(model.device)
+
+        set_seed(0)
+        gen_out = model.generate(**inputs, do_sample=False, max_new_tokens=10)
+        decoded = tokenizer.batch_decode(gen_out, skip_special_tokens=True)
+        with self.subTest(f"{attn_implementation}, dynamic"):
+            self.assertListEqual(decoded, EXPECTED_GENERATION)
+
+        set_seed(0)
+        model.generation_config.cache_implementation = cache_implementation
+        gen_out = model.generate(**inputs, do_sample=False, max_new_tokens=10)
+        decoded = tokenizer.batch_decode(gen_out, skip_special_tokens=True)
+        with self.subTest(f"{attn_implementation}, static, eager"):
+            self.assertListEqual(decoded, EXPECTED_GENERATION)
+
+        set_seed(0)
+        model.forward = torch.compile(model.forward)
+        gen_out = model.generate(**inputs, do_sample=False, max_new_tokens=10)
+        decoded = tokenizer.batch_decode(gen_out, skip_special_tokens=True)
+        with self.subTest(f"{attn_implementation}, static, compiled"):
+            self.assertListEqual(decoded, EXPECTED_GENERATION)
+
+    @require_torch_gpu
+    @parameterized.expand(
+        [
+            ("eager", "static"),
+            ("sdpa", "static"),
+        ]
+    )
+    def test_static_cache_greedy_decoding_pad_right(self, attn_implementation, cache_implementation):
+        EXPECTED_GENERATION = [
+            "The best color isЋ the one that complements the skin tone of",
+            "We should not undermind the issues at hand.\nWe should not undermind the issues",
+        ]
+
+        tokenizer = AutoTokenizer.from_pretrained(
+            "NousResearch/Llama-2-7b-chat-hf", padding_side="right", pad_token="<s>"
+        )
+        model = AutoModelForCausalLM.from_pretrained(
+            "NousResearch/Llama-2-7b-chat-hf",
+            torch_dtype=torch.bfloat16,
+            attn_implementation=attn_implementation,
+        ).to(torch_device)
+        inputs = tokenizer(
+            ["The best color is", "We should not undermind the issues at hand"], padding=True, return_tensors="pt"
+        ).to(model.device)
+
+        set_seed(0)
+        gen_out = model.generate(**inputs, do_sample=False, max_new_tokens=10)
+        decoded = tokenizer.batch_decode(gen_out, skip_special_tokens=True)
+        with self.subTest(f"{attn_implementation}, dynamic"):
+            self.assertListEqual(decoded, EXPECTED_GENERATION)
+
+        set_seed(0)
+        model.generation_config.cache_implementation = cache_implementation
+        gen_out = model.generate(**inputs, do_sample=False, max_new_tokens=10)
+        decoded = tokenizer.batch_decode(gen_out, skip_special_tokens=True)
+        with self.subTest(f"{attn_implementation}, static, eager"):
+            self.assertListEqual(decoded, EXPECTED_GENERATION)
+
+    def test_dynamic_cache_extra_left_padding(self):
+        """Tests that adding extra left-padding does not affect the generation with the dynamic cache"""
+        EXPECTED_GENERATION = [
+            "The best color is the one that complements the skin tone of the",
+            "We should not undermind the issues at hand.\nWe should not undermind the issues",
+        ]
+
+        tokenizer = AutoTokenizer.from_pretrained(
+            "NousResearch/Llama-2-7b-chat-hf", padding_side="left", pad_token="<s>"
+        )
+        model = AutoModelForCausalLM.from_pretrained(
+            "NousResearch/Llama-2-7b-chat-hf",
+            torch_dtype=torch.bfloat16,
+        ).to(torch_device)
+        inputs = tokenizer(
+            ["The best color is", "We should not undermind the issues at hand"], padding=True, return_tensors="pt"
+        ).to(model.device)
+
+        gen_out = model.generate(**inputs, do_sample=False, max_new_tokens=10)
+        decoded = tokenizer.batch_decode(gen_out, skip_special_tokens=True)
+        self.assertListEqual(decoded, EXPECTED_GENERATION)
+
+        # Now with extra left-padding
+        inputs_expanded = tokenizer(
+            ["The best color is", "We should not undermind the issues at hand"],
+            padding=True,
+            return_tensors="pt",
+            pad_to_multiple_of=32,
+        ).to(model.device)
+        self.assertTrue(inputs.input_ids.shape[1] < inputs_expanded.input_ids.shape[1])
+        gen_out = model.generate(**inputs_expanded, do_sample=False, max_new_tokens=10)
+        decoded = tokenizer.batch_decode(gen_out, skip_special_tokens=True)
+        self.assertListEqual(decoded, EXPECTED_GENERATION)
+
+    @parameterized.expand(
+        [
+            "static",
+        ]
+    )
+    def test_static_cache_extra_left_padding(self, cache_implementation):
+        """Tests that adding extra left-padding does not affect the generation with the static cache"""
+        EXPECTED_GENERATION = [
+            "The best color is the one that complements the skin tone of the",
+            "We should not undermind the issues at hand.\nWe should not undermind the issues",
+        ]
+
+        tokenizer = AutoTokenizer.from_pretrained(
+            "NousResearch/Llama-2-7b-chat-hf", padding_side="left", pad_token="<s>"
+        )
+        model = AutoModelForCausalLM.from_pretrained(
+            "NousResearch/Llama-2-7b-chat-hf",
+            torch_dtype=torch.bfloat16,
+        ).to(torch_device)
+        inputs = tokenizer(
+            ["The best color is", "We should not undermind the issues at hand"], padding=True, return_tensors="pt"
+        ).to(model.device)
+
+        model.generation_config.cache_implementation = cache_implementation
+
+        gen_out = model.generate(**inputs, do_sample=False, max_new_tokens=10)
+        decoded = tokenizer.batch_decode(gen_out, skip_special_tokens=True)
+        self.assertListEqual(decoded, EXPECTED_GENERATION)
+
+        # Now with extra left-padding
+        inputs_expanded = tokenizer(
+            ["The best color is", "We should not undermind the issues at hand"],
+            padding=True,
+            return_tensors="pt",
+            pad_to_multiple_of=32,
+        ).to(model.device)
+        self.assertTrue(inputs.input_ids.shape[1] < inputs_expanded.input_ids.shape[1])
+        gen_out = model.generate(**inputs_expanded, do_sample=False, max_new_tokens=10)
+        decoded = tokenizer.batch_decode(gen_out, skip_special_tokens=True)
+        self.assertListEqual(decoded, EXPECTED_GENERATION)
+
+    @unittest.skip(reason="TODO @gante static cache's does not support beam search yet")
+    def test_static_cache_beam_search(self):
+        pass
+
+    @require_torch_accelerator
+    def test_offloaded_cache_equivalent_to_dynamic_cache(self):
+        """Tests that OffloadedCache produces the same result as the default DynamicCache"""
+        model_name = "microsoft/Phi-3-mini-4k-instruct"
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", torch_dtype=torch.float16)
+        device = model.device
+
+        if not is_torch_greater_or_equal("2.7", accept_dev=True) and device.type == "xpu":
+            self.skipTest(reason="This test requires torch >= 2.7 to run on xpu.")
+
+        input_text = "Fun fact:"
+        inputs = tokenizer(input_text, return_tensors="pt").to(device)
+        common = {
+            "num_beams": 4,
+            "num_beam_groups": 2,
+            "num_return_sequences": 4,
+            "diversity_penalty": 1.0,
+            "max_new_tokens": 20,
+            "early_stopping": True,
+        }
+        original = GenerationConfig(**common)
+        offloaded = GenerationConfig(cache_implementation="offloaded", **common)
+        original_outputs = model.generate(generation_config=original, **inputs)
+        offloaded_outputs = model.generate(generation_config=offloaded, **inputs)
+        for original_output, offloaded_output in zip(original_outputs, offloaded_outputs):
+            assert torch.all(original_output == offloaded_output).item()
+
+    @require_torch_accelerator
+    def test_offloaded_cache_uses_less_memory_than_dynamic_cache(self):
+        """Tests that OffloadedCache uses less memory than the default DynamicCache"""
+        model_name = "microsoft/Phi-3-mini-4k-instruct"
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", torch_dtype=torch.float16)
+        device = model.device
+
+        if not is_torch_greater_or_equal("2.7", accept_dev=True) and device.type == "xpu":
+            self.skipTest(reason="This test requires torch >= 2.7 to run on xpu.")
+
+        input_text = "Fun fact:"
+        inputs = tokenizer(input_text, return_tensors="pt").to(device)
+        common = {
+            "num_beams": 4,
+            "num_beam_groups": 2,
+            "num_return_sequences": 4,
+            "diversity_penalty": 1.0,
+            "max_new_tokens": 20,
+            "early_stopping": True,
+        }
+        original = GenerationConfig(**common)
+        offloaded = GenerationConfig(cache_implementation="offloaded", **common)
+
+        torch_accelerator_module = None
+        if device.type == "cuda":
+            torch_accelerator_module = torch.cuda
+        elif device.type == "xpu":
+            torch_accelerator_module = torch.xpu
+
+        torch_accelerator_module.reset_peak_memory_stats(device)
+        model.generate(generation_config=original, **inputs)
+        original_peak_memory = torch_accelerator_module.max_memory_allocated(device)
+        torch_accelerator_module.reset_peak_memory_stats(device)
+        model.generate(generation_config=offloaded, **inputs)
+        offloaded_peak_memory = torch_accelerator_module.max_memory_allocated(device)
+        print(f"original_peak_memory: {original_peak_memory}, offloaded_peak_memory: {offloaded_peak_memory}")
+        assert offloaded_peak_memory < original_peak_memory
+
+    @require_torch_gpu
+    def test_cache_copy(self):
+        model_name = "microsoft/Phi-3-mini-4k-instruct"
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        model = AutoModelForCausalLM.from_pretrained(model_name, device_map="cuda", torch_dtype=torch.bfloat16)
+
+        prompt_cache = StaticCache(
+            config=model.config, max_batch_size=1, max_cache_len=1024, device="cuda", dtype=torch.bfloat16
+        )
+
+        INITIAL_PROMPT = "You are a helpful assistant. "
+        inputs_initial_prompt = tokenizer(INITIAL_PROMPT, return_tensors="pt").to("cuda")
+        # This is the common prompt cached, we need to run forward without grad to be abel to copy
+        with torch.no_grad():
+            prompt_cache = model(**inputs_initial_prompt, past_key_values=prompt_cache).past_key_values
+
+        prompts = ["Help me to write a blogpost about travelling.", "What is the capital of France?"]
+        responses = []
+        for prompt in prompts:
+            new_inputs = tokenizer(INITIAL_PROMPT + prompt, return_tensors="pt").to("cuda")
+            past_key_values = copy.deepcopy(prompt_cache)
+            outputs = model.generate(**new_inputs, past_key_values=past_key_values, max_new_tokens=40)
+            response = tokenizer.batch_decode(outputs)[0]
+            responses.append(response)
+
+        EXPECTED_DECODED_TEXT = [
+            "You are a helpful assistant. Help me to write a blogpost about travelling.\n\nTraveling is an enriching experience that broadens our horizons and exposes us to new cultures, landscapes, and people. Whether it's a week",
+            'You are a helpful assistant. What is the capital of France?\n\n\n## Response:Paris is the capital of France.\n\n\n\n\n\n## Query:\n\nIn a detailed analysis, compare the economic impacts of the introduction of the'
+        ]  # fmt: skip
+        self.assertEqual(responses, EXPECTED_DECODED_TEXT)
+
+    @require_torch_multi_gpu
+    def test_data_parallel_dynamic_cache(self):
+        """
+        Tests that the dynamic cache works with nn.DataParallel. Under the hood, `DynamicCache` is rebuilt from
+        multiple `DynamicCache` in the gather step.
+        """
+
+        model_repo = "hf-internal-testing/tiny-random-MistralForCausalLM"
+        model = AutoModelForCausalLM.from_pretrained(model_repo).to(torch_device)
+        tokenizer = AutoTokenizer.from_pretrained(model_repo)
+
+        # w/o DP: batch_size = num_gpu
+        # w DP: batch_size = 1 (with num_gpus replicas)
+        num_gpus = get_gpu_count()
+        model_inputs = tokenizer(["foo bar"] * num_gpus, return_tensors="pt").to(model.device)
+
+        # w/o DP
+        no_parallelism_cache = model(**model_inputs).past_key_values
+        self.assertIsInstance(no_parallelism_cache, DynamicCache)
+
+        # w DP
+        model = torch.nn.DataParallel(model)
+        parallelism_cache = model(**model_inputs).past_key_values
+        self.assertIsInstance(parallelism_cache, DynamicCache)
+
+        # Check that the caches are the same
+        for layer_idx in range(len(no_parallelism_cache)):
+            for kv_idx in range(2):  # 0 = key, 1 = value
+                torch.testing.assert_close(
+                    actual=parallelism_cache[layer_idx][kv_idx], expected=no_parallelism_cache[layer_idx][kv_idx]
+                )
+
+    @require_torch_gpu
+    def test_static_cache_no_cuda_graph_skips(self):
+        """
+        Tests generating with static cache and compilation doesn't skip cuda graphs. Regression test for #36543.
+
+        (? We set `fullgraph=True`, which according to torch docs means it should raise an exception. Instead,
+        messages are being thrown to stderr?)
+        """
+        model_repo = "hf-internal-testing/tiny-random-MistralForCausalLM"
+        model = AutoModelForCausalLM.from_pretrained(model_repo).to(torch_device)
+        tokenizer = AutoTokenizer.from_pretrained(model_repo)
+        inputs = tokenizer(["foo bar"], return_tensors="pt").to(torch_device)
+
+        # on `main`, prior to #36543, this would send stderr messages about cuda graphs being skipped.
+        with CaptureStderr() as cap:
+            model.generate(**inputs, max_new_tokens=2, cache_implementation="static")
+        self.assertEqual(cap.err, "")
+
+    @require_torch_multi_gpu
+    def test_static_cache_multi_gpu(self):
+        """Regression test for #35164: static cache with multi-gpu"""
+
+        model_id = "google/gemma-2-2b-it"
+        tokenizer = AutoTokenizer.from_pretrained(model_id)
+
+        device_map = {"model.embed_tokens": 0, "model.norm": 1, "model.rotary_emb": 1, "lm_head": 0}
+        num_hidden_layers = 26
+        for i in range(num_hidden_layers):
+            device_map[f"model.layers.{i}"] = 0 if i < 13 else 1
+
+        model = AutoModelForCausalLM.from_pretrained(
+            model_id,
+            torch_dtype="bfloat16",
+            device_map=device_map,
+        )
+        inputs = tokenizer("Today is a beautiful day!", return_tensors="pt").to(0)
+        _ = model(**inputs)
+        _ = model.generate(**inputs, max_new_tokens=2, cache_implementation="hybrid")

docs/transformers/tests/utils/test_chat_template_utils.py

@@ -0,0 +1,501 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+from typing import Optional, Union
+
+from transformers.utils import DocstringParsingException, TypeHintParsingException, get_json_schema
+
+
+class JsonSchemaGeneratorTest(unittest.TestCase):
+    def test_simple_function(self):
+        def fn(x: int):
+            """
+            Test function
+
+            Args:
+                 x: The input
+            """
+            return x
+
+        schema = get_json_schema(fn)
+        expected_schema = {
+            "name": "fn",
+            "description": "Test function",
+            "parameters": {
+                "type": "object",
+                "properties": {"x": {"type": "integer", "description": "The input"}},
+                "required": ["x"],
+            },
+        }
+        self.assertEqual(schema["function"], expected_schema)
+
+    def test_no_arguments(self):
+        def fn():
+            """
+            Test function
+            """
+            return True
+
+        schema = get_json_schema(fn)
+        expected_schema = {
+            "name": "fn",
+            "description": "Test function",
+            "parameters": {"type": "object", "properties": {}},
+        }
+        self.assertEqual(schema["function"], expected_schema)
+
+    def test_union(self):
+        def fn(x: Union[int, float]):
+            """
+            Test function
+
+            Args:
+                x: The input
+            """
+            return x
+
+        schema = get_json_schema(fn)
+        expected_schema = {
+            "name": "fn",
+            "description": "Test function",
+            "parameters": {
+                "type": "object",
+                "properties": {"x": {"type": ["integer", "number"], "description": "The input"}},
+                "required": ["x"],
+            },
+        }
+        self.assertEqual(schema["function"], expected_schema)
+
+    def test_optional(self):
+        def fn(x: Optional[int]):
+            """
+            Test function
+
+            Args:
+                x: The input
+            """
+            return x
+
+        schema = get_json_schema(fn)
+        expected_schema = {
+            "name": "fn",
+            "description": "Test function",
+            "parameters": {
+                "type": "object",
+                "properties": {"x": {"type": "integer", "description": "The input", "nullable": True}},
+                "required": ["x"],
+            },
+        }
+        self.assertEqual(schema["function"], expected_schema)
+
+    def test_default_arg(self):
+        def fn(x: int = 42):
+            """
+            Test function
+
+            Args:
+                 x: The input
+            """
+            return x
+
+        schema = get_json_schema(fn)
+        expected_schema = {
+            "name": "fn",
+            "description": "Test function",
+            "parameters": {"type": "object", "properties": {"x": {"type": "integer", "description": "The input"}}},
+        }
+        self.assertEqual(schema["function"], expected_schema)
+
+    def test_nested_list(self):
+        def fn(x: list[list[Union[str, int]]]):
+            """
+            Test function
+
+            Args:
+                x: The input
+            """
+            return x
+
+        schema = get_json_schema(fn)
+        expected_schema = {
+            "name": "fn",
+            "description": "Test function",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "x": {
+                        "type": "array",
+                        "items": {"type": "array", "items": {"type": ["integer", "string"]}},
+                        "description": "The input",
+                    }
+                },
+                "required": ["x"],
+            },
+        }
+        self.assertEqual(schema["function"], expected_schema)
+
+    def test_multiple_arguments(self):
+        def fn(x: int, y: str):
+            """
+            Test function
+
+            Args:
+                x: The input
+                y: Also the input
+            """
+            return x
+
+        schema = get_json_schema(fn)
+        expected_schema = {
+            "name": "fn",
+            "description": "Test function",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "x": {"type": "integer", "description": "The input"},
+                    "y": {"type": "string", "description": "Also the input"},
+                },
+                "required": ["x", "y"],
+            },
+        }
+        self.assertEqual(schema["function"], expected_schema)
+
+    def test_multiple_complex_arguments(self):
+        def fn(x: list[Union[int, float]], y: Optional[Union[int, str]] = None):
+            """
+            Test function
+
+            Args:
+                x: The input
+                y: Also the input
+            """
+            return x
+
+        schema = get_json_schema(fn)
+        expected_schema = {
+            "name": "fn",
+            "description": "Test function",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "x": {"type": "array", "items": {"type": ["integer", "number"]}, "description": "The input"},
+                    "y": {
+                        "type": ["integer", "string"],
+                        "nullable": True,
+                        "description": "Also the input",
+                    },
+                },
+                "required": ["x"],
+            },
+        }
+        self.assertEqual(schema["function"], expected_schema)
+
+    def test_missing_docstring(self):
+        def fn(x: int):
+            return x
+
+        with self.assertRaises(DocstringParsingException):
+            get_json_schema(fn)
+
+    def test_missing_param_docstring(self):
+        def fn(x: int):
+            """
+            Test function
+            """
+            return x
+
+        with self.assertRaises(DocstringParsingException):
+            get_json_schema(fn)
+
+    def test_missing_type_hint(self):
+        def fn(x):
+            """
+            Test function
+
+            Args:
+                 x: The input
+            """
+            return x
+
+        with self.assertRaises(TypeHintParsingException):
+            get_json_schema(fn)
+
+    def test_return_value(self):
+        def fn(x: int) -> int:
+            """
+            Test function
+
+            Args:
+                x: The input
+            """
+            return x
+
+        schema = get_json_schema(fn)
+        expected_schema = {
+            "name": "fn",
+            "description": "Test function",
+            "parameters": {
+                "type": "object",
+                "properties": {"x": {"type": "integer", "description": "The input"}},
+                "required": ["x"],
+            },
+            "return": {"type": "integer"},
+        }
+        self.assertEqual(schema["function"], expected_schema)
+
+    def test_return_value_docstring(self):
+        def fn(x: int) -> int:
+            """
+            Test function
+
+            Args:
+                x: The input
+
+
+            Returns:
+                The output
+            """
+            return x
+
+        schema = get_json_schema(fn)
+        expected_schema = {
+            "name": "fn",
+            "description": "Test function",
+            "parameters": {
+                "type": "object",
+                "properties": {"x": {"type": "integer", "description": "The input"}},
+                "required": ["x"],
+            },
+            "return": {"type": "integer", "description": "The output"},
+        }
+        self.assertEqual(schema["function"], expected_schema)
+
+    def test_tuple(self):
+        def fn(x: tuple[int, str]):
+            """
+            Test function
+
+            Args:
+                x: The input
+
+
+            Returns:
+                The output
+            """
+            return x
+
+        schema = get_json_schema(fn)
+        expected_schema = {
+            "name": "fn",
+            "description": "Test function",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "x": {
+                        "type": "array",
+                        "prefixItems": [{"type": "integer"}, {"type": "string"}],
+                        "description": "The input",
+                    }
+                },
+                "required": ["x"],
+            },
+        }
+        self.assertEqual(schema["function"], expected_schema)
+
+    def test_single_element_tuple_fails(self):
+        def fn(x: tuple[int]):
+            """
+            Test function
+
+            Args:
+                x: The input
+
+
+            Returns:
+                The output
+            """
+            return x
+
+        # Single-element tuples should just be the type itself, or List[type] for variable-length inputs
+        with self.assertRaises(TypeHintParsingException):
+            get_json_schema(fn)
+
+    def test_ellipsis_type_fails(self):
+        def fn(x: tuple[int, ...]):
+            """
+            Test function
+
+            Args:
+                x: The input
+
+
+            Returns:
+                The output
+            """
+            return x
+
+        # Variable length inputs should be specified with List[type], not Tuple[type, ...]
+        with self.assertRaises(TypeHintParsingException):
+            get_json_schema(fn)
+
+    def test_enum_extraction(self):
+        def fn(temperature_format: str):
+            """
+            Test function
+
+            Args:
+                temperature_format: The temperature format to use (Choices: ["celsius", "fahrenheit"])
+
+
+            Returns:
+                The temperature
+            """
+            return -40.0
+
+        # Let's see if that gets correctly parsed as an enum
+        schema = get_json_schema(fn)
+        expected_schema = {
+            "name": "fn",
+            "description": "Test function",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "temperature_format": {
+                        "type": "string",
+                        "enum": ["celsius", "fahrenheit"],
+                        "description": "The temperature format to use",
+                    }
+                },
+                "required": ["temperature_format"],
+            },
+        }
+
+        self.assertEqual(schema["function"], expected_schema)
+
+    def test_multiline_docstring_with_types(self):
+        def fn(x: int, y: int):
+            """
+            Test function
+
+            Args:
+                x: The first input
+
+                y: The second input. This is a longer description
+                   that spans multiple lines with indentation and stuff.
+
+            Returns:
+                God knows what
+            """
+            pass
+
+        schema = get_json_schema(fn)
+        expected_schema = {
+            "name": "fn",
+            "description": "Test function",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "x": {"type": "integer", "description": "The first input"},
+                    "y": {
+                        "type": "integer",
+                        "description": "The second input. This is a longer description that spans multiple lines with indentation and stuff.",
+                    },
+                },
+                "required": ["x", "y"],
+            },
+        }
+
+        self.assertEqual(schema["function"], expected_schema)
+
+    def test_return_none(self):
+        def fn(x: int) -> None:
+            """
+            Test function
+
+            Args:
+                x: The first input
+            """
+            pass
+
+        schema = get_json_schema(fn)
+        expected_schema = {
+            "name": "fn",
+            "description": "Test function",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "x": {"type": "integer", "description": "The first input"},
+                },
+                "required": ["x"],
+            },
+            "return": {"type": "null"},
+        }
+        self.assertEqual(schema["function"], expected_schema)
+
+    def test_everything_all_at_once(self):
+        def fn(
+            x: str, y: Optional[list[Union[str, int]]], z: tuple[Union[str, int], str] = (42, "hello")
+        ) -> tuple[int, str]:
+            """
+            Test function with multiple args, and docstring args that we have to strip out.
+
+            Args:
+                x: The first input. It's got a big multiline
+                   description and also contains
+                   (choices: ["a", "b", "c"])
+
+                y: The second input. It's a big list with a single-line description.
+
+                z: The third input. It's some kind of tuple with a default arg.
+
+            Returns:
+                The output. The return description is also a big multiline
+                description that spans multiple lines.
+            """
+            pass
+
+        schema = get_json_schema(fn)
+        expected_schema = {
+            "name": "fn",
+            "description": "Test function with multiple args, and docstring args that we have to strip out.",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "x": {
+                        "type": "string",
+                        "enum": ["a", "b", "c"],
+                        "description": "The first input. It's got a big multiline description and also contains",
+                    },
+                    "y": {
+                        "type": "array",
+                        "items": {"type": ["integer", "string"]},
+                        "nullable": True,
+                        "description": "The second input. It's a big list with a single-line description.",
+                    },
+                    "z": {
+                        "type": "array",
+                        "prefixItems": [{"type": ["integer", "string"]}, {"type": "string"}],
+                        "description": "The third input. It's some kind of tuple with a default arg.",
+                    },
+                },
+                "required": ["x", "y"],
+            },
+            "return": {
+                "type": "array",
+                "prefixItems": [{"type": "integer"}, {"type": "string"}],
+                "description": "The output. The return description is also a big multiline\n    description that spans multiple lines.",
+            },
+        }
+        self.assertEqual(schema["function"], expected_schema)

docs/transformers/tests/utils/test_cli.py

@@ -0,0 +1,77 @@
+# Copyright 2019-present, the HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import shutil
+import unittest
+from unittest.mock import patch
+
+from transformers.testing_utils import CaptureStd, require_torch
+
+
+class CLITest(unittest.TestCase):
+    @patch("sys.argv", ["fakeprogrampath", "env"])
+    def test_cli_env(self):
+        # test transformers-cli env
+        import transformers.commands.transformers_cli
+
+        with CaptureStd() as cs:
+            transformers.commands.transformers_cli.main()
+        self.assertIn("Python version", cs.out)
+        self.assertIn("Platform", cs.out)
+        self.assertIn("Using distributed or parallel set-up in script?", cs.out)
+
+    @require_torch
+    @patch("sys.argv", ["fakeprogrampath", "download", "hf-internal-testing/tiny-random-gptj", "--cache-dir", "/tmp"])
+    def test_cli_download(self):
+        import transformers.commands.transformers_cli
+
+        # # remove any previously downloaded model to start clean
+        shutil.rmtree("/tmp/models--hf-internal-testing--tiny-random-gptj", ignore_errors=True)
+
+        # run the command
+        transformers.commands.transformers_cli.main()
+
+        # check if the model files are downloaded correctly on /tmp/models--hf-internal-testing--tiny-random-gptj
+        self.assertTrue(os.path.exists("/tmp/models--hf-internal-testing--tiny-random-gptj/blobs"))
+        self.assertTrue(os.path.exists("/tmp/models--hf-internal-testing--tiny-random-gptj/refs"))
+        self.assertTrue(os.path.exists("/tmp/models--hf-internal-testing--tiny-random-gptj/snapshots"))
+
+    @require_torch
+    @patch(
+        "sys.argv",
+        [
+            "fakeprogrampath",
+            "download",
+            "hf-internal-testing/test_dynamic_model_with_tokenizer",
+            "--trust-remote-code",
+            "--cache-dir",
+            "/tmp",
+        ],
+    )
+    def test_cli_download_trust_remote(self):
+        import transformers.commands.transformers_cli
+
+        # # remove any previously downloaded model to start clean
+        shutil.rmtree("/tmp/models--hf-internal-testing--test_dynamic_model_with_tokenizer", ignore_errors=True)
+
+        # run the command
+        transformers.commands.transformers_cli.main()
+
+        # check if the model files are downloaded correctly on /tmp/models--hf-internal-testing--test_dynamic_model_with_tokenizer
+        self.assertTrue(os.path.exists("/tmp/models--hf-internal-testing--test_dynamic_model_with_tokenizer/blobs"))
+        self.assertTrue(os.path.exists("/tmp/models--hf-internal-testing--test_dynamic_model_with_tokenizer/refs"))
+        self.assertTrue(
+            os.path.exists("/tmp/models--hf-internal-testing--test_dynamic_model_with_tokenizer/snapshots")
+        )

docs/transformers/tests/utils/test_configuration_utils.py

@@ -0,0 +1,302 @@
+# Copyright 2019 HuggingFace Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import json
+import os
+import shutil
+import sys
+import tempfile
+import unittest
+import unittest.mock as mock
+import warnings
+from pathlib import Path
+
+from huggingface_hub import HfFolder
+from requests.exceptions import HTTPError
+
+from transformers import AutoConfig, BertConfig, GPT2Config
+from transformers.configuration_utils import PretrainedConfig
+from transformers.testing_utils import TOKEN, TemporaryHubRepo, is_staging_test
+
+
+sys.path.append(str(Path(__file__).parent.parent.parent / "utils"))
+
+from test_module.custom_configuration import CustomConfig  # noqa E402
+
+
+config_common_kwargs = {
+    "return_dict": False,
+    "output_hidden_states": True,
+    "output_attentions": True,
+    "torchscript": True,
+    "torch_dtype": "float16",
+    "use_bfloat16": True,
+    "tf_legacy_loss": True,
+    "pruned_heads": {"a": 1},
+    "tie_word_embeddings": False,
+    "is_decoder": True,
+    "cross_attention_hidden_size": 128,
+    "add_cross_attention": True,
+    "tie_encoder_decoder": True,
+    "max_length": 50,
+    "min_length": 3,
+    "do_sample": True,
+    "early_stopping": True,
+    "num_beams": 3,
+    "num_beam_groups": 3,
+    "diversity_penalty": 0.5,
+    "temperature": 2.0,
+    "top_k": 10,
+    "top_p": 0.7,
+    "typical_p": 0.2,
+    "repetition_penalty": 0.8,
+    "length_penalty": 0.8,
+    "no_repeat_ngram_size": 5,
+    "encoder_no_repeat_ngram_size": 5,
+    "bad_words_ids": [1, 2, 3],
+    "num_return_sequences": 3,
+    "chunk_size_feed_forward": 5,
+    "output_scores": True,
+    "return_dict_in_generate": True,
+    "forced_bos_token_id": 2,
+    "forced_eos_token_id": 3,
+    "remove_invalid_values": True,
+    "architectures": ["BertModel"],
+    "finetuning_task": "translation",
+    "id2label": {0: "label"},
+    "label2id": {"label": "0"},
+    "tokenizer_class": "BertTokenizerFast",
+    "prefix": "prefix",
+    "bos_token_id": 6,
+    "pad_token_id": 7,
+    "eos_token_id": 8,
+    "sep_token_id": 9,
+    "decoder_start_token_id": 10,
+    "exponential_decay_length_penalty": (5, 1.01),
+    "suppress_tokens": [0, 1],
+    "begin_suppress_tokens": 2,
+    "task_specific_params": {"translation": "some_params"},
+    "problem_type": "regression",
+}
+
+
+@is_staging_test
+class ConfigPushToHubTester(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls._token = TOKEN
+        HfFolder.save_token(TOKEN)
+
+    def test_push_to_hub(self):
+        with TemporaryHubRepo(token=self._token) as tmp_repo:
+            config = BertConfig(
+                vocab_size=99, hidden_size=32, num_hidden_layers=5, num_attention_heads=4, intermediate_size=37
+            )
+            config.push_to_hub(tmp_repo.repo_id, token=self._token)
+
+            new_config = BertConfig.from_pretrained(tmp_repo.repo_id)
+            for k, v in config.to_dict().items():
+                if k != "transformers_version":
+                    self.assertEqual(v, getattr(new_config, k))
+
+    def test_push_to_hub_via_save_pretrained(self):
+        with TemporaryHubRepo(token=self._token) as tmp_repo:
+            config = BertConfig(
+                vocab_size=99, hidden_size=32, num_hidden_layers=5, num_attention_heads=4, intermediate_size=37
+            )
+            # Push to hub via save_pretrained
+            with tempfile.TemporaryDirectory() as tmp_dir:
+                config.save_pretrained(tmp_dir, repo_id=tmp_repo.repo_id, push_to_hub=True, token=self._token)
+
+            new_config = BertConfig.from_pretrained(tmp_repo.repo_id)
+            for k, v in config.to_dict().items():
+                if k != "transformers_version":
+                    self.assertEqual(v, getattr(new_config, k))
+
+    def test_push_to_hub_in_organization(self):
+        with TemporaryHubRepo(namespace="valid_org", token=self._token) as tmp_repo:
+            config = BertConfig(
+                vocab_size=99, hidden_size=32, num_hidden_layers=5, num_attention_heads=4, intermediate_size=37
+            )
+            config.push_to_hub(tmp_repo.repo_id, token=self._token)
+
+            new_config = BertConfig.from_pretrained(tmp_repo.repo_id)
+            for k, v in config.to_dict().items():
+                if k != "transformers_version":
+                    self.assertEqual(v, getattr(new_config, k))
+
+    def test_push_to_hub_in_organization_via_save_pretrained(self):
+        with TemporaryHubRepo(namespace="valid_org", token=self._token) as tmp_repo:
+            config = BertConfig(
+                vocab_size=99, hidden_size=32, num_hidden_layers=5, num_attention_heads=4, intermediate_size=37
+            )
+            # Push to hub via save_pretrained
+            with tempfile.TemporaryDirectory() as tmp_dir:
+                config.save_pretrained(tmp_dir, repo_id=tmp_repo.repo_id, push_to_hub=True, token=self._token)
+
+            new_config = BertConfig.from_pretrained(tmp_repo.repo_id)
+            for k, v in config.to_dict().items():
+                if k != "transformers_version":
+                    self.assertEqual(v, getattr(new_config, k))
+
+    def test_push_to_hub_dynamic_config(self):
+        with TemporaryHubRepo(token=self._token) as tmp_repo:
+            CustomConfig.register_for_auto_class()
+            config = CustomConfig(attribute=42)
+
+            config.push_to_hub(tmp_repo.repo_id, token=self._token)
+
+            # This has added the proper auto_map field to the config
+            self.assertDictEqual(config.auto_map, {"AutoConfig": "custom_configuration.CustomConfig"})
+
+            new_config = AutoConfig.from_pretrained(tmp_repo.repo_id, trust_remote_code=True)
+            # Can't make an isinstance check because the new_config is from the FakeConfig class of a dynamic module
+            self.assertEqual(new_config.__class__.__name__, "CustomConfig")
+            self.assertEqual(new_config.attribute, 42)
+
+
+class ConfigTestUtils(unittest.TestCase):
+    def test_config_from_string(self):
+        c = GPT2Config()
+
+        # attempt to modify each of int/float/bool/str config records and verify they were updated
+        n_embd = c.n_embd + 1  # int
+        resid_pdrop = c.resid_pdrop + 1.0  # float
+        scale_attn_weights = not c.scale_attn_weights  # bool
+        summary_type = c.summary_type + "foo"  # str
+        c.update_from_string(
+            f"n_embd={n_embd},resid_pdrop={resid_pdrop},scale_attn_weights={scale_attn_weights},summary_type={summary_type}"
+        )
+        self.assertEqual(n_embd, c.n_embd, "mismatch for key: n_embd")
+        self.assertEqual(resid_pdrop, c.resid_pdrop, "mismatch for key: resid_pdrop")
+        self.assertEqual(scale_attn_weights, c.scale_attn_weights, "mismatch for key: scale_attn_weights")
+        self.assertEqual(summary_type, c.summary_type, "mismatch for key: summary_type")
+
+    def test_config_common_kwargs_is_complete(self):
+        base_config = PretrainedConfig()
+        missing_keys = [key for key in base_config.__dict__ if key not in config_common_kwargs]
+        # If this part of the test fails, you have arguments to addin config_common_kwargs above.
+        self.assertListEqual(
+            missing_keys,
+            [
+                "is_encoder_decoder",
+                "_name_or_path",
+                "_commit_hash",
+                "_attn_implementation_internal",
+                "_attn_implementation_autoset",
+                "transformers_version",
+            ],
+        )
+        keys_with_defaults = [key for key, value in config_common_kwargs.items() if value == getattr(base_config, key)]
+        if len(keys_with_defaults) > 0:
+            raise ValueError(
+                "The following keys are set with the default values in"
+                " `test_configuration_common.config_common_kwargs` pick another value for them:"
+                f" {', '.join(keys_with_defaults)}."
+            )
+
+    def test_nested_config_load_from_dict(self):
+        config = AutoConfig.from_pretrained(
+            "hf-internal-testing/tiny-random-CLIPModel", text_config={"num_hidden_layers": 2}
+        )
+        self.assertNotIsInstance(config.text_config, dict)
+        self.assertEqual(config.text_config.__class__.__name__, "CLIPTextConfig")
+
+    def test_from_pretrained_subfolder(self):
+        config = BertConfig.from_pretrained("hf-internal-testing/tiny-random-bert-subfolder")
+        self.assertIsNotNone(config)
+
+        config = BertConfig.from_pretrained("hf-internal-testing/tiny-random-bert-subfolder", subfolder="bert")
+        self.assertIsNotNone(config)
+
+    def test_cached_files_are_used_when_internet_is_down(self):
+        # A mock response for an HTTP head request to emulate server down
+        response_mock = mock.Mock()
+        response_mock.status_code = 500
+        response_mock.headers = {}
+        response_mock.raise_for_status.side_effect = HTTPError
+        response_mock.json.return_value = {}
+
+        # Download this model to make sure it's in the cache.
+        _ = BertConfig.from_pretrained("hf-internal-testing/tiny-random-bert")
+
+        # Under the mock environment we get a 500 error when trying to reach the model.
+        with mock.patch("requests.Session.request", return_value=response_mock) as mock_head:
+            _ = BertConfig.from_pretrained("hf-internal-testing/tiny-random-bert")
+            # This check we did call the fake head request
+            mock_head.assert_called()
+
+    def test_local_versioning(self):
+        configuration = AutoConfig.from_pretrained("google-bert/bert-base-cased")
+        configuration.configuration_files = ["config.4.0.0.json"]
+
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            configuration.save_pretrained(tmp_dir)
+            configuration.hidden_size = 2
+            json.dump(configuration.to_dict(), open(os.path.join(tmp_dir, "config.4.0.0.json"), "w"))
+
+            # This should pick the new configuration file as the version of Transformers is > 4.0.0
+            new_configuration = AutoConfig.from_pretrained(tmp_dir)
+            self.assertEqual(new_configuration.hidden_size, 2)
+
+            # Will need to be adjusted if we reach v42 and this test is still here.
+            # Should pick the old configuration file as the version of Transformers is < 4.42.0
+            configuration.configuration_files = ["config.42.0.0.json"]
+            configuration.hidden_size = 768
+            configuration.save_pretrained(tmp_dir)
+            shutil.move(os.path.join(tmp_dir, "config.4.0.0.json"), os.path.join(tmp_dir, "config.42.0.0.json"))
+            new_configuration = AutoConfig.from_pretrained(tmp_dir)
+            self.assertEqual(new_configuration.hidden_size, 768)
+
+    def test_repo_versioning_before(self):
+        # This repo has two configuration files, one for v4.0.0 and above with a different hidden size.
+        repo = "hf-internal-testing/test-two-configs"
+
+        import transformers as new_transformers
+
+        new_transformers.configuration_utils.__version__ = "v4.0.0"
+        new_configuration, kwargs = new_transformers.models.auto.AutoConfig.from_pretrained(
+            repo, return_unused_kwargs=True
+        )
+        self.assertEqual(new_configuration.hidden_size, 2)
+        # This checks `_configuration_file` ia not kept in the kwargs by mistake.
+        self.assertDictEqual(kwargs, {})
+
+        # Testing an older version by monkey-patching the version in the module it's used.
+        import transformers as old_transformers
+
+        old_transformers.configuration_utils.__version__ = "v3.0.0"
+        old_configuration = old_transformers.models.auto.AutoConfig.from_pretrained(repo)
+        self.assertEqual(old_configuration.hidden_size, 768)
+
+    def test_saving_config_with_custom_generation_kwargs_raises_warning(self):
+        config = BertConfig(min_length=3)  # `min_length = 3` is a non-default generation kwarg
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            with self.assertWarns(UserWarning) as cm:
+                config.save_pretrained(tmp_dir)
+            self.assertIn("min_length", str(cm.warning))
+
+    def test_get_non_default_generation_parameters(self):
+        config = BertConfig()
+        self.assertFalse(len(config._get_non_default_generation_parameters()) > 0)
+        config = BertConfig(min_length=3)
+        self.assertTrue(len(config._get_non_default_generation_parameters()) > 0)
+        config = BertConfig(min_length=0)  # `min_length = 0` is a default generation kwarg
+        self.assertFalse(len(config._get_non_default_generation_parameters()) > 0)
+
+    def test_loading_config_do_not_raise_future_warnings(self):
+        """Regression test for https://github.com/huggingface/transformers/issues/31002."""
+        # Loading config should not raise a FutureWarning. It was the case before.
+        with warnings.catch_warnings():
+            warnings.simplefilter("error")
+            PretrainedConfig.from_pretrained("bert-base-uncased")

docs/transformers/tests/utils/test_convert_slow_tokenizer.py

@@ -0,0 +1,35 @@
+import unittest
+import warnings
+from dataclasses import dataclass
+
+from transformers.convert_slow_tokenizer import SpmConverter
+from transformers.testing_utils import get_tests_dir
+
+
+@dataclass
+class FakeOriginalTokenizer:
+    vocab_file: str
+
+
+class ConvertSlowTokenizerTest(unittest.TestCase):
+    def test_spm_converter_bytefallback_warning(self):
+        spm_model_file_without_bytefallback = get_tests_dir("fixtures/test_sentencepiece.model")
+        spm_model_file_with_bytefallback = get_tests_dir("fixtures/test_sentencepiece_with_bytefallback.model")
+
+        original_tokenizer_without_bytefallback = FakeOriginalTokenizer(vocab_file=spm_model_file_without_bytefallback)
+
+        with warnings.catch_warnings(record=True) as w:
+            _ = SpmConverter(original_tokenizer_without_bytefallback)
+        self.assertEqual(len(w), 0)
+
+        original_tokenizer_with_bytefallback = FakeOriginalTokenizer(vocab_file=spm_model_file_with_bytefallback)
+
+        with warnings.catch_warnings(record=True) as w:
+            _ = SpmConverter(original_tokenizer_with_bytefallback)
+        self.assertEqual(len(w), 1)
+
+        self.assertIn(
+            "The sentencepiece tokenizer that you are converting to a fast tokenizer uses the byte fallback option"
+            " which is not implemented in the fast tokenizers.",
+            str(w[0].message),
+        )

docs/transformers/tests/utils/test_deprecation.py

@@ -0,0 +1,195 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+import warnings
+
+from parameterized import parameterized
+
+from transformers import __version__, is_torch_available
+from transformers.testing_utils import require_torch_gpu
+from transformers.utils.deprecation import deprecate_kwarg
+
+
+if is_torch_available():
+    import torch
+
+
+INFINITE_VERSION = "9999.0.0"
+
+
+class DeprecationDecoratorTester(unittest.TestCase):
+    def test_rename_kwarg(self):
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore")
+
+            @deprecate_kwarg("deprecated_name", new_name="new_name", version=INFINITE_VERSION)
+            def dummy_function(new_name=None, other_name=None):
+                return new_name, other_name
+
+            # Test keyword argument is renamed
+            value, other_value = dummy_function(deprecated_name="old_value")
+            self.assertEqual(value, "old_value")
+            self.assertIsNone(other_value)
+
+            # Test deprecated keyword argument not passed
+            value, other_value = dummy_function(new_name="new_value")
+            self.assertEqual(value, "new_value")
+            self.assertIsNone(other_value)
+
+            # Test other keyword argument
+            value, other_value = dummy_function(other_name="other_value")
+            self.assertIsNone(value)
+            self.assertEqual(other_value, "other_value")
+
+            # Test deprecated and new args are passed, the new one should be returned
+            value, other_value = dummy_function(deprecated_name="old_value", new_name="new_value")
+            self.assertEqual(value, "new_value")
+            self.assertIsNone(other_value)
+
+    def test_rename_multiple_kwargs(self):
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore")
+
+            @deprecate_kwarg("deprecated_name1", new_name="new_name1", version=INFINITE_VERSION)
+            @deprecate_kwarg("deprecated_name2", new_name="new_name2", version=INFINITE_VERSION)
+            def dummy_function(new_name1=None, new_name2=None, other_name=None):
+                return new_name1, new_name2, other_name
+
+            # Test keyword argument is renamed
+            value1, value2, other_value = dummy_function(deprecated_name1="old_value1", deprecated_name2="old_value2")
+            self.assertEqual(value1, "old_value1")
+            self.assertEqual(value2, "old_value2")
+            self.assertIsNone(other_value)
+
+            # Test deprecated keyword argument is not passed
+            value1, value2, other_value = dummy_function(new_name1="new_value1", new_name2="new_value2")
+            self.assertEqual(value1, "new_value1")
+            self.assertEqual(value2, "new_value2")
+            self.assertIsNone(other_value)
+
+            # Test other keyword argument is passed and correctly returned
+            value1, value2, other_value = dummy_function(other_name="other_value")
+            self.assertIsNone(value1)
+            self.assertIsNone(value2)
+            self.assertEqual(other_value, "other_value")
+
+    def test_warnings(self):
+        # Test warning is raised for future version
+        @deprecate_kwarg("deprecated_name", new_name="new_name", version=INFINITE_VERSION)
+        def dummy_function(new_name=None, other_name=None):
+            return new_name, other_name
+
+        with self.assertWarns(FutureWarning):
+            dummy_function(deprecated_name="old_value")
+
+        # Test warning is not raised for past version, but arg is still renamed
+        @deprecate_kwarg("deprecated_name", new_name="new_name", version="0.0.0")
+        def dummy_function(new_name=None, other_name=None):
+            return new_name, other_name
+
+        with warnings.catch_warnings(record=True) as raised_warnings:
+            warnings.simplefilter("always")
+
+            value, other_value = dummy_function(deprecated_name="old_value")
+
+            self.assertEqual(value, "old_value")
+            self.assertIsNone(other_value)
+            self.assertEqual(len(raised_warnings), 0, f"Warning raised: {[w.message for w in raised_warnings]}")
+
+        # Test warning is raised for future version if warn_if_greater_or_equal_version is set
+        @deprecate_kwarg("deprecated_name", version="0.0.0", warn_if_greater_or_equal_version=True)
+        def dummy_function(deprecated_name=None):
+            return deprecated_name
+
+        with self.assertWarns(FutureWarning):
+            value = dummy_function(deprecated_name="deprecated_value")
+        self.assertEqual(value, "deprecated_value")
+
+        # Test arg is not renamed if new_name is not specified, but warning is raised
+        @deprecate_kwarg("deprecated_name", version=INFINITE_VERSION)
+        def dummy_function(deprecated_name=None):
+            return deprecated_name
+
+        with self.assertWarns(FutureWarning):
+            value = dummy_function(deprecated_name="deprecated_value")
+        self.assertEqual(value, "deprecated_value")
+
+    def test_raises(self):
+        # Test if deprecated name and new name are both passed and raise_if_both_names is set -> raise error
+        @deprecate_kwarg("deprecated_name", new_name="new_name", version=INFINITE_VERSION, raise_if_both_names=True)
+        def dummy_function(new_name=None, other_name=None):
+            return new_name, other_name
+
+        with self.assertRaises(ValueError):
+            dummy_function(deprecated_name="old_value", new_name="new_value")
+
+        # Test for current version == deprecation version
+        @deprecate_kwarg("deprecated_name", version=__version__, raise_if_greater_or_equal_version=True)
+        def dummy_function(deprecated_name=None):
+            return deprecated_name
+
+        with self.assertRaises(ValueError):
+            dummy_function(deprecated_name="old_value")
+
+        # Test for current version > deprecation version
+        @deprecate_kwarg("deprecated_name", version="0.0.0", raise_if_greater_or_equal_version=True)
+        def dummy_function(deprecated_name=None):
+            return deprecated_name
+
+        with self.assertRaises(ValueError):
+            dummy_function(deprecated_name="old_value")
+
+    def test_additional_message(self):
+        # Test additional message is added to the warning
+        @deprecate_kwarg("deprecated_name", version=INFINITE_VERSION, additional_message="Additional message")
+        def dummy_function(deprecated_name=None):
+            return deprecated_name
+
+        with warnings.catch_warnings(record=True) as raised_warnings:
+            warnings.simplefilter("always")
+            dummy_function(deprecated_name="old_value")
+
+            self.assertTrue("Additional message" in str(raised_warnings[0].message))
+
+    @parameterized.expand(["0.0.0", __version__, INFINITE_VERSION])
+    def test_warning_for_both_names(self, version):
+        # We should raise warning if both names are passed for any specified version
+        @deprecate_kwarg("deprecated_name", new_name="new_name", version=version)
+        def dummy_function(new_name=None, **kwargs):
+            return new_name
+
+        with self.assertWarns(FutureWarning):
+            result = dummy_function(deprecated_name="old_value", new_name="new_value")
+        self.assertEqual(result, "new_value")
+
+    @require_torch_gpu
+    def test_compile_safe(self):
+        @deprecate_kwarg("deprecated_factor", new_name="new_factor", version=INFINITE_VERSION)
+        def dummy_function(new_factor=None, **kwargs):
+            return new_factor * torch.ones(1, device="cuda")
+
+        compiled_function = torch.compile(dummy_function, fullgraph=True)
+
+        # Check that we can correctly call the compiled function with the old name, without raising errors
+        out = compiled_function(deprecated_factor=2)
+        self.assertEqual(out.item(), 2)
+
+        # Check that we can correctly call the compiled function with the new name, without raising errors
+        out = compiled_function(new_factor=2)
+        self.assertEqual(out.item(), 2)
+
+        # Check that we can correctly call the compiled function with both names, without raising errors
+        out = compiled_function(new_factor=2, deprecated_factor=10)
+        self.assertEqual(out.item(), 2)

docs/transformers/tests/utils/test_doc_samples.py

@@ -0,0 +1,112 @@
+# Copyright 2019-present, the HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import doctest
+import logging
+import os
+import unittest
+from pathlib import Path
+from typing import Union
+
+import transformers
+from transformers.testing_utils import require_tf, require_torch, slow
+
+
+logger = logging.getLogger()
+
+
+@unittest.skip(reason="Temporarily disable the doc tests.")
+@require_torch
+@require_tf
+@slow
+class TestCodeExamples(unittest.TestCase):
+    def analyze_directory(
+        self,
+        directory: Path,
+        identifier: Union[str, None] = None,
+        ignore_files: Union[list[str], None] = None,
+        n_identifier: Union[str, list[str], None] = None,
+        only_modules: bool = True,
+    ):
+        """
+        Runs through the specific directory, looking for the files identified with `identifier`. Executes
+        the doctests in those files
+
+        Args:
+            directory (`Path`): Directory containing the files
+            identifier (`str`): Will parse files containing this
+            ignore_files (`List[str]`): List of files to skip
+            n_identifier (`str` or `List[str]`): Will not parse files containing this/these identifiers.
+            only_modules (`bool`): Whether to only analyze modules
+        """
+        files = [file for file in os.listdir(directory) if os.path.isfile(os.path.join(directory, file))]
+
+        if identifier is not None:
+            files = [file for file in files if identifier in file]
+
+        if n_identifier is not None:
+            if isinstance(n_identifier, list):
+                for n_ in n_identifier:
+                    files = [file for file in files if n_ not in file]
+            else:
+                files = [file for file in files if n_identifier not in file]
+
+        ignore_files = ignore_files or []
+        ignore_files.append("__init__.py")
+        files = [file for file in files if file not in ignore_files]
+
+        for file in files:
+            # Open all files
+            print("Testing", file)
+
+            if only_modules:
+                module_identifier = file.split(".")[0]
+                try:
+                    module_identifier = getattr(transformers, module_identifier)
+                    suite = doctest.DocTestSuite(module_identifier)
+                    result = unittest.TextTestRunner().run(suite)
+                    self.assertIs(len(result.failures), 0)
+                except AttributeError:
+                    logger.info(f"{module_identifier} is not a module.")
+            else:
+                result = doctest.testfile(str(".." / directory / file), optionflags=doctest.ELLIPSIS)
+                self.assertIs(result.failed, 0)
+
+    def test_modeling_examples(self):
+        transformers_directory = Path("src/transformers")
+        files = "modeling"
+        ignore_files = [
+            "modeling_ctrl.py",
+            "modeling_tf_ctrl.py",
+        ]
+        self.analyze_directory(transformers_directory, identifier=files, ignore_files=ignore_files)
+
+    def test_tokenization_examples(self):
+        transformers_directory = Path("src/transformers")
+        files = "tokenization"
+        self.analyze_directory(transformers_directory, identifier=files)
+
+    def test_configuration_examples(self):
+        transformers_directory = Path("src/transformers")
+        files = "configuration"
+        self.analyze_directory(transformers_directory, identifier=files)
+
+    def test_remaining_examples(self):
+        transformers_directory = Path("src/transformers")
+        n_identifiers = ["configuration", "modeling", "tokenization"]
+        self.analyze_directory(transformers_directory, n_identifier=n_identifiers)
+
+    def test_doc_sources(self):
+        doc_source_directory = Path("docs/source")
+        ignore_files = ["favicon.ico"]
+        self.analyze_directory(doc_source_directory, ignore_files=ignore_files, only_modules=False)

docs/transformers/tests/utils/test_dynamic_module_utils.py

@@ -0,0 +1,129 @@
+# Copyright 2023 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+
+import pytest
+
+from transformers.dynamic_module_utils import get_imports
+
+
+TOP_LEVEL_IMPORT = """
+import os
+"""
+
+IMPORT_IN_FUNCTION = """
+def foo():
+    import os
+    return False
+"""
+
+DEEPLY_NESTED_IMPORT = """
+def foo():
+    def bar():
+        if True:
+            import os
+        return False
+    return bar()
+"""
+
+TOP_LEVEL_TRY_IMPORT = """
+import os
+
+try:
+    import bar
+except ImportError:
+    raise ValueError()
+"""
+
+TRY_IMPORT_IN_FUNCTION = """
+import os
+
+def foo():
+    try:
+        import bar
+    except ImportError:
+        raise ValueError()
+"""
+
+MULTIPLE_EXCEPTS_IMPORT = """
+import os
+
+try:
+    import bar
+except (ImportError, AttributeError):
+    raise ValueError()
+"""
+
+EXCEPT_AS_IMPORT = """
+import os
+
+try:
+    import bar
+except ImportError as e:
+    raise ValueError()
+"""
+
+GENERIC_EXCEPT_IMPORT = """
+import os
+
+try:
+    import bar
+except:
+    raise ValueError()
+"""
+
+MULTILINE_TRY_IMPORT = """
+import os
+
+try:
+    import bar
+    import baz
+except ImportError:
+    raise ValueError()
+"""
+
+MULTILINE_BOTH_IMPORT = """
+import os
+
+try:
+    import bar
+    import baz
+except ImportError:
+    x = 1
+    raise ValueError()
+"""
+
+CASES = [
+    TOP_LEVEL_IMPORT,
+    IMPORT_IN_FUNCTION,
+    DEEPLY_NESTED_IMPORT,
+    TOP_LEVEL_TRY_IMPORT,
+    GENERIC_EXCEPT_IMPORT,
+    MULTILINE_TRY_IMPORT,
+    MULTILINE_BOTH_IMPORT,
+    MULTIPLE_EXCEPTS_IMPORT,
+    EXCEPT_AS_IMPORT,
+    TRY_IMPORT_IN_FUNCTION,
+]
+
+
+@pytest.mark.parametrize("case", CASES)
+def test_import_parsing(tmp_path, case):
+    tmp_file_path = os.path.join(tmp_path, "test_file.py")
+    with open(tmp_file_path, "w") as _tmp_file:
+        _tmp_file.write(case)
+
+    parsed_imports = get_imports(tmp_file_path)
+    assert parsed_imports == ["os"]

docs/transformers/tests/utils/test_expectations.py

@@ -0,0 +1,34 @@
+import unittest
+
+from transformers.testing_utils import Expectations
+
+
+class ExpectationsTest(unittest.TestCase):
+    def test_expectations(self):
+        expectations = Expectations(
+            {
+                (None, None): 1,
+                ("cuda", 8): 2,
+                ("cuda", 7): 3,
+                ("rocm", 8): 4,
+                ("rocm", None): 5,
+                ("cpu", None): 6,
+                ("xpu", 3): 7,
+            }
+        )
+
+        def check(value, key):
+            assert expectations.find_expectation(key) == value
+
+        # npu has no matches so should find default expectation
+        check(1, ("npu", None))
+        check(7, ("xpu", 3))
+        check(2, ("cuda", 8))
+        check(3, ("cuda", 7))
+        check(4, ("rocm", 9))
+        check(4, ("rocm", None))
+        check(2, ("cuda", 2))
+
+        expectations = Expectations({("cuda", 8): 1})
+        with self.assertRaises(ValueError):
+            expectations.find_expectation(("xpu", None))

docs/transformers/tests/utils/test_feature_extraction_utils.py

@@ -0,0 +1,121 @@
+# Copyright 2021 HuggingFace Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import sys
+import tempfile
+import unittest
+import unittest.mock as mock
+from pathlib import Path
+
+from huggingface_hub import HfFolder
+from requests.exceptions import HTTPError
+
+from transformers import AutoFeatureExtractor, Wav2Vec2FeatureExtractor
+from transformers.testing_utils import TOKEN, TemporaryHubRepo, get_tests_dir, is_staging_test
+
+
+sys.path.append(str(Path(__file__).parent.parent.parent / "utils"))
+
+from test_module.custom_feature_extraction import CustomFeatureExtractor  # noqa E402
+
+
+SAMPLE_FEATURE_EXTRACTION_CONFIG_DIR = get_tests_dir("fixtures")
+
+
+class FeatureExtractorUtilTester(unittest.TestCase):
+    def test_cached_files_are_used_when_internet_is_down(self):
+        # A mock response for an HTTP head request to emulate server down
+        response_mock = mock.Mock()
+        response_mock.status_code = 500
+        response_mock.headers = {}
+        response_mock.raise_for_status.side_effect = HTTPError
+        response_mock.json.return_value = {}
+
+        # Download this model to make sure it's in the cache.
+        _ = Wav2Vec2FeatureExtractor.from_pretrained("hf-internal-testing/tiny-random-wav2vec2")
+        # Under the mock environment we get a 500 error when trying to reach the model.
+        with mock.patch("requests.Session.request", return_value=response_mock) as mock_head:
+            _ = Wav2Vec2FeatureExtractor.from_pretrained("hf-internal-testing/tiny-random-wav2vec2")
+            # This check we did call the fake head request
+            mock_head.assert_called()
+
+
+@is_staging_test
+class FeatureExtractorPushToHubTester(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls._token = TOKEN
+        HfFolder.save_token(TOKEN)
+
+    def test_push_to_hub(self):
+        with TemporaryHubRepo(token=self._token) as tmp_repo:
+            feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(SAMPLE_FEATURE_EXTRACTION_CONFIG_DIR)
+            feature_extractor.push_to_hub(tmp_repo.repo_id, token=self._token)
+
+            new_feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(tmp_repo.repo_id)
+            for k, v in feature_extractor.__dict__.items():
+                self.assertEqual(v, getattr(new_feature_extractor, k))
+
+    def test_push_to_hub_via_save_pretrained(self):
+        with TemporaryHubRepo(token=self._token) as tmp_repo:
+            feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(SAMPLE_FEATURE_EXTRACTION_CONFIG_DIR)
+            # Push to hub via save_pretrained
+            with tempfile.TemporaryDirectory() as tmp_dir:
+                feature_extractor.save_pretrained(
+                    tmp_dir, repo_id=tmp_repo.repo_id, push_to_hub=True, token=self._token
+                )
+
+            new_feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(tmp_repo.repo_id)
+            for k, v in feature_extractor.__dict__.items():
+                self.assertEqual(v, getattr(new_feature_extractor, k))
+
+    def test_push_to_hub_in_organization(self):
+        with TemporaryHubRepo(namespace="valid_org", token=self._token) as tmp_repo:
+            feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(SAMPLE_FEATURE_EXTRACTION_CONFIG_DIR)
+            feature_extractor.push_to_hub(tmp_repo.repo_id, token=self._token)
+
+            new_feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(tmp_repo.repo_id)
+            for k, v in feature_extractor.__dict__.items():
+                self.assertEqual(v, getattr(new_feature_extractor, k))
+
+    def test_push_to_hub_in_organization_via_save_pretrained(self):
+        with TemporaryHubRepo(namespace="valid_org", token=self._token) as tmp_repo:
+            feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(SAMPLE_FEATURE_EXTRACTION_CONFIG_DIR)
+            # Push to hub via save_pretrained
+            with tempfile.TemporaryDirectory() as tmp_dir:
+                feature_extractor.save_pretrained(
+                    tmp_dir, repo_id=tmp_repo.repo_id, push_to_hub=True, token=self._token
+                )
+
+            new_feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(tmp_repo.repo_id)
+            for k, v in feature_extractor.__dict__.items():
+                self.assertEqual(v, getattr(new_feature_extractor, k))
+
+    def test_push_to_hub_dynamic_feature_extractor(self):
+        with TemporaryHubRepo(token=self._token) as tmp_repo:
+            CustomFeatureExtractor.register_for_auto_class()
+            feature_extractor = CustomFeatureExtractor.from_pretrained(SAMPLE_FEATURE_EXTRACTION_CONFIG_DIR)
+
+            feature_extractor.push_to_hub(tmp_repo.repo_id, token=self._token)
+
+            # This has added the proper auto_map field to the config
+            self.assertDictEqual(
+                feature_extractor.auto_map,
+                {"AutoFeatureExtractor": "custom_feature_extraction.CustomFeatureExtractor"},
+            )
+
+            new_feature_extractor = AutoFeatureExtractor.from_pretrained(tmp_repo.repo_id, trust_remote_code=True)
+            # Can't make an isinstance check because the new_feature_extractor is from the CustomFeatureExtractor class of a dynamic module
+            self.assertEqual(new_feature_extractor.__class__.__name__, "CustomFeatureExtractor")

docs/transformers/tests/utils/test_file_utils.py

@@ -0,0 +1,133 @@
+# Copyright 2020 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import contextlib
+import importlib
+import io
+import unittest
+
+import transformers
+
+# Try to import everything from transformers to ensure every object can be loaded.
+from transformers import *  # noqa F406
+from transformers.testing_utils import DUMMY_UNKNOWN_IDENTIFIER, require_flax, require_tf, require_torch
+from transformers.utils import ContextManagers, find_labels, is_flax_available, is_tf_available, is_torch_available
+
+
+if is_torch_available():
+    from transformers import BertForPreTraining, BertForQuestionAnswering, BertForSequenceClassification
+
+if is_tf_available():
+    from transformers import TFBertForPreTraining, TFBertForQuestionAnswering, TFBertForSequenceClassification
+
+if is_flax_available():
+    from transformers import FlaxBertForPreTraining, FlaxBertForQuestionAnswering, FlaxBertForSequenceClassification
+
+
+MODEL_ID = DUMMY_UNKNOWN_IDENTIFIER
+# An actual model hosted on huggingface.co
+
+REVISION_ID_DEFAULT = "main"
+# Default branch name
+REVISION_ID_ONE_SPECIFIC_COMMIT = "f2c752cfc5c0ab6f4bdec59acea69eefbee381c2"
+# One particular commit (not the top of `main`)
+REVISION_ID_INVALID = "aaaaaaa"
+# This commit does not exist, so we should 404.
+
+PINNED_SHA1 = "d9e9f15bc825e4b2c9249e9578f884bbcb5e3684"
+# Sha-1 of config.json on the top of `main`, for checking purposes
+PINNED_SHA256 = "4b243c475af8d0a7754e87d7d096c92e5199ec2fe168a2ee7998e3b8e9bcb1d3"
+# Sha-256 of pytorch_model.bin on the top of `main`, for checking purposes
+
+
+# Dummy contexts to test `ContextManagers`
+@contextlib.contextmanager
+def context_en():
+    print("Welcome!")
+    yield
+    print("Bye!")
+
+
+@contextlib.contextmanager
+def context_fr():
+    print("Bonjour!")
+    yield
+    print("Au revoir!")
+
+
+class TestImportMechanisms(unittest.TestCase):
+    def test_module_spec_available(self):
+        # If the spec is missing, importlib would not be able to import the module dynamically.
+        assert transformers.__spec__ is not None
+        assert importlib.util.find_spec("transformers") is not None
+
+
+class GenericUtilTests(unittest.TestCase):
+    @unittest.mock.patch("sys.stdout", new_callable=io.StringIO)
+    def test_context_managers_no_context(self, mock_stdout):
+        with ContextManagers([]):
+            print("Transformers are awesome!")
+        # The print statement adds a new line at the end of the output
+        self.assertEqual(mock_stdout.getvalue(), "Transformers are awesome!\n")
+
+    @unittest.mock.patch("sys.stdout", new_callable=io.StringIO)
+    def test_context_managers_one_context(self, mock_stdout):
+        with ContextManagers([context_en()]):
+            print("Transformers are awesome!")
+        # The output should be wrapped with an English welcome and goodbye
+        self.assertEqual(mock_stdout.getvalue(), "Welcome!\nTransformers are awesome!\nBye!\n")
+
+    @unittest.mock.patch("sys.stdout", new_callable=io.StringIO)
+    def test_context_managers_two_context(self, mock_stdout):
+        with ContextManagers([context_fr(), context_en()]):
+            print("Transformers are awesome!")
+        # The output should be wrapped with an English and French welcome and goodbye
+        self.assertEqual(mock_stdout.getvalue(), "Bonjour!\nWelcome!\nTransformers are awesome!\nBye!\nAu revoir!\n")
+
+    @require_torch
+    def test_find_labels_pt(self):
+        self.assertEqual(find_labels(BertForSequenceClassification), ["labels"])
+        self.assertEqual(find_labels(BertForPreTraining), ["labels", "next_sentence_label"])
+        self.assertEqual(find_labels(BertForQuestionAnswering), ["start_positions", "end_positions"])
+
+        # find_labels works regardless of the class name (it detects the framework through inheritance)
+        class DummyModel(BertForSequenceClassification):
+            pass
+
+        self.assertEqual(find_labels(DummyModel), ["labels"])
+
+    @require_tf
+    def test_find_labels_tf(self):
+        self.assertEqual(find_labels(TFBertForSequenceClassification), ["labels"])
+        self.assertEqual(find_labels(TFBertForPreTraining), ["labels", "next_sentence_label"])
+        self.assertEqual(find_labels(TFBertForQuestionAnswering), ["start_positions", "end_positions"])
+
+        # find_labels works regardless of the class name (it detects the framework through inheritance)
+        class DummyModel(TFBertForSequenceClassification):
+            pass
+
+        self.assertEqual(find_labels(DummyModel), ["labels"])
+
+    @require_flax
+    def test_find_labels_flax(self):
+        # Flax models don't have labels
+        self.assertEqual(find_labels(FlaxBertForSequenceClassification), [])
+        self.assertEqual(find_labels(FlaxBertForPreTraining), [])
+        self.assertEqual(find_labels(FlaxBertForQuestionAnswering), [])
+
+        # find_labels works regardless of the class name (it detects the framework through inheritance)
+        class DummyModel(FlaxBertForSequenceClassification):
+            pass
+
+        self.assertEqual(find_labels(DummyModel), [])

docs/transformers/tests/utils/test_generic.py

@@ -0,0 +1,463 @@
+# Copyright 2019-present, the HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+import warnings
+
+import numpy as np
+
+from transformers.configuration_utils import PretrainedConfig
+from transformers.modeling_outputs import BaseModelOutput
+from transformers.testing_utils import require_flax, require_tf, require_torch
+from transformers.utils import (
+    can_return_tuple,
+    expand_dims,
+    filter_out_non_signature_kwargs,
+    flatten_dict,
+    is_flax_available,
+    is_tf_available,
+    is_torch_available,
+    reshape,
+    squeeze,
+    to_py_obj,
+    transpose,
+)
+
+
+if is_flax_available():
+    import jax.numpy as jnp
+
+if is_tf_available():
+    import tensorflow as tf
+
+if is_torch_available():
+    import torch
+
+
+class GenericTester(unittest.TestCase):
+    def test_flatten_dict(self):
+        input_dict = {
+            "task_specific_params": {
+                "summarization": {"length_penalty": 1.0, "max_length": 128, "min_length": 12, "num_beams": 4},
+                "summarization_cnn": {"length_penalty": 2.0, "max_length": 142, "min_length": 56, "num_beams": 4},
+                "summarization_xsum": {"length_penalty": 1.0, "max_length": 62, "min_length": 11, "num_beams": 6},
+            }
+        }
+        expected_dict = {
+            "task_specific_params.summarization.length_penalty": 1.0,
+            "task_specific_params.summarization.max_length": 128,
+            "task_specific_params.summarization.min_length": 12,
+            "task_specific_params.summarization.num_beams": 4,
+            "task_specific_params.summarization_cnn.length_penalty": 2.0,
+            "task_specific_params.summarization_cnn.max_length": 142,
+            "task_specific_params.summarization_cnn.min_length": 56,
+            "task_specific_params.summarization_cnn.num_beams": 4,
+            "task_specific_params.summarization_xsum.length_penalty": 1.0,
+            "task_specific_params.summarization_xsum.max_length": 62,
+            "task_specific_params.summarization_xsum.min_length": 11,
+            "task_specific_params.summarization_xsum.num_beams": 6,
+        }
+
+        self.assertEqual(flatten_dict(input_dict), expected_dict)
+
+    def test_transpose_numpy(self):
+        x = np.random.randn(3, 4)
+        self.assertTrue(np.allclose(transpose(x), x.transpose()))
+
+        x = np.random.randn(3, 4, 5)
+        self.assertTrue(np.allclose(transpose(x, axes=(1, 2, 0)), x.transpose((1, 2, 0))))
+
+    @require_torch
+    def test_transpose_torch(self):
+        x = np.random.randn(3, 4)
+        t = torch.tensor(x)
+        self.assertTrue(np.allclose(transpose(x), transpose(t).numpy()))
+
+        x = np.random.randn(3, 4, 5)
+        t = torch.tensor(x)
+        self.assertTrue(np.allclose(transpose(x, axes=(1, 2, 0)), transpose(t, axes=(1, 2, 0)).numpy()))
+
+    @require_tf
+    def test_transpose_tf(self):
+        x = np.random.randn(3, 4)
+        t = tf.constant(x)
+        self.assertTrue(np.allclose(transpose(x), transpose(t).numpy()))
+
+        x = np.random.randn(3, 4, 5)
+        t = tf.constant(x)
+        self.assertTrue(np.allclose(transpose(x, axes=(1, 2, 0)), transpose(t, axes=(1, 2, 0)).numpy()))
+
+    @require_flax
+    def test_transpose_flax(self):
+        x = np.random.randn(3, 4)
+        t = jnp.array(x)
+        self.assertTrue(np.allclose(transpose(x), np.asarray(transpose(t))))
+
+        x = np.random.randn(3, 4, 5)
+        t = jnp.array(x)
+        self.assertTrue(np.allclose(transpose(x, axes=(1, 2, 0)), np.asarray(transpose(t, axes=(1, 2, 0)))))
+
+    def test_reshape_numpy(self):
+        x = np.random.randn(3, 4)
+        self.assertTrue(np.allclose(reshape(x, (4, 3)), np.reshape(x, (4, 3))))
+
+        x = np.random.randn(3, 4, 5)
+        self.assertTrue(np.allclose(reshape(x, (12, 5)), np.reshape(x, (12, 5))))
+
+    @require_torch
+    def test_reshape_torch(self):
+        x = np.random.randn(3, 4)
+        t = torch.tensor(x)
+        self.assertTrue(np.allclose(reshape(x, (4, 3)), reshape(t, (4, 3)).numpy()))
+
+        x = np.random.randn(3, 4, 5)
+        t = torch.tensor(x)
+        self.assertTrue(np.allclose(reshape(x, (12, 5)), reshape(t, (12, 5)).numpy()))
+
+    @require_tf
+    def test_reshape_tf(self):
+        x = np.random.randn(3, 4)
+        t = tf.constant(x)
+        self.assertTrue(np.allclose(reshape(x, (4, 3)), reshape(t, (4, 3)).numpy()))
+
+        x = np.random.randn(3, 4, 5)
+        t = tf.constant(x)
+        self.assertTrue(np.allclose(reshape(x, (12, 5)), reshape(t, (12, 5)).numpy()))
+
+    @require_flax
+    def test_reshape_flax(self):
+        x = np.random.randn(3, 4)
+        t = jnp.array(x)
+        self.assertTrue(np.allclose(reshape(x, (4, 3)), np.asarray(reshape(t, (4, 3)))))
+
+        x = np.random.randn(3, 4, 5)
+        t = jnp.array(x)
+        self.assertTrue(np.allclose(reshape(x, (12, 5)), np.asarray(reshape(t, (12, 5)))))
+
+    def test_squeeze_numpy(self):
+        x = np.random.randn(1, 3, 4)
+        self.assertTrue(np.allclose(squeeze(x), np.squeeze(x)))
+
+        x = np.random.randn(1, 4, 1, 5)
+        self.assertTrue(np.allclose(squeeze(x, axis=2), np.squeeze(x, axis=2)))
+
+    @require_torch
+    def test_squeeze_torch(self):
+        x = np.random.randn(1, 3, 4)
+        t = torch.tensor(x)
+        self.assertTrue(np.allclose(squeeze(x), squeeze(t).numpy()))
+
+        x = np.random.randn(1, 4, 1, 5)
+        t = torch.tensor(x)
+        self.assertTrue(np.allclose(squeeze(x, axis=2), squeeze(t, axis=2).numpy()))
+
+    @require_tf
+    def test_squeeze_tf(self):
+        x = np.random.randn(1, 3, 4)
+        t = tf.constant(x)
+        self.assertTrue(np.allclose(squeeze(x), squeeze(t).numpy()))
+
+        x = np.random.randn(1, 4, 1, 5)
+        t = tf.constant(x)
+        self.assertTrue(np.allclose(squeeze(x, axis=2), squeeze(t, axis=2).numpy()))
+
+    @require_flax
+    def test_squeeze_flax(self):
+        x = np.random.randn(1, 3, 4)
+        t = jnp.array(x)
+        self.assertTrue(np.allclose(squeeze(x), np.asarray(squeeze(t))))
+
+        x = np.random.randn(1, 4, 1, 5)
+        t = jnp.array(x)
+        self.assertTrue(np.allclose(squeeze(x, axis=2), np.asarray(squeeze(t, axis=2))))
+
+    def test_expand_dims_numpy(self):
+        x = np.random.randn(3, 4)
+        self.assertTrue(np.allclose(expand_dims(x, axis=1), np.expand_dims(x, axis=1)))
+
+    @require_torch
+    def test_expand_dims_torch(self):
+        x = np.random.randn(3, 4)
+        t = torch.tensor(x)
+        self.assertTrue(np.allclose(expand_dims(x, axis=1), expand_dims(t, axis=1).numpy()))
+
+    @require_tf
+    def test_expand_dims_tf(self):
+        x = np.random.randn(3, 4)
+        t = tf.constant(x)
+        self.assertTrue(np.allclose(expand_dims(x, axis=1), expand_dims(t, axis=1).numpy()))
+
+    @require_flax
+    def test_expand_dims_flax(self):
+        x = np.random.randn(3, 4)
+        t = jnp.array(x)
+        self.assertTrue(np.allclose(expand_dims(x, axis=1), np.asarray(expand_dims(t, axis=1))))
+
+    def test_to_py_obj_native(self):
+        self.assertTrue(to_py_obj(1) == 1)
+        self.assertTrue(to_py_obj([1, 2, 3]) == [1, 2, 3])
+        self.assertTrue(to_py_obj([((1.0, 1.1), 1.2), (2, 3)]) == [[[1.0, 1.1], 1.2], [2, 3]])
+
+    def test_to_py_obj_numpy(self):
+        x1 = [[1, 2, 3], [4, 5, 6]]
+        t1 = np.array(x1)
+        self.assertTrue(to_py_obj(t1) == x1)
+
+        x2 = [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]
+        t2 = np.array(x2)
+        self.assertTrue(to_py_obj(t2) == x2)
+
+        self.assertTrue(to_py_obj([t1, t2]) == [x1, x2])
+
+    @require_torch
+    def test_to_py_obj_torch(self):
+        x1 = [[1, 2, 3], [4, 5, 6]]
+        t1 = torch.tensor(x1)
+        self.assertTrue(to_py_obj(t1) == x1)
+
+        x2 = [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]
+        t2 = torch.tensor(x2)
+        self.assertTrue(to_py_obj(t2) == x2)
+
+        self.assertTrue(to_py_obj([t1, t2]) == [x1, x2])
+
+    @require_tf
+    def test_to_py_obj_tf(self):
+        x1 = [[1, 2, 3], [4, 5, 6]]
+        t1 = tf.constant(x1)
+        self.assertTrue(to_py_obj(t1) == x1)
+
+        x2 = [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]
+        t2 = tf.constant(x2)
+        self.assertTrue(to_py_obj(t2) == x2)
+
+        self.assertTrue(to_py_obj([t1, t2]) == [x1, x2])
+
+    @require_flax
+    def test_to_py_obj_flax(self):
+        x1 = [[1, 2, 3], [4, 5, 6]]
+        t1 = jnp.array(x1)
+        self.assertTrue(to_py_obj(t1) == x1)
+
+        x2 = [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]
+        t2 = jnp.array(x2)
+        self.assertTrue(to_py_obj(t2) == x2)
+
+        self.assertTrue(to_py_obj([t1, t2]) == [x1, x2])
+
+    @require_torch
+    @require_tf
+    @require_flax
+    def test_to_py_obj_mixed(self):
+        x1 = [[1], [2]]
+        t1 = np.array(x1)
+
+        x2 = [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]
+        t2 = torch.tensor(x2)
+
+        x3 = [1, 2, 3]
+        t3 = tf.constant(x3)
+
+        x4 = [[[1.0, 2.0]]]
+        t4 = jnp.array(x4)
+
+        mixed = [(t1, t2), (t3, t4)]
+        self.assertTrue(to_py_obj(mixed) == [[x1, x2], [x3, x4]])
+
+
+class ValidationDecoratorTester(unittest.TestCase):
+    def test_cases_no_warning(self):
+        with warnings.catch_warnings(record=True) as raised_warnings:
+            warnings.simplefilter("always")
+
+            # basic test
+            @filter_out_non_signature_kwargs()
+            def func1(a):
+                return a
+
+            result = func1(1)
+            self.assertEqual(result, 1)
+
+            # include extra kwarg
+            @filter_out_non_signature_kwargs(extra=["extra_arg"])
+            def func2(a, **kwargs):
+                return a, kwargs
+
+            a, kwargs = func2(1)
+            self.assertEqual(a, 1)
+            self.assertEqual(kwargs, {})
+
+            a, kwargs = func2(1, extra_arg=2)
+            self.assertEqual(a, 1)
+            self.assertEqual(kwargs, {"extra_arg": 2})
+
+            # multiple extra kwargs
+            @filter_out_non_signature_kwargs(extra=["extra_arg", "extra_arg2"])
+            def func3(a, **kwargs):
+                return a, kwargs
+
+            a, kwargs = func3(2)
+            self.assertEqual(a, 2)
+            self.assertEqual(kwargs, {})
+
+            a, kwargs = func3(3, extra_arg2=3)
+            self.assertEqual(a, 3)
+            self.assertEqual(kwargs, {"extra_arg2": 3})
+
+            a, kwargs = func3(1, extra_arg=2, extra_arg2=3)
+            self.assertEqual(a, 1)
+            self.assertEqual(kwargs, {"extra_arg": 2, "extra_arg2": 3})
+
+            # Check that no warnings were raised
+            self.assertEqual(len(raised_warnings), 0, f"Warning raised: {[w.message for w in raised_warnings]}")
+
+    def test_cases_with_warnings(self):
+        @filter_out_non_signature_kwargs()
+        def func1(a):
+            return a
+
+        with self.assertWarns(UserWarning):
+            func1(1, extra_arg=2)
+
+        @filter_out_non_signature_kwargs(extra=["extra_arg"])
+        def func2(a, **kwargs):
+            return kwargs
+
+        with self.assertWarns(UserWarning):
+            kwargs = func2(1, extra_arg=2, extra_arg2=3)
+        self.assertEqual(kwargs, {"extra_arg": 2})
+
+        @filter_out_non_signature_kwargs(extra=["extra_arg", "extra_arg2"])
+        def func3(a, **kwargs):
+            return kwargs
+
+        with self.assertWarns(UserWarning):
+            kwargs = func3(1, extra_arg=2, extra_arg2=3, extra_arg3=4)
+        self.assertEqual(kwargs, {"extra_arg": 2, "extra_arg2": 3})
+
+
+@require_torch
+class CanReturnTupleDecoratorTester(unittest.TestCase):
+    def _get_model(self, config, store_config=True, raise_in_forward=False):
+        # Simple model class for testing can_return_tuple decorator.
+        class SimpleTestModel(torch.nn.Module):
+            def __init__(self, config):
+                super().__init__()
+                if store_config:
+                    self.config = config
+
+            @can_return_tuple
+            def forward(self, x):
+                if raise_in_forward:
+                    raise ValueError("Test error")
+                return BaseModelOutput(
+                    last_hidden_state=x,
+                    hidden_states=None,
+                    attentions=None,
+                )
+
+        return SimpleTestModel(config)
+
+    def test_decorator_eager(self):
+        """Test that the can_return_tuple decorator works with eager mode."""
+
+        # test nothing is set
+        config = PretrainedConfig()
+        model = self._get_model(config)
+        inputs = torch.tensor(10)
+        output = model(inputs)
+        self.assertIsInstance(
+            output, BaseModelOutput, "output should be a BaseModelOutput when return_dict is not set"
+        )
+
+        # test all explicit cases
+        for config_return_dict in [True, False, None]:
+            for return_dict in [True, False, None]:
+                config = PretrainedConfig(return_dict=config_return_dict)
+                model = self._get_model(config)
+                output = model(torch.tensor(10), return_dict=return_dict)
+
+                expected_type = tuple if config_return_dict is False or return_dict is False else BaseModelOutput
+                message = f"output should be a {expected_type.__name__} when config.use_return_dict={config_return_dict} and return_dict={return_dict}"
+                self.assertIsInstance(output, expected_type, message)
+
+    def test_decorator_compiled(self):
+        """Test that the can_return_tuple decorator works with compiled mode."""
+        config = PretrainedConfig()
+
+        # Output object
+        model = self._get_model(config)
+        compiled_model = torch.compile(model)
+        output = compiled_model(torch.tensor(10))
+        self.assertIsInstance(output, BaseModelOutput)
+
+        # Tuple output
+        model = self._get_model(config)
+        compiled_model = torch.compile(model)
+        output = compiled_model(torch.tensor(10), return_dict=False)
+        self.assertIsInstance(output, tuple)
+
+    def test_decorator_torch_export(self):
+        """Test that the can_return_tuple decorator works with torch.export."""
+        config = PretrainedConfig()
+        model = self._get_model(config)
+        torch.export.export(model, args=(torch.tensor(10),))
+
+    def test_decorator_torchscript(self):
+        """Test that the can_return_tuple decorator works with torch.jit.trace."""
+        config = PretrainedConfig(return_dict=False)
+        model = self._get_model(config)
+        inputs = torch.tensor(10)
+        traced_module = torch.jit.trace(model, inputs)
+        output = traced_module(inputs)
+        self.assertIsInstance(output, tuple)
+
+    def test_attribute_cleanup(self):
+        """Test that the `_is_top_level_module` attribute is removed after the forward call."""
+
+        config = PretrainedConfig(return_dict=False)
+        inputs = torch.tensor(10)
+
+        # working case
+        model = self._get_model(config)
+        output = model(inputs)
+
+        self.assertIsInstance(output, tuple)
+        for name, module in model.named_modules():
+            self.assertFalse(
+                hasattr(module, "_is_top_level_module"),
+                f"Module `{name}` should not have `_is_top_level_module` attribute",
+            )
+
+        # model without config
+        no_config_model = self._get_model(config, store_config=False)
+        output = no_config_model(inputs)
+
+        self.assertIsInstance(output, BaseModelOutput)
+        for name, module in no_config_model.named_modules():
+            self.assertFalse(
+                hasattr(module, "_is_top_level_module"),
+                f"Module `{name}` should not have `_is_top_level_module` attribute",
+            )
+
+        # model with raise in forward
+        model_with_raise = self._get_model(config, raise_in_forward=True)
+        with self.assertRaises(ValueError):
+            model_with_raise(inputs)
+
+        for name, module in model_with_raise.named_modules():
+            self.assertFalse(
+                hasattr(module, "_is_top_level_module"),
+                f"Module `{name}` should not have `_is_top_level_module` attribute",
+            )

docs/transformers/tests/utils/test_hf_argparser.py

@@ -0,0 +1,482 @@
+# Copyright 2020 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+import json
+import os
+import sys
+import tempfile
+import unittest
+from argparse import Namespace
+from dataclasses import dataclass, field
+from enum import Enum
+from pathlib import Path
+from typing import List, Literal, Optional, Union, get_args, get_origin
+
+import yaml
+
+from transformers import HfArgumentParser, TrainingArguments
+from transformers.hf_argparser import make_choice_type_function, string_to_bool
+from transformers.testing_utils import require_torch
+
+
+# Since Python 3.10, we can use the builtin `|` operator for Union types
+# See PEP 604: https://peps.python.org/pep-0604
+is_python_no_less_than_3_10 = sys.version_info >= (3, 10)
+
+
+def list_field(default=None, metadata=None):
+    return field(default_factory=lambda: default, metadata=metadata)
+
+
+@dataclass
+class BasicExample:
+    foo: int
+    bar: float
+    baz: str
+    flag: bool
+
+
+@dataclass
+class WithDefaultExample:
+    foo: int = 42
+    baz: str = field(default="toto", metadata={"help": "help message"})
+
+
+@dataclass
+class WithDefaultBoolExample:
+    foo: bool = False
+    baz: bool = True
+    opt: Optional[bool] = None
+
+
+class BasicEnum(Enum):
+    titi = "titi"
+    toto = "toto"
+
+
+class MixedTypeEnum(Enum):
+    titi = "titi"
+    toto = "toto"
+    fourtytwo = 42
+
+
+@dataclass
+class EnumExample:
+    foo: BasicEnum = "toto"
+
+    def __post_init__(self):
+        self.foo = BasicEnum(self.foo)
+
+
+@dataclass
+class MixedTypeEnumExample:
+    foo: MixedTypeEnum = "toto"
+
+    def __post_init__(self):
+        self.foo = MixedTypeEnum(self.foo)
+
+
+@dataclass
+class OptionalExample:
+    foo: Optional[int] = None
+    bar: Optional[float] = field(default=None, metadata={"help": "help message"})
+    baz: Optional[str] = None
+    ces: Optional[list[str]] = list_field(default=[])
+    des: Optional[list[int]] = list_field(default=[])
+
+
+@dataclass
+class ListExample:
+    foo_int: list[int] = list_field(default=[])
+    bar_int: list[int] = list_field(default=[1, 2, 3])
+    foo_str: list[str] = list_field(default=["Hallo", "Bonjour", "Hello"])
+    foo_float: list[float] = list_field(default=[0.1, 0.2, 0.3])
+
+
+@dataclass
+class RequiredExample:
+    required_list: list[int] = field()
+    required_str: str = field()
+    required_enum: BasicEnum = field()
+
+    def __post_init__(self):
+        self.required_enum = BasicEnum(self.required_enum)
+
+
+@dataclass
+class StringLiteralAnnotationExample:
+    foo: int
+    required_enum: "BasicEnum" = field()
+    opt: "Optional[bool]" = None
+    baz: "str" = field(default="toto", metadata={"help": "help message"})
+    foo_str: "List[str]" = list_field(default=["Hallo", "Bonjour", "Hello"])
+
+
+if is_python_no_less_than_3_10:
+
+    @dataclass
+    class WithDefaultBoolExamplePep604:
+        foo: bool = False
+        baz: bool = True
+        opt: bool | None = None
+
+    @dataclass
+    class OptionalExamplePep604:
+        foo: int | None = None
+        bar: float | None = field(default=None, metadata={"help": "help message"})
+        baz: str | None = None
+        ces: list[str] | None = list_field(default=[])
+        des: list[int] | None = list_field(default=[])
+
+
+class HfArgumentParserTest(unittest.TestCase):
+    def argparsersEqual(self, a: argparse.ArgumentParser, b: argparse.ArgumentParser):
+        """
+        Small helper to check pseudo-equality of parsed arguments on `ArgumentParser` instances.
+        """
+        self.assertEqual(len(a._actions), len(b._actions))
+        for x, y in zip(a._actions, b._actions):
+            xx = {k: v for k, v in vars(x).items() if k != "container"}
+            yy = {k: v for k, v in vars(y).items() if k != "container"}
+
+            # Choices with mixed type have custom function as "type"
+            # So we need to compare results directly for equality
+            if xx.get("choices", None) and yy.get("choices", None):
+                for expected_choice in yy["choices"] + xx["choices"]:
+                    self.assertEqual(xx["type"](expected_choice), yy["type"](expected_choice))
+                del xx["type"], yy["type"]
+
+            self.assertEqual(xx, yy)
+
+    def test_basic(self):
+        parser = HfArgumentParser(BasicExample)
+
+        expected = argparse.ArgumentParser()
+        expected.add_argument("--foo", type=int, required=True)
+        expected.add_argument("--bar", type=float, required=True)
+        expected.add_argument("--baz", type=str, required=True)
+        expected.add_argument("--flag", type=string_to_bool, default=False, const=True, nargs="?")
+        self.argparsersEqual(parser, expected)
+
+        args = ["--foo", "1", "--baz", "quux", "--bar", "0.5"]
+        (example,) = parser.parse_args_into_dataclasses(args, look_for_args_file=False)
+        self.assertFalse(example.flag)
+
+    def test_with_default(self):
+        parser = HfArgumentParser(WithDefaultExample)
+
+        expected = argparse.ArgumentParser()
+        expected.add_argument("--foo", default=42, type=int)
+        expected.add_argument("--baz", default="toto", type=str, help="help message")
+        self.argparsersEqual(parser, expected)
+
+    def test_with_default_bool(self):
+        expected = argparse.ArgumentParser()
+        expected.add_argument("--foo", type=string_to_bool, default=False, const=True, nargs="?")
+        expected.add_argument("--baz", type=string_to_bool, default=True, const=True, nargs="?")
+        # A boolean no_* argument always has to come after its "default: True" regular counter-part
+        # and its default must be set to False
+        expected.add_argument("--no_baz", "--no-baz", action="store_false", default=False, dest="baz")
+        expected.add_argument("--opt", type=string_to_bool, default=None)
+
+        dataclass_types = [WithDefaultBoolExample]
+        if is_python_no_less_than_3_10:
+            dataclass_types.append(WithDefaultBoolExamplePep604)
+
+        for dataclass_type in dataclass_types:
+            parser = HfArgumentParser(dataclass_type)
+            self.argparsersEqual(parser, expected)
+
+            args = parser.parse_args([])
+            self.assertEqual(args, Namespace(foo=False, baz=True, opt=None))
+
+            args = parser.parse_args(["--foo", "--no_baz"])
+            self.assertEqual(args, Namespace(foo=True, baz=False, opt=None))
+
+            args = parser.parse_args(["--foo", "--no-baz"])
+            self.assertEqual(args, Namespace(foo=True, baz=False, opt=None))
+
+            args = parser.parse_args(["--foo", "--baz"])
+            self.assertEqual(args, Namespace(foo=True, baz=True, opt=None))
+
+            args = parser.parse_args(["--foo", "True", "--baz", "True", "--opt", "True"])
+            self.assertEqual(args, Namespace(foo=True, baz=True, opt=True))
+
+            args = parser.parse_args(["--foo", "False", "--baz", "False", "--opt", "False"])
+            self.assertEqual(args, Namespace(foo=False, baz=False, opt=False))
+
+    def test_with_enum(self):
+        parser = HfArgumentParser(MixedTypeEnumExample)
+
+        expected = argparse.ArgumentParser()
+        expected.add_argument(
+            "--foo",
+            default="toto",
+            choices=["titi", "toto", 42],
+            type=make_choice_type_function(["titi", "toto", 42]),
+        )
+        self.argparsersEqual(parser, expected)
+
+        args = parser.parse_args([])
+        self.assertEqual(args.foo, "toto")
+        enum_ex = parser.parse_args_into_dataclasses([])[0]
+        self.assertEqual(enum_ex.foo, MixedTypeEnum.toto)
+
+        args = parser.parse_args(["--foo", "titi"])
+        self.assertEqual(args.foo, "titi")
+        enum_ex = parser.parse_args_into_dataclasses(["--foo", "titi"])[0]
+        self.assertEqual(enum_ex.foo, MixedTypeEnum.titi)
+
+        args = parser.parse_args(["--foo", "42"])
+        self.assertEqual(args.foo, 42)
+        enum_ex = parser.parse_args_into_dataclasses(["--foo", "42"])[0]
+        self.assertEqual(enum_ex.foo, MixedTypeEnum.fourtytwo)
+
+    def test_with_literal(self):
+        @dataclass
+        class LiteralExample:
+            foo: Literal["titi", "toto", 42] = "toto"
+
+        parser = HfArgumentParser(LiteralExample)
+
+        expected = argparse.ArgumentParser()
+        expected.add_argument(
+            "--foo",
+            default="toto",
+            choices=("titi", "toto", 42),
+            type=make_choice_type_function(["titi", "toto", 42]),
+        )
+        self.argparsersEqual(parser, expected)
+
+        args = parser.parse_args([])
+        self.assertEqual(args.foo, "toto")
+
+        args = parser.parse_args(["--foo", "titi"])
+        self.assertEqual(args.foo, "titi")
+
+        args = parser.parse_args(["--foo", "42"])
+        self.assertEqual(args.foo, 42)
+
+    def test_with_list(self):
+        parser = HfArgumentParser(ListExample)
+
+        expected = argparse.ArgumentParser()
+        expected.add_argument("--foo_int", "--foo-int", nargs="+", default=[], type=int)
+        expected.add_argument("--bar_int", "--bar-int", nargs="+", default=[1, 2, 3], type=int)
+        expected.add_argument("--foo_str", "--foo-str", nargs="+", default=["Hallo", "Bonjour", "Hello"], type=str)
+        expected.add_argument("--foo_float", "--foo-float", nargs="+", default=[0.1, 0.2, 0.3], type=float)
+
+        self.argparsersEqual(parser, expected)
+
+        args = parser.parse_args([])
+        self.assertEqual(
+            args,
+            Namespace(foo_int=[], bar_int=[1, 2, 3], foo_str=["Hallo", "Bonjour", "Hello"], foo_float=[0.1, 0.2, 0.3]),
+        )
+
+        args = parser.parse_args("--foo_int 1 --bar_int 2 3 --foo_str a b c --foo_float 0.1 0.7".split())
+        self.assertEqual(args, Namespace(foo_int=[1], bar_int=[2, 3], foo_str=["a", "b", "c"], foo_float=[0.1, 0.7]))
+
+        args = parser.parse_args("--foo-int 1 --bar-int 2 3 --foo-str a b c --foo-float 0.1 0.7".split())
+        self.assertEqual(args, Namespace(foo_int=[1], bar_int=[2, 3], foo_str=["a", "b", "c"], foo_float=[0.1, 0.7]))
+
+    def test_with_optional(self):
+        expected = argparse.ArgumentParser()
+        expected.add_argument("--foo", default=None, type=int)
+        expected.add_argument("--bar", default=None, type=float, help="help message")
+        expected.add_argument("--baz", default=None, type=str)
+        expected.add_argument("--ces", nargs="+", default=[], type=str)
+        expected.add_argument("--des", nargs="+", default=[], type=int)
+
+        dataclass_types = [OptionalExample]
+        if is_python_no_less_than_3_10:
+            dataclass_types.append(OptionalExamplePep604)
+
+        for dataclass_type in dataclass_types:
+            parser = HfArgumentParser(dataclass_type)
+
+            self.argparsersEqual(parser, expected)
+
+            args = parser.parse_args([])
+            self.assertEqual(args, Namespace(foo=None, bar=None, baz=None, ces=[], des=[]))
+
+            args = parser.parse_args("--foo 12 --bar 3.14 --baz 42 --ces a b c --des 1 2 3".split())
+            self.assertEqual(args, Namespace(foo=12, bar=3.14, baz="42", ces=["a", "b", "c"], des=[1, 2, 3]))
+
+    def test_with_required(self):
+        parser = HfArgumentParser(RequiredExample)
+
+        expected = argparse.ArgumentParser()
+        expected.add_argument("--required_list", "--required-list", nargs="+", type=int, required=True)
+        expected.add_argument("--required_str", "--required-str", type=str, required=True)
+        expected.add_argument(
+            "--required_enum",
+            "--required-enum",
+            type=make_choice_type_function(["titi", "toto"]),
+            choices=["titi", "toto"],
+            required=True,
+        )
+        self.argparsersEqual(parser, expected)
+
+    def test_with_string_literal_annotation(self):
+        parser = HfArgumentParser(StringLiteralAnnotationExample)
+
+        expected = argparse.ArgumentParser()
+        expected.add_argument("--foo", type=int, required=True)
+        expected.add_argument(
+            "--required_enum",
+            "--required-enum",
+            type=make_choice_type_function(["titi", "toto"]),
+            choices=["titi", "toto"],
+            required=True,
+        )
+        expected.add_argument("--opt", type=string_to_bool, default=None)
+        expected.add_argument("--baz", default="toto", type=str, help="help message")
+        expected.add_argument("--foo_str", "--foo-str", nargs="+", default=["Hallo", "Bonjour", "Hello"], type=str)
+        self.argparsersEqual(parser, expected)
+
+    def test_parse_dict(self):
+        parser = HfArgumentParser(BasicExample)
+
+        args_dict = {
+            "foo": 12,
+            "bar": 3.14,
+            "baz": "42",
+            "flag": True,
+        }
+
+        parsed_args = parser.parse_dict(args_dict)[0]
+        args = BasicExample(**args_dict)
+        self.assertEqual(parsed_args, args)
+
+    def test_parse_dict_extra_key(self):
+        parser = HfArgumentParser(BasicExample)
+
+        args_dict = {
+            "foo": 12,
+            "bar": 3.14,
+            "baz": "42",
+            "flag": True,
+            "extra": 42,
+        }
+
+        self.assertRaises(ValueError, parser.parse_dict, args_dict, allow_extra_keys=False)
+
+    def test_parse_json(self):
+        parser = HfArgumentParser(BasicExample)
+
+        args_dict_for_json = {
+            "foo": 12,
+            "bar": 3.14,
+            "baz": "42",
+            "flag": True,
+        }
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            temp_local_path = os.path.join(tmp_dir, "temp_json")
+            os.mkdir(temp_local_path)
+            with open(temp_local_path + ".json", "w+") as f:
+                json.dump(args_dict_for_json, f)
+            parsed_args = parser.parse_json_file(Path(temp_local_path + ".json"))[0]
+
+        args = BasicExample(**args_dict_for_json)
+        self.assertEqual(parsed_args, args)
+
+    def test_parse_yaml(self):
+        parser = HfArgumentParser(BasicExample)
+
+        args_dict_for_yaml = {
+            "foo": 12,
+            "bar": 3.14,
+            "baz": "42",
+            "flag": True,
+        }
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            temp_local_path = os.path.join(tmp_dir, "temp_yaml")
+            os.mkdir(temp_local_path)
+            with open(temp_local_path + ".yaml", "w+") as f:
+                yaml.dump(args_dict_for_yaml, f)
+            parsed_args = parser.parse_yaml_file(Path(temp_local_path + ".yaml"))[0]
+        args = BasicExample(**args_dict_for_yaml)
+        self.assertEqual(parsed_args, args)
+
+    def test_z_integration_training_args(self):
+        # make sure that this test executes last in the test suite
+        parser = HfArgumentParser(TrainingArguments)
+        self.assertIsNotNone(parser)
+
+    def test_valid_dict_annotation(self):
+        """
+        Tests to make sure that `dict` based annotations
+        are correctly made in the `TrainingArguments`.
+
+        If this fails, a type annotation change is
+        needed on a new input
+        """
+        base_list = TrainingArguments._VALID_DICT_FIELDS.copy()
+        args = TrainingArguments
+
+        # First find any annotations that contain `dict`
+        fields = args.__dataclass_fields__
+
+        raw_dict_fields = []
+        optional_dict_fields = []
+
+        for field in fields.values():
+            # First verify raw dict
+            if field.type in (dict, dict):
+                raw_dict_fields.append(field)
+            # Next check for `Union` or `Optional`
+            elif get_origin(field.type) == Union:
+                if any(arg in (dict, dict) for arg in get_args(field.type)):
+                    optional_dict_fields.append(field)
+
+        # First check: anything in `raw_dict_fields` is very bad
+        self.assertEqual(
+            len(raw_dict_fields),
+            0,
+            "Found invalid raw `dict` types in the `TrainingArgument` typings. "
+            "This leads to issues with the CLI. Please turn this into `typing.Optional[dict,str]`",
+        )
+
+        # Next check raw annotations
+        for field in optional_dict_fields:
+            args = get_args(field.type)
+            # These should be returned as `dict`, `str`, ...
+            # we only care about the first two
+            self.assertIn(args[0], (dict, dict))
+            self.assertEqual(
+                str(args[1]),
+                "<class 'str'>",
+                f"Expected field `{field.name}` to have a type signature of at least `typing.Union[dict,str,...]` for CLI compatibility, "
+                "but `str` not found. Please fix this.",
+            )
+
+        # Second check: anything in `optional_dict_fields` is bad if it's not in `base_list`
+        for field in optional_dict_fields:
+            self.assertIn(
+                field.name,
+                base_list,
+                f"Optional dict field `{field.name}` is not in the base list of valid fields. Please add it to `TrainingArguments._VALID_DICT_FIELDS`",
+            )
+
+    @require_torch
+    def test_valid_dict_input_parsing(self):
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            args = TrainingArguments(
+                output_dir=tmp_dir,
+                accelerator_config='{"split_batches": "True", "gradient_accumulation_kwargs": {"num_steps": 2}}',
+            )
+            self.assertEqual(args.accelerator_config.split_batches, True)
+            self.assertEqual(args.accelerator_config.gradient_accumulation_kwargs["num_steps"], 2)

docs/transformers/tests/utils/test_hub_utils.py

@@ -0,0 +1,200 @@
+# Copyright 2020 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import json
+import os
+import tempfile
+import unittest
+import unittest.mock as mock
+from pathlib import Path
+
+from huggingface_hub import hf_hub_download
+from requests.exceptions import HTTPError
+
+from transformers.utils import (
+    CONFIG_NAME,
+    FLAX_WEIGHTS_NAME,
+    TF2_WEIGHTS_NAME,
+    TRANSFORMERS_CACHE,
+    WEIGHTS_NAME,
+    cached_file,
+    has_file,
+)
+
+
+RANDOM_BERT = "hf-internal-testing/tiny-random-bert"
+TINY_BERT_PT_ONLY = "hf-internal-testing/tiny-bert-pt-only"
+CACHE_DIR = os.path.join(TRANSFORMERS_CACHE, "models--hf-internal-testing--tiny-random-bert")
+FULL_COMMIT_HASH = "9b8c223d42b2188cb49d29af482996f9d0f3e5a6"
+
+GATED_REPO = "hf-internal-testing/dummy-gated-model"
+README_FILE = "README.md"
+
+
+class GetFromCacheTests(unittest.TestCase):
+    def test_cached_file(self):
+        archive_file = cached_file(RANDOM_BERT, CONFIG_NAME)
+        # Should have downloaded the file in here
+        self.assertTrue(os.path.isdir(CACHE_DIR))
+        # Cache should contain at least those three subfolders:
+        for subfolder in ["blobs", "refs", "snapshots"]:
+            self.assertTrue(os.path.isdir(os.path.join(CACHE_DIR, subfolder)))
+        with open(os.path.join(CACHE_DIR, "refs", "main")) as f:
+            main_commit = f.read()
+        self.assertEqual(archive_file, os.path.join(CACHE_DIR, "snapshots", main_commit, CONFIG_NAME))
+        self.assertTrue(os.path.isfile(archive_file))
+
+        # File is cached at the same place the second time.
+        new_archive_file = cached_file(RANDOM_BERT, CONFIG_NAME)
+        self.assertEqual(archive_file, new_archive_file)
+
+        # Using a specific revision to test the full commit hash.
+        archive_file = cached_file(RANDOM_BERT, CONFIG_NAME, revision="9b8c223")
+        self.assertEqual(archive_file, os.path.join(CACHE_DIR, "snapshots", FULL_COMMIT_HASH, CONFIG_NAME))
+
+    def test_cached_file_errors(self):
+        with self.assertRaisesRegex(EnvironmentError, "is not a valid model identifier"):
+            _ = cached_file("tiny-random-bert", CONFIG_NAME)
+
+        with self.assertRaisesRegex(EnvironmentError, "is not a valid git identifier"):
+            _ = cached_file(RANDOM_BERT, CONFIG_NAME, revision="aaaa")
+
+        with self.assertRaisesRegex(EnvironmentError, "does not appear to have a file named"):
+            _ = cached_file(RANDOM_BERT, "conf")
+
+    def test_non_existence_is_cached(self):
+        with self.assertRaisesRegex(EnvironmentError, "does not appear to have a file named"):
+            _ = cached_file(RANDOM_BERT, "conf")
+
+        with open(os.path.join(CACHE_DIR, "refs", "main")) as f:
+            main_commit = f.read()
+        self.assertTrue(os.path.isfile(os.path.join(CACHE_DIR, ".no_exist", main_commit, "conf")))
+
+        path = cached_file(RANDOM_BERT, "conf", _raise_exceptions_for_missing_entries=False)
+        self.assertIsNone(path)
+
+        path = cached_file(RANDOM_BERT, "conf", local_files_only=True, _raise_exceptions_for_missing_entries=False)
+        self.assertIsNone(path)
+
+        # Under the mock environment, hf_hub_download will always raise an HTTPError
+        with mock.patch("transformers.utils.hub.hf_hub_download", side_effect=HTTPError) as mock_head:
+            path = cached_file(RANDOM_BERT, "conf", _raise_exceptions_for_connection_errors=False)
+            self.assertIsNone(path)
+            # This check we did call the fake head request
+            mock_head.assert_called()
+
+    def test_has_file(self):
+        self.assertTrue(has_file(TINY_BERT_PT_ONLY, WEIGHTS_NAME))
+        self.assertFalse(has_file(TINY_BERT_PT_ONLY, TF2_WEIGHTS_NAME))
+        self.assertFalse(has_file(TINY_BERT_PT_ONLY, FLAX_WEIGHTS_NAME))
+
+    def test_has_file_in_cache(self):
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            # Empty cache dir + offline mode => return False
+            assert not has_file(TINY_BERT_PT_ONLY, WEIGHTS_NAME, local_files_only=True, cache_dir=tmp_dir)
+
+            # Populate cache dir
+            hf_hub_download(TINY_BERT_PT_ONLY, WEIGHTS_NAME, cache_dir=tmp_dir)
+
+            # Cache dir + offline mode => return True
+            assert has_file(TINY_BERT_PT_ONLY, WEIGHTS_NAME, local_files_only=True, cache_dir=tmp_dir)
+
+    def test_get_file_from_repo_distant(self):
+        # should return None if the file does not exist
+        self.assertIsNone(
+            cached_file(
+                "google-bert/bert-base-cased",
+                "ahah.txt",
+                _raise_exceptions_for_gated_repo=False,
+                _raise_exceptions_for_missing_entries=False,
+                _raise_exceptions_for_connection_errors=False,
+            )
+        )
+
+        # The function raises if the repository does not exist.
+        with self.assertRaisesRegex(EnvironmentError, "is not a valid model identifier"):
+            cached_file(
+                "bert-base-case",
+                CONFIG_NAME,
+                _raise_exceptions_for_gated_repo=False,
+                _raise_exceptions_for_missing_entries=False,
+                _raise_exceptions_for_connection_errors=False,
+            )
+
+        # The function raises if the revision does not exist.
+        with self.assertRaisesRegex(EnvironmentError, "is not a valid git identifier"):
+            cached_file(
+                "google-bert/bert-base-cased",
+                CONFIG_NAME,
+                revision="ahaha",
+                _raise_exceptions_for_gated_repo=False,
+                _raise_exceptions_for_missing_entries=False,
+                _raise_exceptions_for_connection_errors=False,
+            )
+
+        resolved_file = cached_file(
+            "google-bert/bert-base-cased",
+            CONFIG_NAME,
+            _raise_exceptions_for_gated_repo=False,
+            _raise_exceptions_for_missing_entries=False,
+            _raise_exceptions_for_connection_errors=False,
+        )
+        # The name is the cached name which is not very easy to test, so instead we load the content.
+        config = json.loads(open(resolved_file).read())
+        self.assertEqual(config["hidden_size"], 768)
+
+    def test_get_file_from_repo_local(self):
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            filename = Path(tmp_dir) / "a.txt"
+            filename.touch()
+            self.assertEqual(
+                cached_file(
+                    tmp_dir,
+                    "a.txt",
+                    _raise_exceptions_for_gated_repo=False,
+                    _raise_exceptions_for_missing_entries=False,
+                    _raise_exceptions_for_connection_errors=False,
+                ),
+                str(filename),
+            )
+
+            self.assertIsNone(
+                cached_file(
+                    tmp_dir,
+                    "b.txt",
+                    _raise_exceptions_for_gated_repo=False,
+                    _raise_exceptions_for_missing_entries=False,
+                    _raise_exceptions_for_connection_errors=False,
+                )
+            )
+
+    def test_get_file_gated_repo(self):
+        """Test download file from a gated repo fails with correct message when not authenticated."""
+        with self.assertRaisesRegex(EnvironmentError, "You are trying to access a gated repo."):
+            # All files except README.md are protected on a gated repo.
+            cached_file(GATED_REPO, "gated_file.txt", token=False)
+
+    def test_has_file_gated_repo(self):
+        """Test check file existence from a gated repo fails with correct message when not authenticated."""
+        with self.assertRaisesRegex(EnvironmentError, "is a gated repository"):
+            # All files except README.md are protected on a gated repo.
+            has_file(GATED_REPO, "gated_file.txt", token=False)
+
+    def test_cached_files_exception_raised(self):
+        """Test that unhadled exceptions, e.g. ModuleNotFoundError, is properly re-raised by cached_files when hf_hub_download fails."""
+        with mock.patch(
+            "transformers.utils.hub.hf_hub_download", side_effect=ModuleNotFoundError("No module named 'MockModule'")
+        ):
+            with self.assertRaises(ModuleNotFoundError):
+                # The error should be re-raised by cached_files, not caught in the exception handling block
+                cached_file(RANDOM_BERT, "nonexistent.json")

docs/transformers/tests/utils/test_image_processing_utils.py

@@ -0,0 +1,181 @@
+# Copyright 2024 HuggingFace Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import sys
+import tempfile
+import unittest
+import unittest.mock as mock
+from pathlib import Path
+
+from huggingface_hub import HfFolder
+from requests.exceptions import HTTPError
+
+from transformers import AutoImageProcessor, ViTImageProcessor
+from transformers.image_processing_utils import get_size_dict
+from transformers.testing_utils import TOKEN, TemporaryHubRepo, get_tests_dir, is_staging_test
+
+
+sys.path.append(str(Path(__file__).parent.parent.parent / "utils"))
+
+from test_module.custom_image_processing import CustomImageProcessor  # noqa E402
+
+
+SAMPLE_IMAGE_PROCESSING_CONFIG_DIR = get_tests_dir("fixtures")
+
+
+class ImageProcessorUtilTester(unittest.TestCase):
+    def test_cached_files_are_used_when_internet_is_down(self):
+        # A mock response for an HTTP head request to emulate server down
+        response_mock = mock.Mock()
+        response_mock.status_code = 500
+        response_mock.headers = {}
+        response_mock.raise_for_status.side_effect = HTTPError
+        response_mock.json.return_value = {}
+
+        # Download this model to make sure it's in the cache.
+        _ = ViTImageProcessor.from_pretrained("hf-internal-testing/tiny-random-vit")
+        # Under the mock environment we get a 500 error when trying to reach the model.
+        with mock.patch("requests.Session.request", return_value=response_mock) as mock_head:
+            _ = ViTImageProcessor.from_pretrained("hf-internal-testing/tiny-random-vit")
+            # This check we did call the fake head request
+            mock_head.assert_called()
+
+    def test_image_processor_from_pretrained_subfolder(self):
+        with self.assertRaises(OSError):
+            # config is in subfolder, the following should not work without specifying the subfolder
+            _ = AutoImageProcessor.from_pretrained("hf-internal-testing/stable-diffusion-all-variants")
+
+        config = AutoImageProcessor.from_pretrained(
+            "hf-internal-testing/stable-diffusion-all-variants", subfolder="feature_extractor"
+        )
+
+        self.assertIsNotNone(config)
+
+
+@is_staging_test
+class ImageProcessorPushToHubTester(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls._token = TOKEN
+        HfFolder.save_token(TOKEN)
+
+    def test_push_to_hub(self):
+        with TemporaryHubRepo(token=self._token) as tmp_repo:
+            image_processor = ViTImageProcessor.from_pretrained(SAMPLE_IMAGE_PROCESSING_CONFIG_DIR)
+            image_processor.push_to_hub(tmp_repo.repo_id, token=self._token)
+
+            new_image_processor = ViTImageProcessor.from_pretrained(tmp_repo.repo_id)
+            for k, v in image_processor.__dict__.items():
+                self.assertEqual(v, getattr(new_image_processor, k))
+
+    def test_push_to_hub_via_save_pretrained(self):
+        with TemporaryHubRepo(token=self._token) as tmp_repo:
+            image_processor = ViTImageProcessor.from_pretrained(SAMPLE_IMAGE_PROCESSING_CONFIG_DIR)
+            # Push to hub via save_pretrained
+            with tempfile.TemporaryDirectory() as tmp_dir:
+                image_processor.save_pretrained(tmp_dir, repo_id=tmp_repo.repo_id, push_to_hub=True, token=self._token)
+
+            new_image_processor = ViTImageProcessor.from_pretrained(tmp_repo.repo_id)
+            for k, v in image_processor.__dict__.items():
+                self.assertEqual(v, getattr(new_image_processor, k))
+
+    def test_push_to_hub_in_organization(self):
+        with TemporaryHubRepo(namespace="valid_org", token=self._token) as tmp_repo:
+            image_processor = ViTImageProcessor.from_pretrained(SAMPLE_IMAGE_PROCESSING_CONFIG_DIR)
+            image_processor.push_to_hub(tmp_repo.repo_id, token=self._token)
+
+            new_image_processor = ViTImageProcessor.from_pretrained(tmp_repo.repo_id)
+            for k, v in image_processor.__dict__.items():
+                self.assertEqual(v, getattr(new_image_processor, k))
+
+    def test_push_to_hub_in_organization_via_save_pretrained(self):
+        with TemporaryHubRepo(namespace="valid_org", token=self._token) as tmp_repo:
+            image_processor = ViTImageProcessor.from_pretrained(SAMPLE_IMAGE_PROCESSING_CONFIG_DIR)
+            # Push to hub via save_pretrained
+            with tempfile.TemporaryDirectory() as tmp_dir:
+                image_processor.save_pretrained(tmp_dir, repo_id=tmp_repo.repo_id, push_to_hub=True, token=self._token)
+
+            new_image_processor = ViTImageProcessor.from_pretrained(tmp_repo.repo_id)
+            for k, v in image_processor.__dict__.items():
+                self.assertEqual(v, getattr(new_image_processor, k))
+
+    def test_push_to_hub_dynamic_image_processor(self):
+        with TemporaryHubRepo(token=self._token) as tmp_repo:
+            CustomImageProcessor.register_for_auto_class()
+            image_processor = CustomImageProcessor.from_pretrained(SAMPLE_IMAGE_PROCESSING_CONFIG_DIR)
+
+            image_processor.push_to_hub(tmp_repo.repo_id, token=self._token)
+
+            # This has added the proper auto_map field to the config
+            self.assertDictEqual(
+                image_processor.auto_map,
+                {"AutoImageProcessor": "custom_image_processing.CustomImageProcessor"},
+            )
+
+            new_image_processor = AutoImageProcessor.from_pretrained(tmp_repo.repo_id, trust_remote_code=True)
+            # Can't make an isinstance check because the new_image_processor is from the CustomImageProcessor class of a dynamic module
+            self.assertEqual(new_image_processor.__class__.__name__, "CustomImageProcessor")
+
+
+class ImageProcessingUtilsTester(unittest.TestCase):
+    def test_get_size_dict(self):
+        # Test a dict with the wrong keys raises an error
+        inputs = {"wrong_key": 224}
+        with self.assertRaises(ValueError):
+            get_size_dict(inputs)
+
+        inputs = {"height": 224}
+        with self.assertRaises(ValueError):
+            get_size_dict(inputs)
+
+        inputs = {"width": 224, "shortest_edge": 224}
+        with self.assertRaises(ValueError):
+            get_size_dict(inputs)
+
+        # Test a dict with the correct keys is returned as is
+        inputs = {"height": 224, "width": 224}
+        outputs = get_size_dict(inputs)
+        self.assertEqual(outputs, inputs)
+
+        inputs = {"shortest_edge": 224}
+        outputs = get_size_dict(inputs)
+        self.assertEqual(outputs, {"shortest_edge": 224})
+
+        inputs = {"longest_edge": 224, "shortest_edge": 224}
+        outputs = get_size_dict(inputs)
+        self.assertEqual(outputs, {"longest_edge": 224, "shortest_edge": 224})
+
+        # Test a single int value which  represents (size, size)
+        outputs = get_size_dict(224)
+        self.assertEqual(outputs, {"height": 224, "width": 224})
+
+        # Test a single int value which represents the shortest edge
+        outputs = get_size_dict(224, default_to_square=False)
+        self.assertEqual(outputs, {"shortest_edge": 224})
+
+        # Test a tuple of ints which represents (height, width)
+        outputs = get_size_dict((150, 200))
+        self.assertEqual(outputs, {"height": 150, "width": 200})
+
+        # Test a tuple of ints which represents (width, height)
+        outputs = get_size_dict((150, 200), height_width_order=False)
+        self.assertEqual(outputs, {"height": 200, "width": 150})
+
+        # Test an int representing the shortest edge and max_size which represents the longest edge
+        outputs = get_size_dict(224, max_size=256, default_to_square=False)
+        self.assertEqual(outputs, {"shortest_edge": 224, "longest_edge": 256})
+
+        # Test int with default_to_square=True and max_size fails
+        with self.assertRaises(ValueError):
+            get_size_dict(224, max_size=256, default_to_square=True)

docs/transformers/tests/utils/test_image_utils.py

@@ -0,0 +1,1061 @@
+# Copyright 2021 HuggingFace Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import codecs
+import os
+import tempfile
+import unittest
+from io import BytesIO
+from typing import Optional
+
+import numpy as np
+import pytest
+import requests
+from huggingface_hub.file_download import hf_hub_url, http_get
+from requests import ConnectTimeout, ReadTimeout
+
+from tests.pipelines.test_pipelines_document_question_answering import INVOICE_URL
+from transformers import is_torch_available, is_vision_available
+from transformers.image_utils import (
+    ChannelDimension,
+    get_channel_dimension_axis,
+    make_batched_videos,
+    make_flat_list_of_images,
+    make_list_of_images,
+    make_nested_list_of_images,
+)
+from transformers.testing_utils import is_flaky, require_torch, require_vision
+
+
+if is_torch_available():
+    import torch
+
+if is_vision_available():
+    import PIL.Image
+
+    from transformers import ImageFeatureExtractionMixin
+    from transformers.image_utils import get_image_size, infer_channel_dimension_format, load_image
+
+
+def get_image_from_hub_dataset(dataset_id: str, filename: str, revision: Optional[str] = None) -> "PIL.Image.Image":
+    url = hf_hub_url(dataset_id, filename, repo_type="dataset", revision=revision)
+    return PIL.Image.open(BytesIO(requests.get(url).content))
+
+
+def get_random_image(height, width):
+    random_array = np.random.randint(0, 256, (height, width, 3), dtype=np.uint8)
+    return PIL.Image.fromarray(random_array)
+
+
+@require_vision
+class ImageFeatureExtractionTester(unittest.TestCase):
+    def test_conversion_image_to_array(self):
+        feature_extractor = ImageFeatureExtractionMixin()
+        image = get_random_image(16, 32)
+
+        # Conversion with defaults (rescale + channel first)
+        array1 = feature_extractor.to_numpy_array(image)
+        self.assertTrue(array1.dtype, np.float32)
+        self.assertEqual(array1.shape, (3, 16, 32))
+
+        # Conversion with rescale and not channel first
+        array2 = feature_extractor.to_numpy_array(image, channel_first=False)
+        self.assertTrue(array2.dtype, np.float32)
+        self.assertEqual(array2.shape, (16, 32, 3))
+        self.assertTrue(np.array_equal(array1, array2.transpose(2, 0, 1)))
+
+        # Conversion with no rescale and channel first
+        array3 = feature_extractor.to_numpy_array(image, rescale=False)
+        self.assertTrue(array3.dtype, np.uint8)
+        self.assertEqual(array3.shape, (3, 16, 32))
+        self.assertTrue(np.array_equal(array1, array3.astype(np.float32) * (1 / 255.0)))
+
+        # Conversion with no rescale and not channel first
+        array4 = feature_extractor.to_numpy_array(image, rescale=False, channel_first=False)
+        self.assertTrue(array4.dtype, np.uint8)
+        self.assertEqual(array4.shape, (16, 32, 3))
+        self.assertTrue(np.array_equal(array2, array4.astype(np.float32) * (1 / 255.0)))
+
+    def test_conversion_array_to_array(self):
+        feature_extractor = ImageFeatureExtractionMixin()
+        array = np.random.randint(0, 256, (16, 32, 3), dtype=np.uint8)
+
+        # By default, rescale (for an array of ints) and channel permute
+        array1 = feature_extractor.to_numpy_array(array)
+        self.assertTrue(array1.dtype, np.float32)
+        self.assertEqual(array1.shape, (3, 16, 32))
+        self.assertTrue(np.array_equal(array1, array.transpose(2, 0, 1).astype(np.float32) * (1 / 255.0)))
+
+        # Same with no permute
+        array2 = feature_extractor.to_numpy_array(array, channel_first=False)
+        self.assertTrue(array2.dtype, np.float32)
+        self.assertEqual(array2.shape, (16, 32, 3))
+        self.assertTrue(np.array_equal(array2, array.astype(np.float32) * (1 / 255.0)))
+
+        # Force rescale to False
+        array3 = feature_extractor.to_numpy_array(array, rescale=False)
+        self.assertTrue(array3.dtype, np.uint8)
+        self.assertEqual(array3.shape, (3, 16, 32))
+        self.assertTrue(np.array_equal(array3, array.transpose(2, 0, 1)))
+
+        # Force rescale to False and no channel permute
+        array4 = feature_extractor.to_numpy_array(array, rescale=False, channel_first=False)
+        self.assertTrue(array4.dtype, np.uint8)
+        self.assertEqual(array4.shape, (16, 32, 3))
+        self.assertTrue(np.array_equal(array4, array))
+
+        # Now test the default rescale for a float array (defaults to False)
+        array5 = feature_extractor.to_numpy_array(array2)
+        self.assertTrue(array5.dtype, np.float32)
+        self.assertEqual(array5.shape, (3, 16, 32))
+        self.assertTrue(np.array_equal(array5, array1))
+
+    def test_make_list_of_images_pil(self):
+        # Test a single image is converted to a list of 1 image
+        pil_image = get_random_image(16, 32)
+        images_list = make_list_of_images(pil_image)
+        self.assertIsInstance(images_list, list)
+        self.assertEqual(len(images_list), 1)
+        self.assertIsInstance(images_list[0], PIL.Image.Image)
+
+        # Test a list of images is not modified
+        images = [get_random_image(16, 32) for _ in range(4)]
+        images_list = make_list_of_images(images)
+        self.assertIsInstance(images_list, list)
+        self.assertEqual(len(images_list), 4)
+        self.assertIsInstance(images_list[0], PIL.Image.Image)
+
+    def test_make_list_of_images_numpy(self):
+        # Test a single image is converted to a list of 1 image
+        images = np.random.randint(0, 256, (16, 32, 3))
+        images_list = make_list_of_images(images)
+        self.assertEqual(len(images_list), 1)
+        self.assertTrue(np.array_equal(images_list[0], images))
+        self.assertIsInstance(images_list, list)
+
+        # Test a batch of images is converted to a list of images
+        images = np.random.randint(0, 256, (4, 16, 32, 3))
+        images_list = make_list_of_images(images)
+        self.assertEqual(len(images_list), 4)
+        self.assertTrue(np.array_equal(images_list[0], images[0]))
+        self.assertIsInstance(images_list, list)
+
+        # Test a list of images is not modified
+        images = [np.random.randint(0, 256, (16, 32, 3)) for _ in range(4)]
+        images_list = make_list_of_images(images)
+        self.assertEqual(len(images_list), 4)
+        self.assertTrue(np.array_equal(images_list[0], images[0]))
+        self.assertIsInstance(images_list, list)
+
+        # Test batched masks with no channel dimension are converted to a list of masks
+        masks = np.random.randint(0, 2, (4, 16, 32))
+        masks_list = make_list_of_images(masks, expected_ndims=2)
+        self.assertEqual(len(masks_list), 4)
+        self.assertTrue(np.array_equal(masks_list[0], masks[0]))
+        self.assertIsInstance(masks_list, list)
+
+    @require_torch
+    def test_make_list_of_images_torch(self):
+        # Test a single image is converted to a list of 1 image
+        images = torch.randint(0, 256, (16, 32, 3))
+        images_list = make_list_of_images(images)
+        self.assertEqual(len(images_list), 1)
+        self.assertTrue(np.array_equal(images_list[0], images))
+        self.assertIsInstance(images_list, list)
+
+        # Test a batch of images is converted to a list of images
+        images = torch.randint(0, 256, (4, 16, 32, 3))
+        images_list = make_list_of_images(images)
+        self.assertEqual(len(images_list), 4)
+        self.assertTrue(np.array_equal(images_list[0], images[0]))
+        self.assertIsInstance(images_list, list)
+
+        # Test a list of images is left unchanged
+        images = [torch.randint(0, 256, (16, 32, 3)) for _ in range(4)]
+        images_list = make_list_of_images(images)
+        self.assertEqual(len(images_list), 4)
+        self.assertTrue(np.array_equal(images_list[0], images[0]))
+        self.assertIsInstance(images_list, list)
+
+    def test_make_flat_list_of_images_pil(self):
+        # Test a single image is converted to a list of 1 image
+        pil_image = get_random_image(16, 32)
+        images_list = make_flat_list_of_images(pil_image)
+        self.assertIsInstance(images_list, list)
+        self.assertEqual(len(images_list), 1)
+        self.assertIsInstance(images_list[0], PIL.Image.Image)
+
+        # Test a list of images is not modified
+        images = [get_random_image(16, 32) for _ in range(4)]
+        images_list = make_flat_list_of_images(images)
+        self.assertIsInstance(images_list, list)
+        self.assertEqual(len(images_list), 4)
+        self.assertIsInstance(images_list[0], PIL.Image.Image)
+
+        # Test a nested list of images is flattened
+        images = [[get_random_image(16, 32) for _ in range(2)] for _ in range(2)]
+        images_list = make_flat_list_of_images(images)
+        self.assertIsInstance(images_list, list)
+        self.assertEqual(len(images_list), 4)
+        self.assertIsInstance(images_list[0], PIL.Image.Image)
+
+    def test_make_flat_list_of_images_numpy(self):
+        # Test a single image is converted to a list of 1 image
+        images = np.random.randint(0, 256, (16, 32, 3))
+        images_list = make_flat_list_of_images(images)
+        self.assertEqual(len(images_list), 1)
+        self.assertTrue(np.array_equal(images_list[0], images))
+        self.assertIsInstance(images_list, list)
+
+        # Test a 4d array of images is changed to a list of images
+        images = np.random.randint(0, 256, (4, 16, 32, 3))
+        images_list = make_flat_list_of_images(images)
+        self.assertEqual(len(images_list), 4)
+        self.assertIsInstance(images_list, list)
+        self.assertIsInstance(images_list[0], np.ndarray)
+        self.assertTrue(np.array_equal(images_list[0], images[0]))
+
+        # Test a list of images is not modified
+        images = [np.random.randint(0, 256, (16, 32, 3)) for _ in range(4)]
+        images_list = make_flat_list_of_images(images)
+        self.assertEqual(len(images_list), 4)
+        self.assertTrue(np.array_equal(images_list[0], images[0]))
+        self.assertIsInstance(images_list, list)
+
+        # Test list of 4d array images is flattened
+        images = [np.random.randint(0, 256, (4, 16, 32, 3)) for _ in range(2)]
+        images_list = make_flat_list_of_images(images)
+        self.assertEqual(len(images_list), 8)
+        self.assertTrue(np.array_equal(images_list[0], images[0][0]))
+        self.assertIsInstance(images_list, list)
+        self.assertIsInstance(images_list[0], np.ndarray)
+
+        # Test nested list of images is flattened
+        images = [[np.random.randint(0, 256, (16, 32, 3)) for _ in range(2)] for _ in range(2)]
+        images_list = make_flat_list_of_images(images)
+        self.assertEqual(len(images_list), 4)
+        self.assertTrue(np.array_equal(images_list[0], images[0][0]))
+        self.assertIsInstance(images_list, list)
+
+    @require_torch
+    def test_make_flat_list_of_images_torch(self):
+        # Test a single image is converted to a list of 1 image
+        images = torch.randint(0, 256, (16, 32, 3))
+        images_list = make_flat_list_of_images(images)
+        self.assertEqual(len(images_list), 1)
+        self.assertTrue(np.array_equal(images_list[0], images))
+        self.assertIsInstance(images_list, list)
+
+        # Test a 4d tensors of images is changed to a list of images
+        images = torch.randint(0, 256, (4, 16, 32, 3))
+        images_list = make_flat_list_of_images(images)
+        self.assertEqual(len(images_list), 4)
+        self.assertIsInstance(images_list, list)
+        self.assertIsInstance(images_list[0], torch.Tensor)
+        self.assertTrue(np.array_equal(images_list[0], images[0]))
+
+        # Test a list of images is not modified
+        images = [torch.randint(0, 256, (16, 32, 3)) for _ in range(4)]
+        images_list = make_flat_list_of_images(images)
+        self.assertEqual(len(images_list), 4)
+        self.assertTrue(np.array_equal(images_list[0], images[0]))
+        self.assertIsInstance(images_list, list)
+
+        # Test list of 4d tensors of imagess is flattened
+        images = [torch.randint(0, 256, (4, 16, 32, 3)) for _ in range(2)]
+        images_list = make_flat_list_of_images(images)
+        self.assertEqual(len(images_list), 8)
+        self.assertTrue(np.array_equal(images_list[0], images[0][0]))
+        self.assertIsInstance(images_list, list)
+        self.assertIsInstance(images_list[0], torch.Tensor)
+
+        # Test nested list of images is flattened
+        images = [[torch.randint(0, 256, (16, 32, 3)) for _ in range(2)] for _ in range(2)]
+        images_list = make_flat_list_of_images(images)
+        self.assertEqual(len(images_list), 4)
+        self.assertTrue(np.array_equal(images_list[0], images[0][0]))
+        self.assertIsInstance(images_list, list)
+
+    def test_make_nested_list_of_images_pil(self):
+        # Test a single image is converted to a nested list of 1 image
+        pil_image = get_random_image(16, 32)
+        images_list = make_nested_list_of_images(pil_image)
+        self.assertIsInstance(images_list[0], list)
+        self.assertEqual(len(images_list[0]), 1)
+        self.assertIsInstance(images_list[0][0], PIL.Image.Image)
+
+        # Test a list of images is converted to a nested list of images
+        images = [get_random_image(16, 32) for _ in range(4)]
+        images_list = make_nested_list_of_images(images)
+        self.assertIsInstance(images_list[0], list)
+        self.assertEqual(len(images_list), 1)
+        self.assertEqual(len(images_list[0]), 4)
+        self.assertIsInstance(images_list[0][0], PIL.Image.Image)
+
+        # Test a nested list of images is not modified
+        images = [[get_random_image(16, 32) for _ in range(2)] for _ in range(2)]
+        images_list = make_nested_list_of_images(images)
+        self.assertIsInstance(images_list[0], list)
+        self.assertEqual(len(images_list), 2)
+        self.assertEqual(len(images_list[0]), 2)
+        self.assertIsInstance(images_list[0][0], PIL.Image.Image)
+
+    def test_make_nested_list_of_images_numpy(self):
+        # Test a single image is converted to a nested list of 1 image
+        images = np.random.randint(0, 256, (16, 32, 3))
+        images_list = make_nested_list_of_images(images)
+        self.assertIsInstance(images_list[0], list)
+        self.assertEqual(len(images_list), 1)
+        self.assertTrue(np.array_equal(images_list[0][0], images))
+
+        # Test a 4d array of images is converted to a nested list of images
+        images = np.random.randint(0, 256, (4, 16, 32, 3))
+        images_list = make_nested_list_of_images(images)
+        self.assertIsInstance(images_list[0], list)
+        self.assertIsInstance(images_list[0][0], np.ndarray)
+        self.assertEqual(len(images_list), 1)
+        self.assertEqual(len(images_list[0]), 4)
+        self.assertTrue(np.array_equal(images_list[0][0], images[0]))
+
+        # Test a list of images is converted to a nested list of images
+        images = [np.random.randint(0, 256, (16, 32, 3)) for _ in range(4)]
+        images_list = make_nested_list_of_images(images)
+        self.assertIsInstance(images_list[0], list)
+        self.assertEqual(len(images_list), 1)
+        self.assertEqual(len(images_list[0]), 4)
+        self.assertTrue(np.array_equal(images_list[0][0], images[0]))
+
+        # Test a nested list of images is left unchanged
+        images = [[np.random.randint(0, 256, (16, 32, 3)) for _ in range(2)] for _ in range(2)]
+        images_list = make_nested_list_of_images(images)
+        self.assertIsInstance(images_list[0], list)
+        self.assertEqual(len(images_list), 2)
+        self.assertEqual(len(images_list[0]), 2)
+        self.assertTrue(np.array_equal(images_list[0][0], images[0][0]))
+
+        # Test a list of 4d array images is converted to a nested list of images
+        images = [np.random.randint(0, 256, (4, 16, 32, 3)) for _ in range(2)]
+        images_list = make_nested_list_of_images(images)
+        self.assertIsInstance(images_list[0], list)
+        self.assertIsInstance(images_list[0][0], np.ndarray)
+        self.assertEqual(len(images_list), 2)
+        self.assertEqual(len(images_list[0]), 4)
+        self.assertTrue(np.array_equal(images_list[0][0], images[0][0]))
+
+    @require_torch
+    def test_make_nested_list_of_images_torch(self):
+        # Test a single image is converted to a nested list of 1 image
+        images = torch.randint(0, 256, (16, 32, 3))
+        images_list = make_nested_list_of_images(images)
+        self.assertIsInstance(images_list[0], list)
+        self.assertEqual(len(images_list[0]), 1)
+        self.assertTrue(np.array_equal(images_list[0][0], images))
+
+        # Test a 4d tensor of images is converted to a nested list of images
+        images = torch.randint(0, 256, (4, 16, 32, 3))
+        images_list = make_nested_list_of_images(images)
+        self.assertIsInstance(images_list[0], list)
+        self.assertIsInstance(images_list[0][0], torch.Tensor)
+        self.assertEqual(len(images_list), 1)
+        self.assertEqual(len(images_list[0]), 4)
+        self.assertTrue(np.array_equal(images_list[0][0], images[0]))
+
+        # Test a list of images is converted to a nested list of images
+        images = [torch.randint(0, 256, (16, 32, 3)) for _ in range(4)]
+        images_list = make_nested_list_of_images(images)
+        self.assertIsInstance(images_list[0], list)
+        self.assertEqual(len(images_list), 1)
+        self.assertEqual(len(images_list[0]), 4)
+        self.assertTrue(np.array_equal(images_list[0][0], images[0]))
+
+        # Test a nested list of images is left unchanged
+        images = [[torch.randint(0, 256, (16, 32, 3)) for _ in range(2)] for _ in range(2)]
+        images_list = make_nested_list_of_images(images)
+        self.assertIsInstance(images_list[0], list)
+        self.assertEqual(len(images_list), 2)
+        self.assertEqual(len(images_list[0]), 2)
+        self.assertTrue(np.array_equal(images_list[0][0], images[0][0]))
+
+        # Test a list of 4d tensor images is converted to a nested list of images
+        images = [torch.randint(0, 256, (4, 16, 32, 3)) for _ in range(2)]
+        images_list = make_nested_list_of_images(images)
+        self.assertIsInstance(images_list[0], list)
+        self.assertIsInstance(images_list[0][0], torch.Tensor)
+        self.assertEqual(len(images_list), 2)
+        self.assertEqual(len(images_list[0]), 4)
+        self.assertTrue(np.array_equal(images_list[0][0], images[0][0]))
+
+    def test_make_batched_videos_pil(self):
+        # Test a single image is converted to a list of 1 video with 1 frame
+        pil_image = get_random_image(16, 32)
+        videos_list = make_batched_videos(pil_image)
+        self.assertIsInstance(videos_list[0], list)
+        self.assertEqual(len(videos_list[0]), 1)
+        self.assertIsInstance(videos_list[0][0], PIL.Image.Image)
+
+        # Test a list of images is converted to a list of 1 video
+        images = [get_random_image(16, 32) for _ in range(4)]
+        videos_list = make_batched_videos(images)
+        self.assertIsInstance(videos_list[0], list)
+        self.assertEqual(len(videos_list), 1)
+        self.assertEqual(len(videos_list[0]), 4)
+        self.assertIsInstance(videos_list[0][0], PIL.Image.Image)
+
+        # Test a nested list of images is not modified
+        images = [[get_random_image(16, 32) for _ in range(2)] for _ in range(2)]
+        videos_list = make_nested_list_of_images(images)
+        self.assertIsInstance(videos_list[0], list)
+        self.assertEqual(len(videos_list), 2)
+        self.assertEqual(len(videos_list[0]), 2)
+        self.assertIsInstance(videos_list[0][0], PIL.Image.Image)
+
+    def test_make_batched_videos_numpy(self):
+        # Test a single image is converted to a list of 1 video with 1 frame
+        images = np.random.randint(0, 256, (16, 32, 3))
+        videos_list = make_batched_videos(images)
+        self.assertIsInstance(videos_list[0], list)
+        self.assertEqual(len(videos_list), 1)
+        self.assertTrue(np.array_equal(videos_list[0][0], images))
+
+        # Test a 4d array of images is converted to a list of 1 video
+        images = np.random.randint(0, 256, (4, 16, 32, 3))
+        videos_list = make_batched_videos(images)
+        self.assertIsInstance(videos_list[0], list)
+        self.assertIsInstance(videos_list[0][0], np.ndarray)
+        self.assertEqual(len(videos_list), 1)
+        self.assertEqual(len(videos_list[0]), 4)
+        self.assertTrue(np.array_equal(videos_list[0][0], images[0]))
+
+        # Test a list of images is converted to a list of videos
+        images = [np.random.randint(0, 256, (16, 32, 3)) for _ in range(4)]
+        videos_list = make_batched_videos(images)
+        self.assertIsInstance(videos_list[0], list)
+        self.assertEqual(len(videos_list), 1)
+        self.assertEqual(len(videos_list[0]), 4)
+        self.assertTrue(np.array_equal(videos_list[0][0], images[0]))
+
+        # Test a nested list of images is left unchanged
+        images = [[np.random.randint(0, 256, (16, 32, 3)) for _ in range(2)] for _ in range(2)]
+        videos_list = make_batched_videos(images)
+        self.assertIsInstance(videos_list[0], list)
+        self.assertEqual(len(videos_list), 2)
+        self.assertEqual(len(videos_list[0]), 2)
+        self.assertTrue(np.array_equal(videos_list[0][0], images[0][0]))
+
+        # Test a list of 4d array images is converted to a list of videos
+        images = [np.random.randint(0, 256, (4, 16, 32, 3)) for _ in range(2)]
+        videos_list = make_batched_videos(images)
+        self.assertIsInstance(videos_list[0], list)
+        self.assertIsInstance(videos_list[0][0], np.ndarray)
+        self.assertEqual(len(videos_list), 2)
+        self.assertEqual(len(videos_list[0]), 4)
+        self.assertTrue(np.array_equal(videos_list[0][0], images[0][0]))
+
+        # Test a batch of list of 4d array images is converted to a list of videos
+        images = [[np.random.randint(0, 256, (4, 16, 32, 3)) for _ in range(2)] for _ in range(2)]
+        videos_list = make_batched_videos(images)
+        self.assertIsInstance(videos_list[0], list)
+        self.assertIsInstance(videos_list[0][0], np.ndarray)
+        self.assertEqual(len(videos_list), 2)
+        self.assertEqual(len(videos_list[0]), 8)
+        self.assertTrue(np.array_equal(videos_list[0][0], images[0][0][0]))
+
+    @require_torch
+    def test_make_batched_videos_torch(self):
+        # Test a single image is converted to a list of 1 video with 1 frame
+        images = torch.randint(0, 256, (16, 32, 3))
+        videos_list = make_batched_videos(images)
+        self.assertIsInstance(videos_list[0], list)
+        self.assertEqual(len(videos_list[0]), 1)
+        self.assertTrue(np.array_equal(videos_list[0][0], images))
+
+        # Test a 4d tensor of images is converted to a list of 1 video
+        images = torch.randint(0, 256, (4, 16, 32, 3))
+        videos_list = make_batched_videos(images)
+        self.assertIsInstance(videos_list[0], list)
+        self.assertIsInstance(videos_list[0][0], torch.Tensor)
+        self.assertEqual(len(videos_list), 1)
+        self.assertEqual(len(videos_list[0]), 4)
+        self.assertTrue(np.array_equal(videos_list[0][0], images[0]))
+
+        # Test a list of images is converted to a list of videos
+        images = [torch.randint(0, 256, (16, 32, 3)) for _ in range(4)]
+        videos_list = make_batched_videos(images)
+        self.assertIsInstance(videos_list[0], list)
+        self.assertEqual(len(videos_list), 1)
+        self.assertEqual(len(videos_list[0]), 4)
+        self.assertTrue(np.array_equal(videos_list[0][0], images[0]))
+
+        # Test a nested list of images is left unchanged
+        images = [[torch.randint(0, 256, (16, 32, 3)) for _ in range(2)] for _ in range(2)]
+        videos_list = make_batched_videos(images)
+        self.assertIsInstance(videos_list[0], list)
+        self.assertEqual(len(videos_list), 2)
+        self.assertEqual(len(videos_list[0]), 2)
+        self.assertTrue(np.array_equal(videos_list[0][0], images[0][0]))
+
+        # Test a list of 4d tensor images is converted to a list of videos
+        images = [torch.randint(0, 256, (4, 16, 32, 3)) for _ in range(2)]
+        videos_list = make_batched_videos(images)
+        self.assertIsInstance(videos_list[0], list)
+        self.assertIsInstance(videos_list[0][0], torch.Tensor)
+        self.assertEqual(len(videos_list), 2)
+        self.assertEqual(len(videos_list[0]), 4)
+        self.assertTrue(np.array_equal(videos_list[0][0], images[0][0]))
+
+        # Test a batch of list of 4d tensor images is converted to a list of videos
+        images = [[torch.randint(0, 256, (4, 16, 32, 3)) for _ in range(2)] for _ in range(2)]
+        videos_list = make_batched_videos(images)
+        self.assertIsInstance(videos_list[0], list)
+        self.assertIsInstance(videos_list[0][0], torch.Tensor)
+        self.assertEqual(len(videos_list), 2)
+        self.assertEqual(len(videos_list[0]), 8)
+        self.assertTrue(np.array_equal(videos_list[0][0], images[0][0][0]))
+
+    @require_torch
+    def test_conversion_torch_to_array(self):
+        feature_extractor = ImageFeatureExtractionMixin()
+        tensor = torch.randint(0, 256, (16, 32, 3))
+        array = tensor.numpy()
+
+        # By default, rescale (for a tensor of ints) and channel permute
+        array1 = feature_extractor.to_numpy_array(array)
+        self.assertTrue(array1.dtype, np.float32)
+        self.assertEqual(array1.shape, (3, 16, 32))
+        self.assertTrue(np.array_equal(array1, array.transpose(2, 0, 1).astype(np.float32) * (1 / 255.0)))
+
+        # Same with no permute
+        array2 = feature_extractor.to_numpy_array(array, channel_first=False)
+        self.assertTrue(array2.dtype, np.float32)
+        self.assertEqual(array2.shape, (16, 32, 3))
+        self.assertTrue(np.array_equal(array2, array.astype(np.float32) * (1 / 255.0)))
+
+        # Force rescale to False
+        array3 = feature_extractor.to_numpy_array(array, rescale=False)
+        self.assertTrue(array3.dtype, np.uint8)
+        self.assertEqual(array3.shape, (3, 16, 32))
+        self.assertTrue(np.array_equal(array3, array.transpose(2, 0, 1)))
+
+        # Force rescale to False and no channel permute
+        array4 = feature_extractor.to_numpy_array(array, rescale=False, channel_first=False)
+        self.assertTrue(array4.dtype, np.uint8)
+        self.assertEqual(array4.shape, (16, 32, 3))
+        self.assertTrue(np.array_equal(array4, array))
+
+        # Now test the default rescale for a float tensor (defaults to False)
+        array5 = feature_extractor.to_numpy_array(array2)
+        self.assertTrue(array5.dtype, np.float32)
+        self.assertEqual(array5.shape, (3, 16, 32))
+        self.assertTrue(np.array_equal(array5, array1))
+
+    def test_conversion_image_to_image(self):
+        feature_extractor = ImageFeatureExtractionMixin()
+        image = get_random_image(16, 32)
+
+        # On an image, `to_pil_image1` is a noop.
+        image1 = feature_extractor.to_pil_image(image)
+        self.assertTrue(isinstance(image, PIL.Image.Image))
+        self.assertTrue(np.array_equal(np.array(image), np.array(image1)))
+
+    def test_conversion_array_to_image(self):
+        feature_extractor = ImageFeatureExtractionMixin()
+        array = np.random.randint(0, 256, (16, 32, 3), dtype=np.uint8)
+
+        # By default, no rescale (for an array of ints)
+        image1 = feature_extractor.to_pil_image(array)
+        self.assertTrue(isinstance(image1, PIL.Image.Image))
+        self.assertTrue(np.array_equal(np.array(image1), array))
+
+        # If the array is channel-first, proper reordering of the channels is done.
+        image2 = feature_extractor.to_pil_image(array.transpose(2, 0, 1))
+        self.assertTrue(isinstance(image2, PIL.Image.Image))
+        self.assertTrue(np.array_equal(np.array(image2), array))
+
+        # If the array has floating type, it's rescaled by default.
+        image3 = feature_extractor.to_pil_image(array.astype(np.float32) * (1 / 255.0))
+        self.assertTrue(isinstance(image3, PIL.Image.Image))
+        self.assertTrue(np.array_equal(np.array(image3), array))
+
+        # You can override the default to rescale.
+        image4 = feature_extractor.to_pil_image(array.astype(np.float32), rescale=False)
+        self.assertTrue(isinstance(image4, PIL.Image.Image))
+        self.assertTrue(np.array_equal(np.array(image4), array))
+
+        # And with floats + channel first.
+        image5 = feature_extractor.to_pil_image(array.transpose(2, 0, 1).astype(np.float32) * (1 / 255.0))
+        self.assertTrue(isinstance(image5, PIL.Image.Image))
+        self.assertTrue(np.array_equal(np.array(image5), array))
+
+    @require_torch
+    def test_conversion_tensor_to_image(self):
+        feature_extractor = ImageFeatureExtractionMixin()
+        tensor = torch.randint(0, 256, (16, 32, 3))
+        array = tensor.numpy()
+
+        # By default, no rescale (for a tensor of ints)
+        image1 = feature_extractor.to_pil_image(tensor)
+        self.assertTrue(isinstance(image1, PIL.Image.Image))
+        self.assertTrue(np.array_equal(np.array(image1), array))
+
+        # If the tensor is channel-first, proper reordering of the channels is done.
+        image2 = feature_extractor.to_pil_image(tensor.permute(2, 0, 1))
+        self.assertTrue(isinstance(image2, PIL.Image.Image))
+        self.assertTrue(np.array_equal(np.array(image2), array))
+
+        # If the tensor has floating type, it's rescaled by default.
+        image3 = feature_extractor.to_pil_image(tensor.float() / 255.0)
+        self.assertTrue(isinstance(image3, PIL.Image.Image))
+        self.assertTrue(np.array_equal(np.array(image3), array))
+
+        # You can override the default to rescale.
+        image4 = feature_extractor.to_pil_image(tensor.float(), rescale=False)
+        self.assertTrue(isinstance(image4, PIL.Image.Image))
+        self.assertTrue(np.array_equal(np.array(image4), array))
+
+        # And with floats + channel first.
+        image5 = feature_extractor.to_pil_image(tensor.permute(2, 0, 1).float() * (1 / 255.0))
+        self.assertTrue(isinstance(image5, PIL.Image.Image))
+        self.assertTrue(np.array_equal(np.array(image5), array))
+
+    def test_resize_image_and_array(self):
+        feature_extractor = ImageFeatureExtractionMixin()
+        image = get_random_image(16, 32)
+        array = np.array(image)
+
+        # Size can be an int or a tuple of ints.
+        resized_image = feature_extractor.resize(image, 8)
+        self.assertTrue(isinstance(resized_image, PIL.Image.Image))
+        self.assertEqual(resized_image.size, (8, 8))
+
+        resized_image1 = feature_extractor.resize(image, (8, 16))
+        self.assertTrue(isinstance(resized_image1, PIL.Image.Image))
+        self.assertEqual(resized_image1.size, (8, 16))
+
+        # Passing an array converts it to a PIL Image.
+        resized_image2 = feature_extractor.resize(array, 8)
+        self.assertTrue(isinstance(resized_image2, PIL.Image.Image))
+        self.assertEqual(resized_image2.size, (8, 8))
+        self.assertTrue(np.array_equal(np.array(resized_image), np.array(resized_image2)))
+
+        resized_image3 = feature_extractor.resize(image, (8, 16))
+        self.assertTrue(isinstance(resized_image3, PIL.Image.Image))
+        self.assertEqual(resized_image3.size, (8, 16))
+        self.assertTrue(np.array_equal(np.array(resized_image1), np.array(resized_image3)))
+
+    def test_resize_image_and_array_non_default_to_square(self):
+        feature_extractor = ImageFeatureExtractionMixin()
+
+        heights_widths = [
+            # height, width
+            # square image
+            (28, 28),
+            (27, 27),
+            # rectangular image: h < w
+            (28, 34),
+            (29, 35),
+            # rectangular image: h > w
+            (34, 28),
+            (35, 29),
+        ]
+
+        # single integer or single integer in tuple/list
+        sizes = [22, 27, 28, 36, [22], (27,)]
+
+        for (height, width), size in zip(heights_widths, sizes):
+            for max_size in (None, 37, 1000):
+                image = get_random_image(height, width)
+                array = np.array(image)
+
+                size = size[0] if isinstance(size, (list, tuple)) else size
+                # Size can be an int or a tuple of ints.
+                # If size is an int, smaller edge of the image will be matched to this number.
+                # i.e, if height > width, then image will be rescaled to (size * height / width, size).
+                if height < width:
+                    exp_w, exp_h = (int(size * width / height), size)
+                    if max_size is not None and max_size < exp_w:
+                        exp_w, exp_h = max_size, int(max_size * exp_h / exp_w)
+                elif width < height:
+                    exp_w, exp_h = (size, int(size * height / width))
+                    if max_size is not None and max_size < exp_h:
+                        exp_w, exp_h = int(max_size * exp_w / exp_h), max_size
+                else:
+                    exp_w, exp_h = (size, size)
+                    if max_size is not None and max_size < size:
+                        exp_w, exp_h = max_size, max_size
+
+                resized_image = feature_extractor.resize(image, size=size, default_to_square=False, max_size=max_size)
+                self.assertTrue(isinstance(resized_image, PIL.Image.Image))
+                self.assertEqual(resized_image.size, (exp_w, exp_h))
+
+                # Passing an array converts it to a PIL Image.
+                resized_image2 = feature_extractor.resize(array, size=size, default_to_square=False, max_size=max_size)
+                self.assertTrue(isinstance(resized_image2, PIL.Image.Image))
+                self.assertEqual(resized_image2.size, (exp_w, exp_h))
+                self.assertTrue(np.array_equal(np.array(resized_image), np.array(resized_image2)))
+
+    @require_torch
+    def test_resize_tensor(self):
+        feature_extractor = ImageFeatureExtractionMixin()
+        tensor = torch.randint(0, 256, (16, 32, 3))
+        array = tensor.numpy()
+
+        # Size can be an int or a tuple of ints.
+        resized_image = feature_extractor.resize(tensor, 8)
+        self.assertTrue(isinstance(resized_image, PIL.Image.Image))
+        self.assertEqual(resized_image.size, (8, 8))
+
+        resized_image1 = feature_extractor.resize(tensor, (8, 16))
+        self.assertTrue(isinstance(resized_image1, PIL.Image.Image))
+        self.assertEqual(resized_image1.size, (8, 16))
+
+        # Check we get the same results as with NumPy arrays.
+        resized_image2 = feature_extractor.resize(array, 8)
+        self.assertTrue(np.array_equal(np.array(resized_image), np.array(resized_image2)))
+
+        resized_image3 = feature_extractor.resize(array, (8, 16))
+        self.assertTrue(np.array_equal(np.array(resized_image1), np.array(resized_image3)))
+
+    def test_normalize_image(self):
+        feature_extractor = ImageFeatureExtractionMixin()
+        image = get_random_image(16, 32)
+        array = np.array(image)
+        mean = [0.1, 0.5, 0.9]
+        std = [0.2, 0.4, 0.6]
+
+        # PIL Image are converted to NumPy arrays for the normalization
+        normalized_image = feature_extractor.normalize(image, mean, std)
+        self.assertTrue(isinstance(normalized_image, np.ndarray))
+        self.assertEqual(normalized_image.shape, (3, 16, 32))
+
+        # During the conversion rescale and channel first will be applied.
+        expected = array.transpose(2, 0, 1).astype(np.float32) * (1 / 255.0)
+        np_mean = np.array(mean).astype(np.float32)[:, None, None]
+        np_std = np.array(std).astype(np.float32)[:, None, None]
+        expected = (expected - np_mean) / np_std
+        self.assertTrue(np.array_equal(normalized_image, expected))
+
+    def test_normalize_array(self):
+        feature_extractor = ImageFeatureExtractionMixin()
+        array = np.random.random((16, 32, 3))
+        mean = [0.1, 0.5, 0.9]
+        std = [0.2, 0.4, 0.6]
+
+        # mean and std can be passed as lists or NumPy arrays.
+        expected = (array - np.array(mean)) / np.array(std)
+        normalized_array = feature_extractor.normalize(array, mean, std)
+        self.assertTrue(np.array_equal(normalized_array, expected))
+
+        normalized_array = feature_extractor.normalize(array, np.array(mean), np.array(std))
+        self.assertTrue(np.array_equal(normalized_array, expected))
+
+        # Normalize will detect automatically if channel first or channel last is used.
+        array = np.random.random((3, 16, 32))
+        expected = (array - np.array(mean)[:, None, None]) / np.array(std)[:, None, None]
+        normalized_array = feature_extractor.normalize(array, mean, std)
+        self.assertTrue(np.array_equal(normalized_array, expected))
+
+        normalized_array = feature_extractor.normalize(array, np.array(mean), np.array(std))
+        self.assertTrue(np.array_equal(normalized_array, expected))
+
+    @require_torch
+    def test_normalize_tensor(self):
+        feature_extractor = ImageFeatureExtractionMixin()
+        tensor = torch.rand(16, 32, 3)
+        mean = [0.1, 0.5, 0.9]
+        std = [0.2, 0.4, 0.6]
+
+        # mean and std can be passed as lists or tensors.
+        expected = (tensor - torch.tensor(mean)) / torch.tensor(std)
+        normalized_tensor = feature_extractor.normalize(tensor, mean, std)
+        self.assertTrue(torch.equal(normalized_tensor, expected))
+
+        normalized_tensor = feature_extractor.normalize(tensor, torch.tensor(mean), torch.tensor(std))
+        self.assertTrue(torch.equal(normalized_tensor, expected))
+
+        # Normalize will detect automatically if channel first or channel last is used.
+        tensor = torch.rand(3, 16, 32)
+        expected = (tensor - torch.tensor(mean)[:, None, None]) / torch.tensor(std)[:, None, None]
+        normalized_tensor = feature_extractor.normalize(tensor, mean, std)
+        self.assertTrue(torch.equal(normalized_tensor, expected))
+
+        normalized_tensor = feature_extractor.normalize(tensor, torch.tensor(mean), torch.tensor(std))
+        self.assertTrue(torch.equal(normalized_tensor, expected))
+
+    def test_center_crop_image(self):
+        feature_extractor = ImageFeatureExtractionMixin()
+        image = get_random_image(16, 32)
+
+        # Test various crop sizes: bigger on all dimensions, on one of the dimensions only and on both dimensions.
+        crop_sizes = [8, (8, 64), 20, (32, 64)]
+        for size in crop_sizes:
+            cropped_image = feature_extractor.center_crop(image, size)
+            self.assertTrue(isinstance(cropped_image, PIL.Image.Image))
+
+            # PIL Image.size is transposed compared to NumPy or PyTorch (width first instead of height first).
+            expected_size = (size, size) if isinstance(size, int) else (size[1], size[0])
+            self.assertEqual(cropped_image.size, expected_size)
+
+    def test_center_crop_array(self):
+        feature_extractor = ImageFeatureExtractionMixin()
+        image = get_random_image(16, 32)
+        array = feature_extractor.to_numpy_array(image)
+
+        # Test various crop sizes: bigger on all dimensions, on one of the dimensions only and on both dimensions.
+        crop_sizes = [8, (8, 64), 20, (32, 64)]
+        for size in crop_sizes:
+            cropped_array = feature_extractor.center_crop(array, size)
+            self.assertTrue(isinstance(cropped_array, np.ndarray))
+
+            expected_size = (size, size) if isinstance(size, int) else size
+            self.assertEqual(cropped_array.shape[-2:], expected_size)
+
+            # Check result is consistent with PIL.Image.crop
+            cropped_image = feature_extractor.center_crop(image, size)
+            self.assertTrue(np.array_equal(cropped_array, feature_extractor.to_numpy_array(cropped_image)))
+
+    @require_torch
+    def test_center_crop_tensor(self):
+        feature_extractor = ImageFeatureExtractionMixin()
+        image = get_random_image(16, 32)
+        array = feature_extractor.to_numpy_array(image)
+        tensor = torch.tensor(array)
+
+        # Test various crop sizes: bigger on all dimensions, on one of the dimensions only and on both dimensions.
+        crop_sizes = [8, (8, 64), 20, (32, 64)]
+        for size in crop_sizes:
+            cropped_tensor = feature_extractor.center_crop(tensor, size)
+            self.assertTrue(isinstance(cropped_tensor, torch.Tensor))
+
+            expected_size = (size, size) if isinstance(size, int) else size
+            self.assertEqual(cropped_tensor.shape[-2:], expected_size)
+
+            # Check result is consistent with PIL.Image.crop
+            cropped_image = feature_extractor.center_crop(image, size)
+            self.assertTrue(torch.equal(cropped_tensor, torch.tensor(feature_extractor.to_numpy_array(cropped_image))))
+
+
+@require_vision
+class LoadImageTester(unittest.TestCase):
+    def test_load_img_url(self):
+        img = load_image(INVOICE_URL)
+        img_arr = np.array(img)
+
+        self.assertEqual(img_arr.shape, (1061, 750, 3))
+
+    @is_flaky()
+    def test_load_img_url_timeout(self):
+        with self.assertRaises((ReadTimeout, ConnectTimeout)):
+            load_image(INVOICE_URL, timeout=0.001)
+
+    def test_load_img_local(self):
+        img = load_image("./tests/fixtures/tests_samples/COCO/000000039769.png")
+        img_arr = np.array(img)
+
+        self.assertEqual(
+            img_arr.shape,
+            (480, 640, 3),
+        )
+
+    def test_load_img_base64_prefix(self):
+        try:
+            tmp_file = tempfile.NamedTemporaryFile(delete=False).name
+            with open(tmp_file, "wb") as f:
+                http_get(
+                    "https://huggingface.co/datasets/hf-internal-testing/dummy-base64-images/raw/main/image_0.txt", f
+                )
+
+            with open(tmp_file, encoding="utf-8") as b64:
+                img = load_image(b64.read())
+                img_arr = np.array(img)
+
+        finally:
+            os.remove(tmp_file)
+
+        self.assertEqual(img_arr.shape, (64, 32, 3))
+
+    def test_load_img_base64(self):
+        try:
+            tmp_file = tempfile.NamedTemporaryFile(delete=False).name
+            with open(tmp_file, "wb") as f:
+                http_get(
+                    "https://huggingface.co/datasets/hf-internal-testing/dummy-base64-images/raw/main/image_1.txt", f
+                )
+
+            with open(tmp_file, encoding="utf-8") as b64:
+                img = load_image(b64.read())
+                img_arr = np.array(img)
+
+        finally:
+            os.remove(tmp_file)
+
+        self.assertEqual(img_arr.shape, (64, 32, 3))
+
+    def test_load_img_base64_encoded_bytes(self):
+        try:
+            tmp_file = tempfile.NamedTemporaryFile(delete=False).name
+            with open(tmp_file, "wb") as f:
+                http_get(
+                    "https://huggingface.co/datasets/hf-internal-testing/dummy-base64-images/raw/main/image_2.txt", f
+                )
+
+            with codecs.open(tmp_file, encoding="unicode_escape") as b64:
+                img = load_image(b64.read())
+                img_arr = np.array(img)
+
+        finally:
+            os.remove(tmp_file)
+
+        self.assertEqual(img_arr.shape, (256, 256, 3))
+
+    def test_load_img_rgba(self):
+        # we use revision="refs/pr/1" until the PR is merged
+        # https://hf.co/datasets/hf-internal-testing/fixtures_image_utils/discussions/1
+        img = get_image_from_hub_dataset(
+            "hf-internal-testing/fixtures_image_utils", "0-test-lena.png", revision="refs/pr/1"
+        )
+
+        img = load_image(img)  # img with mode RGBA
+        img_arr = np.array(img)
+
+        self.assertEqual(
+            img_arr.shape,
+            (512, 512, 3),
+        )
+
+    def test_load_img_la(self):
+        # we use revision="refs/pr/1" until the PR is merged
+        # https://hf.co/datasets/hf-internal-testing/fixtures_image_utils/discussions/1
+        img = get_image_from_hub_dataset(
+            "hf-internal-testing/fixtures_image_utils", "1-test-parrots.png", revision="refs/pr/1"
+        )
+
+        img = load_image(img)  # img with mode LA
+        img_arr = np.array(img)
+
+        self.assertEqual(
+            img_arr.shape,
+            (512, 768, 3),
+        )
+
+    def test_load_img_l(self):
+        # we use revision="refs/pr/1" until the PR is merged
+        # https://hf.co/datasets/hf-internal-testing/fixtures_image_utils/discussions/1
+        img = get_image_from_hub_dataset(
+            "hf-internal-testing/fixtures_image_utils", "2-test-tree.png", revision="refs/pr/1"
+        )
+
+        img = load_image(img)  # img with mode L
+        img_arr = np.array(img)
+
+        self.assertEqual(
+            img_arr.shape,
+            (381, 225, 3),
+        )
+
+    def test_load_img_exif_transpose(self):
+        # we use revision="refs/pr/1" until the PR is merged
+        # https://hf.co/datasets/hf-internal-testing/fixtures_image_utils/discussions/1
+
+        img_without_exif_transpose = get_image_from_hub_dataset(
+            "hf-internal-testing/fixtures_image_utils", "3-test-cat-rotated.jpg", revision="refs/pr/1"
+        )
+        img_arr_without_exif_transpose = np.array(img_without_exif_transpose)
+
+        self.assertEqual(
+            img_arr_without_exif_transpose.shape,
+            (333, 500, 3),
+        )
+
+        img_with_exif_transpose = load_image(img_without_exif_transpose)
+        img_arr_with_exif_transpose = np.array(img_with_exif_transpose)
+
+        self.assertEqual(
+            img_arr_with_exif_transpose.shape,
+            (500, 333, 3),
+        )
+
+
+class UtilFunctionTester(unittest.TestCase):
+    def test_get_image_size(self):
+        # Test we can infer the size and channel dimension of an image.
+        image = np.random.randint(0, 256, (32, 64, 3))
+        self.assertEqual(get_image_size(image), (32, 64))
+
+        image = np.random.randint(0, 256, (3, 32, 64))
+        self.assertEqual(get_image_size(image), (32, 64))
+
+        # Test the channel dimension can be overridden
+        image = np.random.randint(0, 256, (3, 32, 64))
+        self.assertEqual(get_image_size(image, channel_dim=ChannelDimension.LAST), (3, 32))
+
+    def test_infer_channel_dimension(self):
+        # Test we fail with invalid input
+        with pytest.raises(ValueError):
+            infer_channel_dimension_format(np.random.randint(0, 256, (10, 10)))
+
+        with pytest.raises(ValueError):
+            infer_channel_dimension_format(np.random.randint(0, 256, (10, 10, 10, 10, 10)))
+
+        # Test we fail if neither first not last dimension is of size 3 or 1
+        with pytest.raises(ValueError):
+            infer_channel_dimension_format(np.random.randint(0, 256, (10, 1, 50)))
+
+        # But if we explicitly set one of the number of channels to 50 it works
+        inferred_dim = infer_channel_dimension_format(np.random.randint(0, 256, (10, 1, 50)), num_channels=50)
+        self.assertEqual(inferred_dim, ChannelDimension.LAST)
+
+        # Test we correctly identify the channel dimension
+        image = np.random.randint(0, 256, (3, 4, 5))
+        inferred_dim = infer_channel_dimension_format(image)
+        self.assertEqual(inferred_dim, ChannelDimension.FIRST)
+
+        image = np.random.randint(0, 256, (1, 4, 5))
+        inferred_dim = infer_channel_dimension_format(image)
+        self.assertEqual(inferred_dim, ChannelDimension.FIRST)
+
+        image = np.random.randint(0, 256, (4, 5, 3))
+        inferred_dim = infer_channel_dimension_format(image)
+        self.assertEqual(inferred_dim, ChannelDimension.LAST)
+
+        image = np.random.randint(0, 256, (4, 5, 1))
+        inferred_dim = infer_channel_dimension_format(image)
+        self.assertEqual(inferred_dim, ChannelDimension.LAST)
+
+        # We can take a batched array of images and find the dimension
+        image = np.random.randint(0, 256, (1, 3, 4, 5))
+        inferred_dim = infer_channel_dimension_format(image)
+        self.assertEqual(inferred_dim, ChannelDimension.FIRST)
+
+    def test_get_channel_dimension_axis(self):
+        # Test we correctly identify the channel dimension
+        image = np.random.randint(0, 256, (3, 4, 5))
+        inferred_axis = get_channel_dimension_axis(image)
+        self.assertEqual(inferred_axis, 0)
+
+        image = np.random.randint(0, 256, (1, 4, 5))
+        inferred_axis = get_channel_dimension_axis(image)
+        self.assertEqual(inferred_axis, 0)
+
+        image = np.random.randint(0, 256, (4, 5, 3))
+        inferred_axis = get_channel_dimension_axis(image)
+        self.assertEqual(inferred_axis, 2)
+
+        image = np.random.randint(0, 256, (4, 5, 1))
+        inferred_axis = get_channel_dimension_axis(image)
+        self.assertEqual(inferred_axis, 2)
+
+        # We can take a batched array of images and find the dimension
+        image = np.random.randint(0, 256, (1, 3, 4, 5))
+        inferred_axis = get_channel_dimension_axis(image)
+        self.assertEqual(inferred_axis, 1)

docs/transformers/tests/utils/test_import_structure.py

@@ -0,0 +1,104 @@
+import os
+import unittest
+from pathlib import Path
+
+from transformers.utils.import_utils import define_import_structure, spread_import_structure
+
+
+import_structures = Path("import_structures")
+
+
+def fetch__all__(file_content):
+    """
+    Returns the content of the __all__ variable in the file content.
+    Returns None if not defined, otherwise returns a list of strings.
+    """
+    lines = file_content.split("\n")
+    for line_index in range(len(lines)):
+        line = lines[line_index]
+        if line.startswith("__all__ = "):
+            # __all__ is defined on a single line
+            if line.endswith("]"):
+                return [obj.strip("\"' ") for obj in line.split("=")[1].strip(" []").split(",")]
+
+            # __all__ is defined on multiple lines
+            else:
+                _all = []
+                for __all__line_index in range(line_index + 1, len(lines)):
+                    if lines[__all__line_index].strip() == "]":
+                        return _all
+                    else:
+                        _all.append(lines[__all__line_index].strip("\"', "))
+
+
+class TestImportStructures(unittest.TestCase):
+    base_transformers_path = Path(__file__).parent.parent.parent
+    models_path = base_transformers_path / "src" / "transformers" / "models"
+    models_import_structure = spread_import_structure(define_import_structure(models_path))
+
+    # TODO: Lysandre
+    # See https://app.circleci.com/pipelines/github/huggingface/transformers/104762/workflows/7ba9c6f7-a3b2-44e6-8eaf-749c7b7261f7/jobs/1393260/tests
+    @unittest.skip(reason="failing")
+    def test_definition(self):
+        import_structure = define_import_structure(import_structures)
+        import_structure_definition = {
+            frozenset(()): {
+                "import_structure_raw_register": {"A0", "a0", "A4"},
+                "import_structure_register_with_comments": {"B0", "b0"},
+            },
+            frozenset(("tf", "torch")): {
+                "import_structure_raw_register": {"A1", "a1", "A2", "a2", "A3", "a3"},
+                "import_structure_register_with_comments": {"B1", "b1", "B2", "b2", "B3", "b3"},
+            },
+            frozenset(("torch",)): {
+                "import_structure_register_with_duplicates": {"C0", "c0", "C1", "c1", "C2", "c2", "C3", "c3"},
+            },
+        }
+
+        self.assertDictEqual(import_structure, import_structure_definition)
+
+    def test_transformers_specific_model_import(self):
+        """
+        This test ensures that there is equivalence between what is written down in __all__ and what is
+        written down with register().
+
+        It doesn't test the backends attributed to register().
+        """
+        for architecture in os.listdir(self.models_path):
+            if (
+                os.path.isfile(self.models_path / architecture)
+                or architecture.startswith("_")
+                or architecture == "deprecated"
+            ):
+                continue
+
+            with self.subTest(f"Testing arch {architecture}"):
+                import_structure = define_import_structure(self.models_path / architecture)
+                backend_agnostic_import_structure = {}
+                for requirement, module_object_mapping in import_structure.items():
+                    for module, objects in module_object_mapping.items():
+                        if module not in backend_agnostic_import_structure:
+                            backend_agnostic_import_structure[module] = []
+
+                        backend_agnostic_import_structure[module].extend(objects)
+
+                for module, objects in backend_agnostic_import_structure.items():
+                    with open(self.models_path / architecture / f"{module}.py") as f:
+                        content = f.read()
+                        _all = fetch__all__(content)
+
+                        if _all is None:
+                            raise ValueError(f"{module} doesn't have __all__ defined.")
+
+                        error_message = (
+                            f"self.models_path / architecture / f'{module}.py doesn't seem to be defined correctly:\n"
+                            f"Defined in __all__: {sorted(_all)}\nDefined with register: {sorted(objects)}"
+                        )
+                        self.assertListEqual(sorted(objects), sorted(_all), msg=error_message)
+
+    # TODO: Lysandre
+    # See https://app.circleci.com/pipelines/github/huggingface/transformers/104762/workflows/7ba9c6f7-a3b2-44e6-8eaf-749c7b7261f7/jobs/1393260/tests
+    @unittest.skip(reason="failing")
+    def test_export_backend_should_be_defined(self):
+        with self.assertRaisesRegex(ValueError, "Backend should be defined in the BACKENDS_MAPPING"):
+            pass

docs/transformers/tests/utils/test_import_utils.py

@@ -0,0 +1,26 @@
+import sys
+
+from transformers.testing_utils import run_test_using_subprocess
+from transformers.utils.import_utils import clear_import_cache
+
+
+@run_test_using_subprocess
+def test_clear_import_cache():
+    """Test the clear_import_cache function."""
+
+    # Save initial state
+    initial_modules = {name: mod for name, mod in sys.modules.items() if name.startswith("transformers.")}
+    assert len(initial_modules) > 0, "No transformers modules loaded before test"
+
+    # Execute clear_import_cache() function
+    clear_import_cache()
+
+    # Verify modules were removed
+    remaining_modules = {name: mod for name, mod in sys.modules.items() if name.startswith("transformers.")}
+    assert len(remaining_modules) < len(initial_modules), "No modules were removed"
+
+    # Import and verify module exists
+    from transformers.models.auto import modeling_auto
+
+    assert "transformers.models.auto.modeling_auto" in sys.modules
+    assert modeling_auto.__name__ == "transformers.models.auto.modeling_auto"

docs/transformers/tests/utils/test_logging.py

@@ -0,0 +1,135 @@
+# Copyright 2020 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import unittest
+
+from huggingface_hub.utils import are_progress_bars_disabled
+
+import transformers.models.bart.tokenization_bart
+from transformers import logging
+from transformers.testing_utils import CaptureLogger, mockenv, mockenv_context
+from transformers.utils.logging import disable_progress_bar, enable_progress_bar
+
+
+class HfArgumentParserTest(unittest.TestCase):
+    def test_set_level(self):
+        logger = logging.get_logger()
+
+        # the current default level is logging.WARNING
+        level_origin = logging.get_verbosity()
+
+        logging.set_verbosity_error()
+        self.assertEqual(logger.getEffectiveLevel(), logging.get_verbosity())
+
+        logging.set_verbosity_warning()
+        self.assertEqual(logger.getEffectiveLevel(), logging.get_verbosity())
+
+        logging.set_verbosity_info()
+        self.assertEqual(logger.getEffectiveLevel(), logging.get_verbosity())
+
+        logging.set_verbosity_debug()
+        self.assertEqual(logger.getEffectiveLevel(), logging.get_verbosity())
+
+        # restore to the original level
+        logging.set_verbosity(level_origin)
+
+    def test_integration(self):
+        level_origin = logging.get_verbosity()
+
+        logger = logging.get_logger("transformers.models.bart.tokenization_bart")
+        msg = "Testing 1, 2, 3"
+
+        # should be able to log warnings (if default settings weren't overridden by `pytest --log-level-all`)
+        if level_origin <= logging.WARNING:
+            with CaptureLogger(logger) as cl:
+                logger.warning(msg)
+            self.assertEqual(cl.out, msg + "\n")
+
+        # this is setting the level for all of `transformers.*` loggers
+        logging.set_verbosity_error()
+
+        # should not be able to log warnings
+        with CaptureLogger(logger) as cl:
+            logger.warning(msg)
+        self.assertEqual(cl.out, "")
+
+        # should be able to log warnings again
+        logging.set_verbosity_warning()
+        with CaptureLogger(logger) as cl:
+            logger.warning(msg)
+        self.assertEqual(cl.out, msg + "\n")
+
+        # restore to the original level
+        logging.set_verbosity(level_origin)
+
+    @mockenv(TRANSFORMERS_VERBOSITY="error")
+    def test_env_override(self):
+        # reset for the env var to take effect, next time some logger call is made
+        transformers.utils.logging._reset_library_root_logger()
+        # this action activates the env var
+        _ = logging.get_logger("transformers.models.bart.tokenization_bart")
+
+        env_level_str = os.getenv("TRANSFORMERS_VERBOSITY", None)
+        env_level = logging.log_levels[env_level_str]
+
+        current_level = logging.get_verbosity()
+        self.assertEqual(
+            env_level,
+            current_level,
+            f"TRANSFORMERS_VERBOSITY={env_level_str}/{env_level}, but internal verbosity is {current_level}",
+        )
+
+        # restore to the original level
+        os.environ["TRANSFORMERS_VERBOSITY"] = ""
+        transformers.utils.logging._reset_library_root_logger()
+
+    @mockenv(TRANSFORMERS_VERBOSITY="super-error")
+    def test_env_invalid_override(self):
+        # reset for the env var to take effect, next time some logger call is made
+        transformers.utils.logging._reset_library_root_logger()
+        logger = logging.logging.getLogger()
+        with CaptureLogger(logger) as cl:
+            # this action activates the env var
+            logging.get_logger("transformers.models.bart.tokenization_bart")
+        self.assertIn("Unknown option TRANSFORMERS_VERBOSITY=super-error", cl.out)
+
+        # no need to restore as nothing was changed
+
+    def test_advisory_warnings(self):
+        # testing `logger.warning_advice()`
+        transformers.utils.logging._reset_library_root_logger()
+
+        logger = logging.get_logger("transformers.models.bart.tokenization_bart")
+        msg = "Testing 1, 2, 3"
+
+        with mockenv_context(TRANSFORMERS_NO_ADVISORY_WARNINGS="1"):
+            # nothing should be logged as env var disables this method
+            with CaptureLogger(logger) as cl:
+                logger.warning_advice(msg)
+            self.assertEqual(cl.out, "")
+
+        with mockenv_context(TRANSFORMERS_NO_ADVISORY_WARNINGS=""):
+            # should log normally as TRANSFORMERS_NO_ADVISORY_WARNINGS is unset
+            with CaptureLogger(logger) as cl:
+                logger.warning_advice(msg)
+            self.assertEqual(cl.out, msg + "\n")
+
+
+def test_set_progress_bar_enabled():
+    disable_progress_bar()
+    assert are_progress_bars_disabled()
+
+    enable_progress_bar()
+    assert not are_progress_bars_disabled()

docs/transformers/tests/utils/test_model_card.py

@@ -0,0 +1,88 @@
+# Copyright 2019 HuggingFace Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import json
+import os
+import tempfile
+import unittest
+
+from transformers.modelcard import ModelCard, TrainingSummary
+
+
+class ModelCardTester(unittest.TestCase):
+    def setUp(self):
+        self.inputs_dict = {
+            "model_details": {
+                "Organization": "testing",
+                "Model date": "today",
+                "Model version": "v2.1, Developed by Test Corp in 2019.",
+                "Architecture": "Convolutional Neural Network.",
+            },
+            "metrics": "BLEU and ROUGE-1",
+            "evaluation_data": {
+                "Datasets": {"BLEU": "My-great-dataset-v1", "ROUGE-1": "My-short-dataset-v2.1"},
+                "Preprocessing": "See details on https://arxiv.org/pdf/1810.03993.pdf",
+            },
+            "training_data": {
+                "Dataset": "English Wikipedia dump dated 2018-12-01",
+                "Preprocessing": (
+                    "Using SentencePiece vocabulary of size 52k tokens. See details on"
+                    " https://arxiv.org/pdf/1810.03993.pdf"
+                ),
+            },
+            "quantitative_analyses": {"BLEU": 55.1, "ROUGE-1": 76},
+        }
+
+    def test_model_card_common_properties(self):
+        modelcard = ModelCard.from_dict(self.inputs_dict)
+        self.assertTrue(hasattr(modelcard, "model_details"))
+        self.assertTrue(hasattr(modelcard, "intended_use"))
+        self.assertTrue(hasattr(modelcard, "factors"))
+        self.assertTrue(hasattr(modelcard, "metrics"))
+        self.assertTrue(hasattr(modelcard, "evaluation_data"))
+        self.assertTrue(hasattr(modelcard, "training_data"))
+        self.assertTrue(hasattr(modelcard, "quantitative_analyses"))
+        self.assertTrue(hasattr(modelcard, "ethical_considerations"))
+        self.assertTrue(hasattr(modelcard, "caveats_and_recommendations"))
+
+    def test_model_card_to_json_string(self):
+        modelcard = ModelCard.from_dict(self.inputs_dict)
+        obj = json.loads(modelcard.to_json_string())
+        for key, value in self.inputs_dict.items():
+            self.assertEqual(obj[key], value)
+
+    def test_model_card_to_json_file(self):
+        model_card_first = ModelCard.from_dict(self.inputs_dict)
+
+        with tempfile.TemporaryDirectory() as tmpdirname:
+            filename = os.path.join(tmpdirname, "modelcard.json")
+            model_card_first.to_json_file(filename)
+            model_card_second = ModelCard.from_json_file(filename)
+
+        self.assertEqual(model_card_second.to_dict(), model_card_first.to_dict())
+
+    def test_model_card_from_and_save_pretrained(self):
+        model_card_first = ModelCard.from_dict(self.inputs_dict)
+
+        with tempfile.TemporaryDirectory() as tmpdirname:
+            model_card_first.save_pretrained(tmpdirname)
+            model_card_second = ModelCard.from_pretrained(tmpdirname)
+
+        self.assertEqual(model_card_second.to_dict(), model_card_first.to_dict())
+
+    def test_model_summary_modelcard_base_metadata(self):
+        metadata = TrainingSummary("Model name").create_metadata()
+        self.assertTrue("library_name" in metadata)
+        self.assertTrue(metadata["library_name"] == "transformers")

docs/transformers/tests/utils/test_model_debugging_utils.py

@@ -0,0 +1,122 @@
+# coding=utf-8
+# Copyright 2025 The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import gc
+import json
+import os
+import tempfile
+import unittest
+from pathlib import Path
+
+from transformers import is_torch_available
+from transformers.model_debugging_utils import model_addition_debugger_context
+
+
+if is_torch_available():
+    import torch
+    from torch import nn
+
+    class ToyModel(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.embed = nn.Embedding(10, 4)
+            self.linear_1 = nn.Linear(4, 8)
+            self.linear_2 = nn.Linear(8, 2)
+            self.act = nn.ReLU()
+
+        def forward(self, input_ids: str):
+            hidden_states = self.embed(input_ids).mean(dim=1)
+            hidden_states = self.act(self.linear_1(hidden_states))
+            return self.linear_2(hidden_states)
+
+    class TestModelAdditionDebugger(unittest.TestCase):
+        def setUp(self):
+            self.model = ToyModel()
+            self.inputs = {"input_ids": torch.randint(0, 10, (1, 3))}
+
+        def tearDown(self):
+            gc.collect()
+
+        def test_debugger_outputs(self):
+            with tempfile.TemporaryDirectory() as tmpdir:
+                with model_addition_debugger_context(self.model, debug_path=str(tmpdir)):
+                    _ = self.model.forward(**self.inputs)
+
+                base = f"{self.model.__class__.__name__}_debug_tree"
+                summary = Path(os.path.join(tmpdir, f"{base}_SUMMARY.json"))
+                full = Path(os.path.join(tmpdir, f"{base}_FULL_TENSORS.json"))
+                self.assertTrue(os.path.isfile(summary) and os.path.isfile(full))
+                data = json.loads(summary.read_text())
+                self.assertTrue({"module_path", "inputs", "children"} <= data.keys())
+                self.assertTrue(data["children"])
+
+    class ToyLayer(nn.Module):
+        def __init__(self, layer_index):
+            super().__init__()
+            self.layer_index = layer_index
+            self.layer_operation = nn.Linear(4, 4)
+
+        def forward(self, hidden_states):
+            return self.layer_operation(hidden_states)
+
+    class ToyModelWithLayers(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.input_proj = nn.Linear(4, 4)
+            self.layers = nn.ModuleList([ToyLayer(layer_index) for layer_index in range(6)])
+            self.output_proj = nn.Linear(4, 2)
+
+        def forward(self, x):
+            x = self.input_proj(x)
+            for layer in self.layers:
+                x = layer(x)
+            return self.output_proj(x)
+
+    class TestModelWithLayers(unittest.TestCase):
+        def setUp(self):
+            self.inputs = {"input_ids": torch.randint(0, 10, (1, 3))}
+            self.model_with_layers = ToyModelWithLayers()
+            self.dense_input = {"x": torch.randn(1, 4)}
+
+        def tearDown(self):
+            gc.collect()
+
+        def test_layer_pruning_behavior(self):
+            # No pruning: expect all 6 layers
+            with tempfile.TemporaryDirectory() as tmpdir:
+                with model_addition_debugger_context(self.model_with_layers, debug_path=tmpdir, do_prune_layers=False):
+                    _ = self.model_with_layers(**self.dense_input)
+
+                summary_path = os.path.join(tmpdir, "ToyModelWithLayers_debug_tree_SUMMARY.json")
+                with open(summary_path) as f:
+                    data = json.load(f)
+                self.assertEqual(set(data.keys()), {"module_path", "inputs", "children"})
+                for layer_index in range(6):
+                    self.assertEqual(
+                        data["children"][layer_index + 1]["module_path"],
+                        f"ToyModelWithLayers.layers.{int(layer_index)}",
+                    )
+
+            # Pruning: expect only 2 layers (0 and 5)
+            with tempfile.TemporaryDirectory() as tmpdir:
+                with model_addition_debugger_context(self.model_with_layers, debug_path=tmpdir, do_prune_layers=True):
+                    _ = self.model_with_layers(**self.dense_input)
+
+                summary_path = os.path.join(tmpdir, "ToyModelWithLayers_debug_tree_SUMMARY.json")
+                with open(summary_path) as f:
+                    data = json.load(f)
+                self.assertEqual(set(data.keys()), {"module_path", "inputs", "children"})
+                self.assertEqual(data["children"][1]["module_path"], "ToyModelWithLayers.layers.0")
+                self.assertEqual(data["children"][2]["module_path"], "ToyModelWithLayers.layers.5")

docs/transformers/tests/utils/test_model_output.py

@@ -0,0 +1,201 @@
+# Copyright 2020 The Hugging Face Team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import io
+import unittest
+from dataclasses import dataclass
+from typing import Optional
+
+from transformers import AlbertForMaskedLM
+from transformers.testing_utils import require_torch
+from transformers.utils import ModelOutput, is_torch_available
+
+
+if is_torch_available():
+    import torch
+
+    from transformers.pytorch_utils import is_torch_greater_or_equal_than_2_2
+
+
+@dataclass
+class ModelOutputTest(ModelOutput):
+    a: float
+    b: Optional[float] = None
+    c: Optional[float] = None
+
+
+class ModelOutputTester(unittest.TestCase):
+    def test_get_attributes(self):
+        x = ModelOutputTest(a=30)
+        self.assertEqual(x.a, 30)
+        self.assertIsNone(x.b)
+        self.assertIsNone(x.c)
+        with self.assertRaises(AttributeError):
+            _ = x.d
+
+    def test_index_with_ints_and_slices(self):
+        x = ModelOutputTest(a=30, b=10)
+        self.assertEqual(x[0], 30)
+        self.assertEqual(x[1], 10)
+        self.assertEqual(x[:2], (30, 10))
+        self.assertEqual(x[:], (30, 10))
+
+        x = ModelOutputTest(a=30, c=10)
+        self.assertEqual(x[0], 30)
+        self.assertEqual(x[1], 10)
+        self.assertEqual(x[:2], (30, 10))
+        self.assertEqual(x[:], (30, 10))
+
+    def test_index_with_strings(self):
+        x = ModelOutputTest(a=30, b=10)
+        self.assertEqual(x["a"], 30)
+        self.assertEqual(x["b"], 10)
+        with self.assertRaises(KeyError):
+            _ = x["c"]
+
+        x = ModelOutputTest(a=30, c=10)
+        self.assertEqual(x["a"], 30)
+        self.assertEqual(x["c"], 10)
+        with self.assertRaises(KeyError):
+            _ = x["b"]
+
+    def test_dict_like_properties(self):
+        x = ModelOutputTest(a=30)
+        self.assertEqual(list(x.keys()), ["a"])
+        self.assertEqual(list(x.values()), [30])
+        self.assertEqual(list(x.items()), [("a", 30)])
+        self.assertEqual(list(x), ["a"])
+
+        x = ModelOutputTest(a=30, b=10)
+        self.assertEqual(list(x.keys()), ["a", "b"])
+        self.assertEqual(list(x.values()), [30, 10])
+        self.assertEqual(list(x.items()), [("a", 30), ("b", 10)])
+        self.assertEqual(list(x), ["a", "b"])
+
+        x = ModelOutputTest(a=30, c=10)
+        self.assertEqual(list(x.keys()), ["a", "c"])
+        self.assertEqual(list(x.values()), [30, 10])
+        self.assertEqual(list(x.items()), [("a", 30), ("c", 10)])
+        self.assertEqual(list(x), ["a", "c"])
+
+        with self.assertRaises(Exception):
+            x = x.update({"d": 20})
+        with self.assertRaises(Exception):
+            del x["a"]
+        with self.assertRaises(Exception):
+            _ = x.pop("a")
+        with self.assertRaises(Exception):
+            _ = x.setdefault("d", 32)
+
+    def test_set_attributes(self):
+        x = ModelOutputTest(a=30)
+        x.a = 10
+        self.assertEqual(x.a, 10)
+        self.assertEqual(x["a"], 10)
+
+    def test_set_keys(self):
+        x = ModelOutputTest(a=30)
+        x["a"] = 10
+        self.assertEqual(x.a, 10)
+        self.assertEqual(x["a"], 10)
+
+    def test_instantiate_from_dict(self):
+        x = ModelOutputTest({"a": 30, "b": 10})
+        self.assertEqual(list(x.keys()), ["a", "b"])
+        self.assertEqual(x.a, 30)
+        self.assertEqual(x.b, 10)
+
+    def test_instantiate_from_iterator(self):
+        x = ModelOutputTest([("a", 30), ("b", 10)])
+        self.assertEqual(list(x.keys()), ["a", "b"])
+        self.assertEqual(x.a, 30)
+        self.assertEqual(x.b, 10)
+
+        with self.assertRaises(ValueError):
+            _ = ModelOutputTest([("a", 30), (10, 10)])
+
+        x = ModelOutputTest(a=(30, 30))
+        self.assertEqual(list(x.keys()), ["a"])
+        self.assertEqual(x.a, (30, 30))
+
+    @require_torch
+    def test_torch_pytree(self):
+        # ensure torch.utils._pytree treats ModelOutput subclasses as nodes (and not leaves)
+        # this is important for DistributedDataParallel gradient synchronization with static_graph=True
+        import torch.utils._pytree as pytree
+
+        x = ModelOutput({"a": 1.0, "c": 2.0})
+        self.assertFalse(pytree._is_leaf(x))
+
+        x = ModelOutputTest(a=1.0, c=2.0)
+        self.assertFalse(pytree._is_leaf(x))
+
+        expected_flat_outs = [1.0, 2.0]
+        expected_tree_spec = pytree.TreeSpec(ModelOutputTest, ["a", "c"], [pytree.LeafSpec(), pytree.LeafSpec()])
+
+        actual_flat_outs, actual_tree_spec = pytree.tree_flatten(x)
+        self.assertEqual(expected_flat_outs, actual_flat_outs)
+        self.assertEqual(expected_tree_spec, actual_tree_spec)
+
+        unflattened_x = pytree.tree_unflatten(actual_flat_outs, actual_tree_spec)
+        self.assertEqual(x, unflattened_x)
+
+        if is_torch_greater_or_equal_than_2_2:
+            self.assertEqual(
+                pytree.treespec_dumps(actual_tree_spec),
+                '[1, {"type": "tests.utils.test_model_output.ModelOutputTest", "context": "[\\"a\\", \\"c\\"]", "children_spec": [{"type": null, "context": null, "children_spec": []}, {"type": null, "context": null, "children_spec": []}]}]',
+            )
+
+    # TODO: @ydshieh
+    @unittest.skip(reason="CPU OOM")
+    @require_torch
+    def test_export_serialization(self):
+        if not is_torch_greater_or_equal_than_2_2:
+            self.skipTest(reason="Export serialization requires torch >= 2.2.0")
+
+        model_cls = AlbertForMaskedLM
+        model_config = model_cls.config_class()
+        model = model_cls(model_config)
+
+        input_dict = {"input_ids": torch.randint(0, 30000, (1, 512), dtype=torch.int64, requires_grad=False)}
+
+        ep = torch.export.export(model, (), input_dict)
+
+        buffer = io.BytesIO()
+        torch.export.save(ep, buffer)
+        buffer.seek(0)
+        loaded_ep = torch.export.load(buffer)
+
+        input_dict = {"input_ids": torch.randint(0, 30000, (1, 512), dtype=torch.int64, requires_grad=False)}
+        assert torch.allclose(model(**input_dict).logits, loaded_ep(**input_dict).logits)
+
+
+class ModelOutputTestNoDataclass(ModelOutput):
+    """Invalid test subclass of ModelOutput where @dataclass decorator is not used"""
+
+    a: float
+    b: Optional[float] = None
+    c: Optional[float] = None
+
+
+class ModelOutputSubclassTester(unittest.TestCase):
+    def test_direct_model_output(self):
+        # Check that direct usage of ModelOutput instantiates without errors
+        ModelOutput({"a": 1.1})
+
+    def test_subclass_no_dataclass(self):
+        # Check that a subclass of ModelOutput without @dataclass is invalid
+        # A valid subclass is inherently tested other unit tests above.
+        with self.assertRaises(TypeError):
+            ModelOutputTestNoDataclass(a=1.1, b=2.2, c=3.3)

docs/transformers/tests/utils/test_modeling_flax_utils.py

@@ -0,0 +1,285 @@
+# Copyright 2020 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import tempfile
+import unittest
+
+import numpy as np
+from huggingface_hub import HfFolder, snapshot_download
+
+from transformers import BertConfig, is_flax_available
+from transformers.testing_utils import (
+    TOKEN,
+    CaptureLogger,
+    TemporaryHubRepo,
+    is_staging_test,
+    require_flax,
+    require_safetensors,
+)
+from transformers.utils import FLAX_WEIGHTS_NAME, SAFE_WEIGHTS_NAME, logging
+
+
+if is_flax_available():
+    import os
+
+    from flax.core.frozen_dict import unfreeze
+    from flax.traverse_util import flatten_dict
+
+    from transformers import FlaxBertModel
+
+    os.environ["XLA_PYTHON_CLIENT_MEM_FRACTION"] = "0.12"  # assumed parallelism: 8
+
+
+@require_flax
+@is_staging_test
+class FlaxModelPushToHubTester(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls._token = TOKEN
+        HfFolder.save_token(TOKEN)
+
+    def test_push_to_hub(self):
+        with TemporaryHubRepo(token=self._token) as tmp_repo:
+            config = BertConfig(
+                vocab_size=99, hidden_size=32, num_hidden_layers=5, num_attention_heads=4, intermediate_size=37
+            )
+            model = FlaxBertModel(config)
+            model.push_to_hub(tmp_repo.repo_id, token=self._token)
+
+            new_model = FlaxBertModel.from_pretrained(tmp_repo.repo_id)
+
+            base_params = flatten_dict(unfreeze(model.params))
+            new_params = flatten_dict(unfreeze(new_model.params))
+
+            for key in base_params.keys():
+                max_diff = (base_params[key] - new_params[key]).sum().item()
+                self.assertLessEqual(max_diff, 1e-3, msg=f"{key} not identical")
+
+    def test_push_to_hub_via_save_pretrained(self):
+        with TemporaryHubRepo(token=self._token) as tmp_repo:
+            config = BertConfig(
+                vocab_size=99, hidden_size=32, num_hidden_layers=5, num_attention_heads=4, intermediate_size=37
+            )
+            model = FlaxBertModel(config)
+            # Push to hub via save_pretrained
+            with tempfile.TemporaryDirectory() as tmp_dir:
+                model.save_pretrained(tmp_dir, repo_id=tmp_repo.repo_id, push_to_hub=True, token=self._token)
+
+            new_model = FlaxBertModel.from_pretrained(tmp_repo.repo_id)
+
+            base_params = flatten_dict(unfreeze(model.params))
+            new_params = flatten_dict(unfreeze(new_model.params))
+
+            for key in base_params.keys():
+                max_diff = (base_params[key] - new_params[key]).sum().item()
+                self.assertLessEqual(max_diff, 1e-3, msg=f"{key} not identical")
+
+    def test_push_to_hub_in_organization(self):
+        with TemporaryHubRepo(namespace="valid_org", token=self._token) as tmp_repo:
+            config = BertConfig(
+                vocab_size=99, hidden_size=32, num_hidden_layers=5, num_attention_heads=4, intermediate_size=37
+            )
+            model = FlaxBertModel(config)
+            model.push_to_hub(tmp_repo.repo_id, token=self._token)
+
+            new_model = FlaxBertModel.from_pretrained(tmp_repo.repo_id)
+
+            base_params = flatten_dict(unfreeze(model.params))
+            new_params = flatten_dict(unfreeze(new_model.params))
+
+            for key in base_params.keys():
+                max_diff = (base_params[key] - new_params[key]).sum().item()
+                self.assertLessEqual(max_diff, 1e-3, msg=f"{key} not identical")
+
+    def test_push_to_hub_in_organization_via_save_pretrained(self):
+        with TemporaryHubRepo(namespace="valid_org", token=self._token) as tmp_repo:
+            config = BertConfig(
+                vocab_size=99, hidden_size=32, num_hidden_layers=5, num_attention_heads=4, intermediate_size=37
+            )
+            model = FlaxBertModel(config)
+            # Push to hub via save_pretrained
+            with tempfile.TemporaryDirectory() as tmp_dir:
+                model.save_pretrained(tmp_dir, repo_id=tmp_repo.repo_id, push_to_hub=True, token=self._token)
+
+            new_model = FlaxBertModel.from_pretrained(tmp_repo.repo_id)
+
+            base_params = flatten_dict(unfreeze(model.params))
+            new_params = flatten_dict(unfreeze(new_model.params))
+
+            for key in base_params.keys():
+                max_diff = (base_params[key] - new_params[key]).sum().item()
+                self.assertLessEqual(max_diff, 1e-3, msg=f"{key} not identical")
+
+
+def check_models_equal(model1, model2):
+    models_are_equal = True
+    flat_params_1 = flatten_dict(model1.params)
+    flat_params_2 = flatten_dict(model2.params)
+    for key in flat_params_1.keys():
+        if np.sum(np.abs(flat_params_1[key] - flat_params_2[key])) > 1e-4:
+            models_are_equal = False
+
+    return models_are_equal
+
+
+@require_flax
+class FlaxModelUtilsTest(unittest.TestCase):
+    def test_model_from_pretrained_subfolder(self):
+        config = BertConfig.from_pretrained("hf-internal-testing/tiny-bert-flax-only")
+        model = FlaxBertModel(config)
+
+        subfolder = "bert"
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            model.save_pretrained(os.path.join(tmp_dir, subfolder))
+
+            with self.assertRaises(OSError):
+                _ = FlaxBertModel.from_pretrained(tmp_dir)
+
+            model_loaded = FlaxBertModel.from_pretrained(tmp_dir, subfolder=subfolder)
+
+        self.assertTrue(check_models_equal(model, model_loaded))
+
+    def test_model_from_pretrained_subfolder_sharded(self):
+        config = BertConfig.from_pretrained("hf-internal-testing/tiny-bert-flax-only")
+        model = FlaxBertModel(config)
+
+        subfolder = "bert"
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            model.save_pretrained(os.path.join(tmp_dir, subfolder), max_shard_size="10KB")
+
+            with self.assertRaises(OSError):
+                _ = FlaxBertModel.from_pretrained(tmp_dir)
+
+            model_loaded = FlaxBertModel.from_pretrained(tmp_dir, subfolder=subfolder)
+
+        self.assertTrue(check_models_equal(model, model_loaded))
+
+    def test_model_from_pretrained_hub_subfolder(self):
+        subfolder = "bert"
+        model_id = "hf-internal-testing/tiny-random-bert-subfolder"
+
+        with self.assertRaises(OSError):
+            _ = FlaxBertModel.from_pretrained(model_id)
+
+        model = FlaxBertModel.from_pretrained(model_id, subfolder=subfolder)
+
+        self.assertIsNotNone(model)
+
+    def test_model_from_pretrained_hub_subfolder_sharded(self):
+        subfolder = "bert"
+        model_id = "hf-internal-testing/tiny-random-bert-sharded-subfolder"
+        with self.assertRaises(OSError):
+            _ = FlaxBertModel.from_pretrained(model_id)
+
+        model = FlaxBertModel.from_pretrained(model_id, subfolder=subfolder)
+
+        self.assertIsNotNone(model)
+
+    @require_safetensors
+    def test_safetensors_save_and_load(self):
+        model = FlaxBertModel.from_pretrained("hf-internal-testing/tiny-bert-flax-only")
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            model.save_pretrained(tmp_dir, safe_serialization=True)
+
+            # No msgpack file, only a model.safetensors
+            self.assertTrue(os.path.isfile(os.path.join(tmp_dir, SAFE_WEIGHTS_NAME)))
+            self.assertFalse(os.path.isfile(os.path.join(tmp_dir, FLAX_WEIGHTS_NAME)))
+
+            new_model = FlaxBertModel.from_pretrained(tmp_dir)
+
+        self.assertTrue(check_models_equal(model, new_model))
+
+    @require_safetensors
+    def test_safetensors_load_from_hub(self):
+        """
+        This test checks that we can load safetensors from a checkpoint that only has those on the Hub
+        """
+        flax_model = FlaxBertModel.from_pretrained("hf-internal-testing/tiny-bert-flax-only")
+
+        # Can load from the Flax-formatted checkpoint
+        safetensors_model = FlaxBertModel.from_pretrained("hf-internal-testing/tiny-bert-flax-safetensors-only")
+        self.assertTrue(check_models_equal(flax_model, safetensors_model))
+
+    @require_safetensors
+    def test_safetensors_load_from_local(self):
+        """
+        This test checks that we can load safetensors from a checkpoint that only has those on the Hub
+        """
+        with tempfile.TemporaryDirectory() as tmp:
+            location = snapshot_download("hf-internal-testing/tiny-bert-flax-only", cache_dir=tmp)
+            flax_model = FlaxBertModel.from_pretrained(location)
+
+        with tempfile.TemporaryDirectory() as tmp:
+            location = snapshot_download("hf-internal-testing/tiny-bert-flax-safetensors-only", cache_dir=tmp)
+            safetensors_model = FlaxBertModel.from_pretrained(location)
+
+        self.assertTrue(check_models_equal(flax_model, safetensors_model))
+
+    @require_safetensors
+    def test_safetensors_load_from_hub_msgpack_before_safetensors(self):
+        """
+        This test checks that we'll first download msgpack weights before safetensors
+        The safetensors file on that repo is a pt safetensors and therefore cannot be loaded without PyTorch
+        """
+        FlaxBertModel.from_pretrained("hf-internal-testing/tiny-bert-pt-safetensors-msgpack")
+
+    @require_safetensors
+    def test_safetensors_load_from_local_msgpack_before_safetensors(self):
+        """
+        This test checks that we'll first download msgpack weights before safetensors
+        The safetensors file on that repo is a pt safetensors and therefore cannot be loaded without PyTorch
+        """
+        with tempfile.TemporaryDirectory() as tmp:
+            location = snapshot_download("hf-internal-testing/tiny-bert-pt-safetensors-msgpack", cache_dir=tmp)
+            FlaxBertModel.from_pretrained(location)
+
+    @require_safetensors
+    def test_safetensors_flax_from_flax(self):
+        model = FlaxBertModel.from_pretrained("hf-internal-testing/tiny-bert-flax-only")
+
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            model.save_pretrained(tmp_dir, safe_serialization=True)
+            new_model = FlaxBertModel.from_pretrained(tmp_dir)
+
+        self.assertTrue(check_models_equal(model, new_model))
+
+    @require_safetensors
+    def test_safetensors_flax_from_sharded_msgpack_with_sharded_safetensors_local(self):
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            path = snapshot_download(
+                "hf-internal-testing/tiny-bert-flax-safetensors-msgpack-sharded", cache_dir=tmp_dir
+            )
+
+            # This should not raise even if there are two types of sharded weights
+            FlaxBertModel.from_pretrained(path)
+
+    @require_safetensors
+    def test_safetensors_flax_from_sharded_msgpack_with_sharded_safetensors_hub(self):
+        # This should not raise even if there are two types of sharded weights
+        # This should discard the safetensors weights in favor of the msgpack sharded weights
+        FlaxBertModel.from_pretrained("hf-internal-testing/tiny-bert-flax-safetensors-msgpack-sharded")
+
+    @require_safetensors
+    def test_safetensors_from_pt_bf16(self):
+        # This should not raise; should be able to load bf16-serialized torch safetensors without issue
+        # and without torch.
+        logger = logging.get_logger("transformers.modeling_flax_utils")
+
+        with CaptureLogger(logger) as cl:
+            FlaxBertModel.from_pretrained("hf-internal-testing/tiny-bert-pt-safetensors-bf16")
+
+        self.assertTrue(
+            "Some of the weights of FlaxBertModel were initialized in bfloat16 precision from the model checkpoint"
+            in cl.out
+        )

docs/transformers/tests/utils/test_modeling_rope_utils.py

@@ -0,0 +1,453 @@
+# Copyright 2024 HuggingFace Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import math
+import unittest
+
+from transformers import LlamaConfig
+from transformers.testing_utils import is_torch_available, require_torch, torch_device
+
+
+if is_torch_available():
+    import torch
+
+    from transformers import ROPE_INIT_FUNCTIONS
+    from transformers.modeling_rope_utils import rope_config_validation
+
+
+@require_torch
+class RopeTest(unittest.TestCase):
+    def test_rope_validation(self):
+        config = LlamaConfig()
+        all_rope_types = ROPE_INIT_FUNCTIONS.keys()
+
+        # The base config is always valid (default RoPE)
+        rope_config_validation(config)
+
+        # If we explicitly set the other RoPE types, then validation should fail
+        for rope_type in all_rope_types:
+            if rope_type != "default":
+                config.rope_scaling = {"rope_type": rope_type}
+                with self.assertRaises(KeyError):
+                    rope_config_validation(config)
+
+        # Parameters are exclusive to their own RoPE type, and should raise an exception if incorrectly passed
+        valid_param_mapping = {
+            "factor": ["linear", "dynamic", "yarn", "longrope"],
+            "attention_factor": ["yarn", "longrope"],
+            "beta_fast": ["yarn"],
+            "beta_slow": ["yarn"],
+            "short_factor": ["longrope"],
+            "long_factor": ["longrope"],
+        }
+        for rope_type in all_rope_types:
+            if rope_type == "default":
+                continue  # checked above
+            for param, valid_rope_types in valid_param_mapping.items():
+                # Set `param` with a dummy value -- we want to test the dict key
+                config.rope_scaling = {"rope_type": rope_type, param: True}
+                if rope_type in valid_rope_types:
+                    continue
+                else:
+                    with self.assertRaises(KeyError):
+                        rope_config_validation(config)
+
+        # Any other parameters passed to RoPE will raise a warning that a particular key is not used
+        # But sometimes we can have model-specific RoPE kwargs and bypass warning with `ignore_keys`
+        model_specific_kwarg = "mrope_sections"  # e,g in Qwen2-VL
+
+        for rope_type in all_rope_types:
+            if rope_type == "default":
+                config.rope_scaling = {"rope_type": rope_type, model_specific_kwarg: True}
+                rope_config_validation(config, ignore_keys={model_specific_kwarg})
+                with self.assertLogs("transformers.modeling_rope_utils", level="WARNING") as logs:
+                    rope_config_validation(config)
+                    self.assertEqual(len(logs.output), 1)
+                    self.assertIn(model_specific_kwarg, logs.output[0])
+
+    def test_default_rope_function_bc(self):
+        config = LlamaConfig()
+        device = torch_device
+
+        rope_kwargs = {
+            "rope_type": "default",
+            "dim": config.hidden_size // config.num_attention_heads,
+            "max_position_embeddings": config.max_position_embeddings,
+            "base": config.rope_theta,
+        }
+
+        rope_fn = ROPE_INIT_FUNCTIONS["default"]
+        config_freqs = rope_fn(config=config, device=device)[0]
+        kwargs_freqs = rope_fn(**rope_kwargs, device=device)[0]
+        torch.testing.assert_close(config_freqs, kwargs_freqs)
+
+    def test_linear_rope_function_bc(self):
+        config = LlamaConfig()
+        config.rope_scaling = {"rope_type": "linear", "factor": 10.0}
+        device = torch_device
+
+        rope_kwargs = {
+            "rope_type": "linear",
+            "dim": config.hidden_size // config.num_attention_heads,
+            "max_position_embeddings": config.max_position_embeddings,
+            "base": config.rope_theta,
+            "factor": 10.0,
+        }
+
+        rope_fn = ROPE_INIT_FUNCTIONS["linear"]
+        config_freqs = rope_fn(config=config, device=device)[0]
+        kwargs_freqs = rope_fn(**rope_kwargs, device=device)[0]
+        torch.testing.assert_close(config_freqs, kwargs_freqs)
+
+    def test_dynamic_rope_function_bc(self):
+        config = LlamaConfig()
+        config.rope_scaling = {"rope_type": "dynamic", "factor": 10.0}
+        device = torch_device
+
+        rope_kwargs = {
+            "rope_type": "dynamic",
+            "dim": config.hidden_size // config.num_attention_heads,
+            "max_position_embeddings": config.max_position_embeddings,
+            "base": config.rope_theta,
+            "factor": 10.0,
+        }
+
+        rope_fn = ROPE_INIT_FUNCTIONS["dynamic"]
+        config_freqs = rope_fn(config=config, device=device)[0]
+        kwargs_freqs = rope_fn(**rope_kwargs, device=device)[0]
+        torch.testing.assert_close(config_freqs, kwargs_freqs)
+
+    def test_default_rope_numerically(self):
+        # Note: some RoPE scaling methods start off by calling the default RoPE frequencies. If this test fails, then
+        # multiple RoPE strategies will fail.
+        # fmt: off
+        EXPECTED_INV_FREQ = torch.tensor(
+            [
+                1.0000e+00, 8.6596e-01, 7.4989e-01, 6.4938e-01, 5.6234e-01, 4.8697e-01,
+                4.2170e-01, 3.6517e-01, 3.1623e-01, 2.7384e-01, 2.3714e-01, 2.0535e-01,
+                1.7783e-01, 1.5399e-01, 1.3335e-01, 1.1548e-01, 1.0000e-01, 8.6596e-02,
+                7.4989e-02, 6.4938e-02, 5.6234e-02, 4.8697e-02, 4.2170e-02, 3.6517e-02,
+                3.1623e-02, 2.7384e-02, 2.3714e-02, 2.0535e-02, 1.7783e-02, 1.5399e-02,
+                1.3335e-02, 1.1548e-02, 1.0000e-02, 8.6596e-03, 7.4989e-03, 6.4938e-03,
+                5.6234e-03, 4.8697e-03, 4.2170e-03, 3.6517e-03, 3.1623e-03, 2.7384e-03,
+                2.3714e-03, 2.0535e-03, 1.7783e-03, 1.5399e-03, 1.3335e-03, 1.1548e-03,
+                1.0000e-03, 8.6596e-04, 7.4989e-04, 6.4938e-04, 5.6234e-04, 4.8697e-04,
+                4.2170e-04, 3.6517e-04, 3.1623e-04, 2.7384e-04, 2.3714e-04, 2.0535e-04,
+                1.7783e-04, 1.5399e-04, 1.3335e-04, 1.1548e-04
+            ], device=torch_device
+        )
+        # fmt: on
+
+        # input sanity checks: if these change, the output will also change
+        config = LlamaConfig()
+        self.assertEqual(config.rope_scaling, None)
+        self.assertEqual(config.hidden_size, 4096)
+        self.assertEqual(config.num_attention_heads, 32)
+        self.assertEqual(config.rope_theta, 10000.0)
+        self.assertFalse(hasattr(config, "partial_rotary_factor"))
+
+        rope_fn = ROPE_INIT_FUNCTIONS["default"]
+        inv_freq, attention_scale = rope_fn(config=config, device=torch_device)
+
+        self.assertEqual(attention_scale, 1.0)  # attention scale is always 1 for default RoPE
+        torch.testing.assert_close(inv_freq, EXPECTED_INV_FREQ)
+
+    def test_linear_rope_numerically(self):
+        # This is a linear scaling strategy, the **frequencies** are scaled linearly with respect to the default
+        # frequencies (= the inverse frequencies are scaled **inversely**)
+        config = LlamaConfig()
+        default_rope_fn = ROPE_INIT_FUNCTIONS["default"]
+        default_inv_freq, _ = default_rope_fn(config=config, device=torch_device)
+
+        rope_fn = ROPE_INIT_FUNCTIONS["linear"]
+        for factor in (2.0, 10.0, 20.0):
+            config.rope_scaling = {"rope_type": "linear", "factor": factor}
+            inv_freq, attention_scale = rope_fn(config=config, device=torch_device)
+            self.assertEqual(attention_scale, 1.0)  # attention scale is always 1 for linear RoPE
+            torch.testing.assert_close(inv_freq, default_inv_freq / factor)
+
+    def test_dynamic_rope_numerically(self):
+        # fmt: off
+        EXPECTED_INV_FREQ = torch.tensor(
+            [
+                1.0000e+00, 8.0931e-01, 6.5498e-01, 5.3008e-01, 4.2900e-01, 3.4720e-01,
+                2.8099e-01, 2.2741e-01, 1.8404e-01, 1.4895e-01, 1.2055e-01, 9.7558e-02,
+                7.8955e-02, 6.3899e-02, 5.1714e-02, 4.1853e-02, 3.3872e-02, 2.7413e-02,
+                2.2185e-02, 1.7955e-02, 1.4531e-02, 1.1760e-02, 9.5176e-03, 7.7027e-03,
+                6.2339e-03, 5.0451e-03, 4.0831e-03, 3.3045e-03, 2.6744e-03, 2.1644e-03,
+                1.7517e-03, 1.4176e-03, 1.1473e-03, 9.2852e-04, 7.5146e-04, 6.0817e-04,
+                4.9220e-04, 3.9834e-04, 3.2238e-04, 2.6091e-04, 2.1115e-04, 1.7089e-04,
+                1.3830e-04, 1.1193e-04, 9.0585e-05, 7.3312e-05, 5.9332e-05, 4.8018e-05,
+                3.8861e-05, 3.1451e-05, 2.5453e-05, 2.0600e-05, 1.6672e-05, 1.3492e-05,
+                1.0920e-05, 8.8374e-06, 7.1522e-06, 5.7883e-06, 4.6845e-06, 3.7912e-06,
+                3.0683e-06, 2.4832e-06, 2.0097e-06, 1.6265e-06
+            ], device=torch_device
+        )
+        # fmt: on
+
+        # input sanity checks: if these change, the output will also change
+        config = LlamaConfig()
+        self.assertEqual(config.rope_scaling, None)
+        self.assertEqual(config.hidden_size, 4096)
+        self.assertEqual(config.num_attention_heads, 32)
+        self.assertEqual(config.rope_theta, 10000.0)
+        self.assertFalse(hasattr(config, "partial_rotary_factor"))
+
+        rope_fn = ROPE_INIT_FUNCTIONS["default"]
+        default_inv_freq, _ = rope_fn(config=config, device=torch_device)
+
+        # Check 1: this is a dynamic scaling strategy, it will not scale unless we provide `seq_len` larger than the
+        # model's original training sequence length
+        rope_fn = ROPE_INIT_FUNCTIONS["dynamic"]
+        for factor in (2.0, 10.0, 20.0):
+            config.rope_scaling = {"rope_type": "dynamic", "factor": factor}
+            inv_freq, attention_scale = rope_fn(config=config, device=torch_device)
+            self.assertEqual(attention_scale, 1.0)  # attention scale is always 1 for dynamic RoPE
+            torch.testing.assert_close(inv_freq, default_inv_freq)
+
+            inv_freq, _ = rope_fn(config=config, device=torch_device, seq_len=1)
+            torch.testing.assert_close(inv_freq, default_inv_freq)
+
+        # Check 2: if we provide `seq_len` larger than the model's original training sequence length, the frequencies
+        # will scale up (i.e., the inverse frequencies will scale down).
+        factor = 10.0
+        config.rope_scaling = {"rope_type": "dynamic", "factor": factor}
+        inv_freq, _ = rope_fn(config=config, device=torch_device, seq_len=16384)
+        with self.assertRaises(AssertionError):  # It is NOT a linear factor
+            torch.testing.assert_close(inv_freq, default_inv_freq / factor)
+        torch.testing.assert_close(inv_freq, EXPECTED_INV_FREQ)
+
+    def test_yarn_rope_numerically(self):
+        # fmt: off
+        EXPECTED_INV_FREQ = torch.tensor(
+            [
+                1.0000e+00, 8.6596e-01, 7.4989e-01, 6.4938e-01, 5.6234e-01, 4.8697e-01,
+                4.2170e-01, 3.6517e-01, 3.1623e-01, 2.7384e-01, 2.3714e-01, 2.0535e-01,
+                1.7783e-01, 1.5399e-01, 1.3335e-01, 1.1548e-01, 1.0000e-01, 8.3479e-02,
+                6.9590e-02, 5.7925e-02, 4.8136e-02, 3.9931e-02, 3.3061e-02, 2.7315e-02,
+                2.2515e-02, 1.8512e-02, 1.5177e-02, 1.2403e-02, 1.0101e-02, 8.1924e-03,
+                6.6143e-03, 5.3120e-03, 4.2400e-03, 3.3599e-03, 2.6396e-03, 2.0520e-03,
+                1.5746e-03, 1.1882e-03, 8.7713e-04, 6.2810e-04, 4.3007e-04, 2.7384e-04,
+                2.3714e-04, 2.0535e-04, 1.7783e-04, 1.5399e-04, 1.3335e-04, 1.1548e-04,
+                1.0000e-04, 8.6596e-05, 7.4989e-05, 6.4938e-05, 5.6234e-05, 4.8697e-05,
+                4.2170e-05, 3.6517e-05, 3.1623e-05, 2.7384e-05, 2.3714e-05, 2.0535e-05,
+                1.7783e-05, 1.5399e-05, 1.3335e-05, 1.1548e-05
+            ], device=torch_device
+        )
+        # fmt: on
+
+        # input sanity checks: if these change, the output will also change
+        config = LlamaConfig()
+        self.assertEqual(config.rope_scaling, None)
+        self.assertEqual(config.hidden_size, 4096)
+        self.assertEqual(config.num_attention_heads, 32)
+        self.assertEqual(config.rope_theta, 10000.0)
+        self.assertFalse(hasattr(config, "partial_rotary_factor"))
+
+        rope_fn = ROPE_INIT_FUNCTIONS["default"]
+        default_inv_freq, _ = rope_fn(config=config, device=torch_device)
+
+        # Check 1: according to the paper, if `attention_factor` is not specified, then it has a specific default --
+        # `0.1 * math.log(factor) + 1.0`
+        rope_fn = ROPE_INIT_FUNCTIONS["yarn"]
+        for factor in (2.0, 10.0, 20.0):
+            config.rope_scaling = {"rope_type": "yarn", "factor": factor}
+            _, attention_scale = rope_fn(config=config, device=torch_device)
+            self.assertEqual(attention_scale, 0.1 * math.log(factor) + 1.0)
+
+            config.rope_scaling = {"rope_type": "yarn", "factor": factor, "attention_factor": 0.5}
+            _, attention_scale = rope_fn(config=config, device=torch_device, seq_len=1)
+            self.assertEqual(attention_scale, 0.5)
+
+        # Check 2: based on `beta_fast` and `beta_slow`, the frequencies will be scaled between 1 and `factor`.
+        # Increasing `beta_fast` will make RoPE more interpolative (apply scaling), and the other way around.
+        # `beta_slow` behaves the opposite way. Remember: `beta_fast` > `beta_slow`
+        # (note: adds a margin to the test for numerical stability)
+        factor = 10.0
+        margin = 1e-8
+        config.rope_scaling = {"rope_type": "yarn", "factor": factor, "beta_fast": 32, "beta_slow": 1}
+        inv_freq, _ = rope_fn(config=config, device=torch_device)
+        is_bounded_by_factor = [
+            ((default_inv_freq[idx] / factor) - margin) <= yarn_inv_freq_value <= (default_inv_freq[idx] + margin)
+            for idx, yarn_inv_freq_value in enumerate(inv_freq)
+        ]
+        self.assertTrue(all(is_bounded_by_factor))
+
+        # super high beta_fast = interpolation (i.e. scaling) in all but the first inverse frequency. The last ~20
+        # values (empirically checked for `beta_fast` = 1000) should be very small to linear scaling
+        config.rope_scaling = {"rope_type": "yarn", "factor": factor, "beta_fast": 1000, "beta_slow": 1}
+        inv_freq, _ = rope_fn(config=config, device=torch_device)
+        is_interpolating = [
+            yarn_inv_freq_value < (default_inv_freq[idx] + margin) for idx, yarn_inv_freq_value in enumerate(inv_freq)
+        ]
+        self.assertFalse(is_interpolating[0])
+        self.assertTrue(all(is_interpolating[1:]))
+        torch.testing.assert_close(inv_freq[-20:], default_inv_freq[-20:] / factor)
+
+        # Check 3: numerical snapshot to avoid regressions
+        config.rope_scaling = {"rope_type": "yarn", "factor": factor, "beta_fast": 32, "beta_slow": 1}
+        inv_freq, _ = rope_fn(config=config, device=torch_device)
+        torch.testing.assert_close(inv_freq, EXPECTED_INV_FREQ)
+
+    def test_longrope_rope_numerically(self):
+        # input sanity checks: if these change, the output will also change
+        config = LlamaConfig()
+        self.assertEqual(config.rope_scaling, None)
+        self.assertEqual(config.hidden_size, 4096)
+        self.assertEqual(config.num_attention_heads, 32)
+        self.assertEqual(config.rope_theta, 10000.0)
+        self.assertFalse(hasattr(config, "partial_rotary_factor"))
+
+        # longrope applies scaling on EACH inv frequency, `short_factor` or `long_factor`, depending on the seq_len
+        dim = config.hidden_size // config.num_attention_heads
+        short_factor = [2.0] * (dim // 2)  # scaling applied when seq_len <= max_position_embeddings
+        long_factor = torch.ones(dim // 2).cumsum(0).tolist()  # scaling applied when seq_len > max_position_embeddings
+
+        rope_fn = ROPE_INIT_FUNCTIONS["default"]
+        default_inv_freq, _ = rope_fn(config=config, device=torch_device)
+
+        # Check 1: according to the paper, if `attention_factor` is not specified, then it has a specific default --
+        # `math.sqrt(1 + math.log(factor) / math.log(max_position_embeddings))`
+        rope_fn = ROPE_INIT_FUNCTIONS["longrope"]
+        max_position_embeddings = config.max_position_embeddings
+        for factor in (2.0, 10.0, 20.0):
+            config.rope_scaling = {
+                "rope_type": "longrope",
+                "factor": factor,
+                "short_factor": short_factor,
+                "long_factor": long_factor,
+            }
+            _, attention_scale = rope_fn(config=config, device=torch_device)
+            self.assertEqual(attention_scale, math.sqrt(1 + math.log(factor) / math.log(max_position_embeddings)))
+
+            config.rope_scaling = {
+                "rope_type": "longrope",
+                "factor": factor,
+                "short_factor": short_factor,
+                "long_factor": long_factor,
+                "attention_factor": 0.5,
+            }
+            _, attention_scale = rope_fn(config=config, device=torch_device, seq_len=1)
+            self.assertEqual(attention_scale, 0.5)
+
+            config.rope_scaling = {
+                "rope_type": "longrope",
+                "factor": factor,
+                "short_factor": short_factor,
+                "long_factor": long_factor,
+            }
+            self.assertEqual(config.rope_scaling.get("attention_factor"), None)
+            # Verify that "TypeError: '<' not supported between instances of 'NoneType' and 'int'" is not raised.
+            rope_config_validation(config)
+
+        # Check 2: seq_len == 0 -> short factor is applied to the default frequencies
+        config.rope_scaling = {
+            "rope_type": "longrope",
+            "factor": 1.0,
+            "short_factor": short_factor,
+            "long_factor": long_factor,
+        }
+        inv_freq, _ = rope_fn(config=config, device=torch_device, seq_len=0)
+        torch.testing.assert_close(inv_freq, default_inv_freq / torch.tensor(short_factor).to(torch_device))
+
+        # Check 3: seq_len > max_position_embeddings -> long factor is applied to the default frequencies
+        inv_freq, _ = rope_fn(config=config, device=torch_device, seq_len=config.max_position_embeddings + 1)
+        torch.testing.assert_close(inv_freq, default_inv_freq / torch.tensor(long_factor).to(torch_device))
+
+    def test_llama3_rope_numerically(self):
+        # fmt: off
+        EXPECTED_INV_FREQ = torch.tensor(
+            [
+                1.0000e+00, 8.6596e-01, 7.4989e-01, 6.4938e-01, 5.6234e-01, 4.8697e-01,
+                4.2170e-01, 3.6517e-01, 3.1623e-01, 2.7384e-01, 2.3714e-01, 2.0535e-01,
+                1.7783e-01, 1.5399e-01, 1.3335e-01, 1.1548e-01, 1.0000e-01, 8.6596e-02,
+                7.4989e-02, 6.4938e-02, 5.6234e-02, 4.8697e-02, 4.2170e-02, 3.6517e-02,
+                3.1623e-02, 2.7384e-02, 2.3714e-02, 2.0535e-02, 1.7783e-02, 1.5399e-02,
+                1.3335e-02, 1.0730e-02, 7.7785e-03, 5.6009e-03, 3.9991e-03, 2.8248e-03,
+                1.9675e-03, 1.3449e-03, 8.9549e-04, 5.7363e-04, 3.4539e-04, 2.7384e-04,
+                2.3714e-04, 2.0535e-04, 1.7783e-04, 1.5399e-04, 1.3335e-04, 1.1548e-04,
+                1.0000e-04, 8.6596e-05, 7.4989e-05, 6.4938e-05, 5.6234e-05, 4.8697e-05,
+                4.2170e-05, 3.6517e-05, 3.1623e-05, 2.7384e-05, 2.3714e-05, 2.0535e-05,
+                1.7783e-05, 1.5399e-05, 1.3335e-05, 1.1548e-05
+            ], device=torch_device
+        )
+        # fmt: on
+
+        # input sanity checks: if these change, the output will also change
+        config = LlamaConfig()
+        self.assertEqual(config.rope_scaling, None)
+        self.assertEqual(config.hidden_size, 4096)
+        self.assertEqual(config.num_attention_heads, 32)
+        self.assertEqual(config.rope_theta, 10000.0)
+        self.assertFalse(hasattr(config, "partial_rotary_factor"))
+
+        rope_fn = ROPE_INIT_FUNCTIONS["default"]
+        default_inv_freq, _ = rope_fn(config=config, device=torch_device)
+
+        # Check 1: `attention_factor` is always 1
+        rope_fn = ROPE_INIT_FUNCTIONS["llama3"]
+        for factor in (2.0, 10.0, 20.0):
+            config.rope_scaling = {
+                "rope_type": "llama3",
+                "factor": factor,
+                "original_max_position_embeddings": 2048,
+                "low_freq_factor": 1,
+                "high_freq_factor": 4,
+            }
+            _, attention_scale = rope_fn(config=config, device=torch_device)
+            self.assertEqual(attention_scale, 1.0)
+
+        # Check 2: based on `low_freq_factor` and `high_freq_factor`, the frequencies will be scaled between 1 and
+        # `factor` (similar to yarn). Low frequencies get scaled by `factor`, high frequencies see no change, medium
+        # frequencies are scaled by a value in between. Changing `low_freq_factor` and `high_freq_factor` changes what
+        # is considered low, medium, and high frequencies.
+        factor = 10.0
+        config.rope_scaling = {
+            "rope_type": "llama3",
+            "factor": factor,
+            "original_max_position_embeddings": 2048,
+            "low_freq_factor": 1,
+            "high_freq_factor": 4,
+        }
+        inv_freq, _ = rope_fn(config=config, device=torch_device)
+        is_bounded_by_factor = [
+            (default_inv_freq[idx] / factor) <= llama3_inv_freq_value <= default_inv_freq[idx]
+            for idx, llama3_inv_freq_value in enumerate(inv_freq)
+        ]
+        self.assertTrue(all(is_bounded_by_factor))
+
+        # if we change `high_freq_factor` to a very high value, none is considered high-frequency -> ALL values will be
+        # scaled
+        config.rope_scaling = config.rope_scaling = {
+            "rope_type": "llama3",
+            "factor": factor,
+            "original_max_position_embeddings": 2048,
+            "low_freq_factor": 1,
+            "high_freq_factor": 1000,
+        }
+        inv_freq, _ = rope_fn(config=config, device=torch_device)
+        is_scaled = [yarn_inv_freq_value < default_inv_freq[idx] for idx, yarn_inv_freq_value in enumerate(inv_freq)]
+        self.assertTrue(all(is_scaled))
+
+        # Check 3: numerical snapshot to avoid regressions
+        config.rope_scaling = {
+            "rope_type": "llama3",
+            "factor": factor,
+            "original_max_position_embeddings": 2048,
+            "low_freq_factor": 1,
+            "high_freq_factor": 4,
+        }
+        inv_freq, _ = rope_fn(config=config, device=torch_device)
+        torch.testing.assert_close(inv_freq, EXPECTED_INV_FREQ)

docs/transformers/tests/utils/test_modeling_tf_core.py

@@ -0,0 +1,403 @@
+# Copyright 2019 HuggingFace Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from __future__ import annotations
+
+import copy
+import os
+import tempfile
+from importlib import import_module
+from math import isnan
+
+from transformers import is_tf_available
+from transformers.models.auto import get_values
+from transformers.testing_utils import require_tf, slow
+
+from ..test_modeling_tf_common import ids_tensor
+
+
+if is_tf_available():
+    import numpy as np
+    import tensorflow as tf
+
+    from transformers import (
+        TF_MODEL_FOR_CAUSAL_LM_MAPPING,
+        TF_MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING,
+        TF_MODEL_FOR_MASKED_LM_MAPPING,
+        TF_MODEL_FOR_MULTIPLE_CHOICE_MAPPING,
+        TF_MODEL_FOR_NEXT_SENTENCE_PREDICTION_MAPPING,
+        TF_MODEL_FOR_PRETRAINING_MAPPING,
+        TF_MODEL_FOR_QUESTION_ANSWERING_MAPPING,
+        TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING,
+        TF_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING,
+        TF_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING,
+        TFSharedEmbeddings,
+    )
+    from transformers.modeling_tf_utils import keras
+
+
+@require_tf
+class TFCoreModelTesterMixin:
+    model_tester = None
+    all_model_classes = ()
+    all_generative_model_classes = ()
+    test_mismatched_shapes = True
+    test_resize_embeddings = True
+    test_head_masking = True
+    is_encoder_decoder = False
+
+    def _prepare_for_class(self, inputs_dict, model_class, return_labels=False) -> dict:
+        inputs_dict = copy.deepcopy(inputs_dict)
+
+        if model_class in get_values(TF_MODEL_FOR_MULTIPLE_CHOICE_MAPPING):
+            inputs_dict = {
+                k: tf.tile(tf.expand_dims(v, 1), (1, self.model_tester.num_choices) + (1,) * (v.ndim - 1))
+                if isinstance(v, tf.Tensor) and v.ndim > 0
+                else v
+                for k, v in inputs_dict.items()
+            }
+
+        if return_labels:
+            if model_class in get_values(TF_MODEL_FOR_MULTIPLE_CHOICE_MAPPING):
+                inputs_dict["labels"] = tf.ones(self.model_tester.batch_size, dtype=tf.int32)
+            elif model_class in get_values(TF_MODEL_FOR_QUESTION_ANSWERING_MAPPING):
+                inputs_dict["start_positions"] = tf.zeros(self.model_tester.batch_size, dtype=tf.int32)
+                inputs_dict["end_positions"] = tf.zeros(self.model_tester.batch_size, dtype=tf.int32)
+            elif model_class in [
+                *get_values(TF_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING),
+                *get_values(TF_MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING),
+            ]:
+                inputs_dict["labels"] = tf.zeros(self.model_tester.batch_size, dtype=tf.int32)
+            elif model_class in get_values(TF_MODEL_FOR_NEXT_SENTENCE_PREDICTION_MAPPING):
+                inputs_dict["next_sentence_label"] = tf.zeros(self.model_tester.batch_size, dtype=tf.int32)
+            elif model_class in [
+                *get_values(TF_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING),
+                *get_values(TF_MODEL_FOR_CAUSAL_LM_MAPPING),
+                *get_values(TF_MODEL_FOR_MASKED_LM_MAPPING),
+                *get_values(TF_MODEL_FOR_PRETRAINING_MAPPING),
+                *get_values(TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING),
+            ]:
+                inputs_dict["labels"] = tf.zeros(
+                    (self.model_tester.batch_size, self.model_tester.seq_length), dtype=tf.int32
+                )
+        return inputs_dict
+
+    @slow
+    def test_graph_mode(self):
+        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
+        for model_class in self.all_model_classes[:2]:
+            inputs = self._prepare_for_class(inputs_dict, model_class)
+            model = model_class(config)
+
+            @tf.function
+            def run_in_graph_mode():
+                return model(inputs)
+
+            outputs = run_in_graph_mode()
+            self.assertIsNotNone(outputs)
+
+    @slow
+    def test_xla_mode(self):
+        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
+        for model_class in self.all_model_classes[:2]:
+            inputs = self._prepare_for_class(inputs_dict, model_class)
+            model = model_class(config)
+
+            @tf.function(experimental_compile=True)
+            def run_in_graph_mode():
+                return model(inputs)
+
+            outputs = run_in_graph_mode()
+            self.assertIsNotNone(outputs)
+
+    @slow
+    def test_xla_fit(self):
+        # This is a copy of the test_keras_fit method, but we use XLA compilation instead of eager
+        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
+        for model_class in self.all_model_classes[:2]:
+            model = model_class(config)
+            if getattr(model, "hf_compute_loss", None):
+                # Test that model correctly compute the loss with kwargs
+                prepared_for_class = self._prepare_for_class(inputs_dict.copy(), model_class, return_labels=True)
+                # Is there a better way to remove these decoder inputs?
+                prepared_for_class = {
+                    key: val
+                    for key, val in prepared_for_class.items()
+                    if key not in ("head_mask", "decoder_head_mask", "cross_attn_head_mask", "decoder_input_ids")
+                }
+
+                possible_label_cols = {
+                    "labels",
+                    "label",
+                    "label_ids",
+                    "start_positions",
+                    "start_position",
+                    "end_positions",
+                    "end_position",
+                    "next_sentence_label",
+                }
+                label_names = possible_label_cols.intersection(set(prepared_for_class))
+                self.assertGreater(len(label_names), 0, msg="No matching label names found!")
+                labels = {key: val for key, val in prepared_for_class.items() if key in label_names}
+                inputs_minus_labels = {key: val for key, val in prepared_for_class.items() if key not in label_names}
+                self.assertGreater(len(inputs_minus_labels), 0)
+
+                # Make sure it works with XLA!
+                model.compile(optimizer=keras.optimizers.SGD(0.0), jit_compile=True)
+                # Make sure the model fits without crashing regardless of where we pass the labels
+                history = model.fit(
+                    prepared_for_class,
+                    validation_data=prepared_for_class,
+                    steps_per_epoch=1,
+                    validation_steps=1,
+                    shuffle=False,
+                    verbose=0,
+                )
+                loss = history.history["loss"][0]
+                self.assertTrue(not isnan(loss))
+                val_loss = history.history["val_loss"][0]
+                self.assertTrue(not isnan(val_loss))
+
+                # Now test it with separate labels, to make sure that path works in XLA too.
+                model = model_class(config)
+                model.compile(optimizer=keras.optimizers.SGD(0.0), jit_compile=True)
+                history = model.fit(
+                    inputs_minus_labels,
+                    labels,
+                    validation_data=(inputs_minus_labels, labels),
+                    steps_per_epoch=1,
+                    validation_steps=1,
+                    shuffle=False,
+                    verbose=0,
+                )
+
+                loss = history.history["loss"][0]
+                self.assertTrue(not isnan(loss))
+                val_loss = history.history["val_loss"][0]
+                self.assertTrue(not isnan(val_loss))
+
+    @slow
+    def test_saved_model_creation_extended(self):
+        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
+        config.output_hidden_states = True
+        config.output_attentions = True
+
+        if hasattr(config, "use_cache"):
+            config.use_cache = True
+
+        encoder_seq_length = getattr(self.model_tester, "encoder_seq_length", self.model_tester.seq_length)
+        encoder_key_length = getattr(self.model_tester, "key_length", encoder_seq_length)
+
+        for model_class in self.all_model_classes[:2]:
+            class_inputs_dict = self._prepare_for_class(inputs_dict, model_class)
+            model = model_class(config)
+            model.build_in_name_scope()
+            num_out = len(model(class_inputs_dict))
+
+            for key in list(class_inputs_dict.keys()):
+                # Remove keys not in the serving signature, as the SavedModel will not be compiled to deal with them
+                if key not in model.input_signature:
+                    del class_inputs_dict[key]
+                # Check it's a tensor, in case the inputs dict has some bools in it too
+                elif isinstance(class_inputs_dict[key], tf.Tensor) and class_inputs_dict[key].dtype.is_integer:
+                    class_inputs_dict[key] = tf.cast(class_inputs_dict[key], tf.int32)
+
+            if set(class_inputs_dict.keys()) != set(model.input_signature.keys()):
+                continue  # Some models have inputs that the preparation functions don't create, we skip those
+
+            with tempfile.TemporaryDirectory() as tmpdirname:
+                model.save_pretrained(tmpdirname, saved_model=True)
+                saved_model_dir = os.path.join(tmpdirname, "saved_model", "1")
+                model = keras.models.load_model(saved_model_dir)
+                outputs = model(class_inputs_dict)
+
+                if self.is_encoder_decoder:
+                    output_hidden_states = outputs["encoder_hidden_states"]
+                    output_attentions = outputs["encoder_attentions"]
+                else:
+                    output_hidden_states = outputs["hidden_states"]
+                    output_attentions = outputs["attentions"]
+
+                self.assertEqual(len(outputs), num_out)
+
+                expected_num_layers = getattr(
+                    self.model_tester, "expected_num_hidden_layers", self.model_tester.num_hidden_layers + 1
+                )
+
+                self.assertEqual(len(output_hidden_states), expected_num_layers)
+                self.assertListEqual(
+                    list(output_hidden_states[0].shape[-2:]),
+                    [self.model_tester.seq_length, self.model_tester.hidden_size],
+                )
+
+                self.assertEqual(len(output_attentions), self.model_tester.num_hidden_layers)
+                self.assertListEqual(
+                    list(output_attentions[0].shape[-3:]),
+                    [self.model_tester.num_attention_heads, encoder_seq_length, encoder_key_length],
+                )
+
+    @slow
+    def test_mixed_precision(self):
+        keras.mixed_precision.set_global_policy("mixed_float16")
+
+        # try/finally block to ensure subsequent tests run in float32
+        try:
+            config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
+            for model_class in self.all_model_classes[:2]:
+                class_inputs_dict = self._prepare_for_class(inputs_dict, model_class)
+                model = model_class(config)
+                outputs = model(class_inputs_dict)
+
+                self.assertIsNotNone(outputs)
+        finally:
+            keras.mixed_precision.set_global_policy("float32")
+
+    @slow
+    def test_train_pipeline_custom_model(self):
+        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
+        # head_mask and decoder_head_mask has different shapes than other input args
+        if "head_mask" in inputs_dict:
+            del inputs_dict["head_mask"]
+        if "decoder_head_mask" in inputs_dict:
+            del inputs_dict["decoder_head_mask"]
+        if "cross_attn_head_mask" in inputs_dict:
+            del inputs_dict["cross_attn_head_mask"]
+        tf_main_layer_classes = {
+            module_member
+            for model_class in self.all_model_classes
+            for module in (import_module(model_class.__module__),)
+            for module_member_name in dir(module)
+            if module_member_name.endswith("MainLayer")
+            for module_member in (getattr(module, module_member_name),)
+            if isinstance(module_member, type)
+            and keras.layers.Layer in module_member.__bases__
+            and getattr(module_member, "_keras_serializable", False)
+        }
+
+        for main_layer_class in tf_main_layer_classes:
+            # T5MainLayer needs an embed_tokens parameter when called without the inputs_embeds parameter
+            if "T5" in main_layer_class.__name__:
+                # Take the same values than in TFT5ModelTester for this shared layer
+                shared = TFSharedEmbeddings(self.model_tester.vocab_size, self.model_tester.hidden_size, name="shared")
+                config.use_cache = False
+                main_layer = main_layer_class(config, embed_tokens=shared)
+            else:
+                main_layer = main_layer_class(config)
+
+            symbolic_inputs = {
+                name: keras.Input(tensor.shape[1:], dtype=tensor.dtype) for name, tensor in inputs_dict.items()
+            }
+
+            if hasattr(self.model_tester, "num_labels"):
+                num_labels = self.model_tester.num_labels
+            else:
+                num_labels = 2
+
+            X = tf.data.Dataset.from_tensor_slices(
+                (inputs_dict, np.ones((self.model_tester.batch_size, self.model_tester.seq_length, num_labels, 1)))
+            ).batch(1)
+
+            hidden_states = main_layer(symbolic_inputs)[0]
+            outputs = keras.layers.Dense(num_labels, activation="softmax", name="outputs")(hidden_states)
+            model = keras.models.Model(inputs=symbolic_inputs, outputs=[outputs])
+
+            model.compile(loss="binary_crossentropy", optimizer="adam", metrics=["binary_accuracy"])
+            model.fit(X, epochs=1)
+
+            with tempfile.TemporaryDirectory() as tmpdirname:
+                filepath = os.path.join(tmpdirname, "keras_model.h5")
+                model.save(filepath)
+                if "T5" in main_layer_class.__name__:
+                    model = keras.models.load_model(
+                        filepath,
+                        custom_objects={
+                            main_layer_class.__name__: main_layer_class,
+                            "TFSharedEmbeddings": TFSharedEmbeddings,
+                        },
+                    )
+                else:
+                    model = keras.models.load_model(
+                        filepath, custom_objects={main_layer_class.__name__: main_layer_class}
+                    )
+                assert isinstance(model, keras.Model)
+                model(inputs_dict)
+
+    @slow
+    def test_graph_mode_with_inputs_embeds(self):
+        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
+
+        for model_class in self.all_model_classes[:2]:
+            model = model_class(config)
+
+            inputs = copy.deepcopy(inputs_dict)
+
+            if not self.is_encoder_decoder:
+                input_ids = inputs["input_ids"]
+                del inputs["input_ids"]
+            else:
+                encoder_input_ids = inputs["input_ids"]
+                decoder_input_ids = inputs.get("decoder_input_ids", encoder_input_ids)
+                del inputs["input_ids"]
+                inputs.pop("decoder_input_ids", None)
+
+            if not self.is_encoder_decoder:
+                inputs["inputs_embeds"] = model.get_input_embeddings()(input_ids)
+            else:
+                inputs["inputs_embeds"] = model.get_input_embeddings()(encoder_input_ids)
+                inputs["decoder_inputs_embeds"] = model.get_input_embeddings()(decoder_input_ids)
+
+            inputs = self._prepare_for_class(inputs, model_class)
+
+            @tf.function
+            def run_in_graph_mode():
+                return model(inputs)
+
+            outputs = run_in_graph_mode()
+            self.assertIsNotNone(outputs)
+
+    def _generate_random_bad_tokens(self, num_bad_tokens, model):
+        # special tokens cannot be bad tokens
+        special_tokens = []
+        if model.config.bos_token_id is not None:
+            special_tokens.append(model.config.bos_token_id)
+        if model.config.pad_token_id is not None:
+            special_tokens.append(model.config.pad_token_id)
+        if model.config.eos_token_id is not None:
+            special_tokens.append(model.config.eos_token_id)
+
+        # create random bad tokens that are not special tokens
+        bad_tokens = []
+        while len(bad_tokens) < num_bad_tokens:
+            token = tf.squeeze(ids_tensor((1, 1), self.model_tester.vocab_size), 0).numpy()[0]
+            if token not in special_tokens:
+                bad_tokens.append(token)
+        return bad_tokens
+
+    def _check_generated_ids(self, output_ids):
+        for token_id in output_ids[0].numpy().tolist():
+            self.assertGreaterEqual(token_id, 0)
+            self.assertLess(token_id, self.model_tester.vocab_size)
+
+    def _check_match_tokens(self, generated_ids, bad_words_ids):
+        # for all bad word tokens
+        for bad_word_ids in bad_words_ids:
+            # for all slices in batch
+            for generated_ids_slice in generated_ids:
+                # for all word idx
+                for i in range(len(bad_word_ids), len(generated_ids_slice)):
+                    # if tokens match
+                    if generated_ids_slice[i - len(bad_word_ids) : i] == bad_word_ids:
+                        return True
+        return False

docs/transformers/tests/utils/test_modeling_tf_utils.py

@@ -0,0 +1,662 @@
+# Copyright 2019 HuggingFace Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from __future__ import annotations
+
+import json
+import os
+import random
+import tempfile
+import unittest
+import unittest.mock as mock
+
+from huggingface_hub import HfFolder, snapshot_download
+from requests.exceptions import HTTPError
+
+from transformers import is_tf_available
+from transformers.configuration_utils import PretrainedConfig
+from transformers.testing_utils import (  # noqa: F401
+    TOKEN,
+    USER,
+    CaptureLogger,
+    TemporaryHubRepo,
+    is_staging_test,
+    require_safetensors,
+    require_tf,
+    slow,
+)
+from transformers.utils import (
+    SAFE_WEIGHTS_INDEX_NAME,
+    SAFE_WEIGHTS_NAME,
+    TF2_WEIGHTS_INDEX_NAME,
+    TF2_WEIGHTS_NAME,
+    logging,
+)
+
+
+logger = logging.get_logger(__name__)
+
+
+if is_tf_available():
+    import h5py
+    import numpy as np
+    import tensorflow as tf
+
+    from transformers import (
+        BertConfig,
+        RagRetriever,
+        TFBertForSequenceClassification,
+        TFBertModel,
+        TFRagModel,
+    )
+    from transformers.modeling_tf_utils import keras, tf_shard_checkpoint, unpack_inputs
+    from transformers.tf_utils import stable_softmax
+
+    tf.config.experimental.enable_tensor_float_32_execution(False)
+
+
+@require_tf
+class TFModelUtilsTest(unittest.TestCase):
+    def test_cached_files_are_used_when_internet_is_down(self):
+        # A mock response for an HTTP head request to emulate server down
+        response_mock = mock.Mock()
+        response_mock.status_code = 500
+        response_mock.headers = {}
+        response_mock.raise_for_status.side_effect = HTTPError
+        response_mock.json.return_value = {}
+
+        # Download this model to make sure it's in the cache.
+        _ = TFBertModel.from_pretrained("hf-internal-testing/tiny-random-bert")
+
+        # Under the mock environment we get a 500 error when trying to reach the model.
+        with mock.patch("requests.Session.request", return_value=response_mock) as mock_head:
+            _ = TFBertModel.from_pretrained("hf-internal-testing/tiny-random-bert")
+            # This check we did call the fake head request
+            mock_head.assert_called()
+
+    # tests whether the unpack_inputs function behaves as expected
+    def test_unpack_inputs(self):
+        class DummyModel:
+            def __init__(self):
+                config_kwargs = {"output_attentions": False, "output_hidden_states": False, "return_dict": False}
+                self.config = PretrainedConfig(**config_kwargs)
+                self.main_input_name = "input_ids"
+
+            @unpack_inputs
+            def call(
+                self,
+                input_ids=None,
+                past_key_values=None,
+                output_attentions=None,
+                output_hidden_states=None,
+                return_dict=None,
+            ):
+                return input_ids, past_key_values, output_attentions, output_hidden_states, return_dict
+
+            @unpack_inputs
+            def foo(self, pixel_values, output_attentions=None, output_hidden_states=None, return_dict=None):
+                return pixel_values, output_attentions, output_hidden_states, return_dict
+
+        dummy_model = DummyModel()
+        input_ids = tf.constant([0, 1, 2, 3], dtype=tf.int32)
+        past_key_values = tf.constant([4, 5, 6, 7], dtype=tf.int32)
+        pixel_values = tf.constant([8, 9, 10, 11], dtype=tf.int32)
+
+        # test case 1: Pass inputs as keyword arguments; Booleans are inherited from the config.
+        output = dummy_model.call(input_ids=input_ids, past_key_values=past_key_values)
+        tf.debugging.assert_equal(output[0], input_ids)
+        tf.debugging.assert_equal(output[1], past_key_values)
+        self.assertFalse(output[2])
+        self.assertFalse(output[3])
+        self.assertFalse(output[4])
+
+        # test case 2: Same as above, but with positional arguments.
+        output = dummy_model.call(input_ids, past_key_values)
+        tf.debugging.assert_equal(output[0], input_ids)
+        tf.debugging.assert_equal(output[1], past_key_values)
+        self.assertFalse(output[2])
+        self.assertFalse(output[3])
+        self.assertFalse(output[4])
+
+        # test case 3: We can also pack everything in the first input.
+        output = dummy_model.call(input_ids={"input_ids": input_ids, "past_key_values": past_key_values})
+        tf.debugging.assert_equal(output[0], input_ids)
+        tf.debugging.assert_equal(output[1], past_key_values)
+        self.assertFalse(output[2])
+        self.assertFalse(output[3])
+        self.assertFalse(output[4])
+
+        # test case 4: Explicit boolean arguments should override the config.
+        output = dummy_model.call(
+            input_ids=input_ids, past_key_values=past_key_values, output_attentions=False, return_dict=True
+        )
+        tf.debugging.assert_equal(output[0], input_ids)
+        tf.debugging.assert_equal(output[1], past_key_values)
+        self.assertFalse(output[2])
+        self.assertFalse(output[3])
+        self.assertTrue(output[4])
+
+        # test case 5: Unexpected arguments should raise an exception.
+        with self.assertRaises(ValueError):
+            output = dummy_model.call(input_ids=input_ids, past_key_values=past_key_values, foo="bar")
+
+        # test case 6: the decorator is independent from `main_input_name` -- it treats the first argument of the
+        # decorated function as its main input.
+        output = dummy_model.foo(pixel_values=pixel_values)
+        tf.debugging.assert_equal(output[0], pixel_values)
+        self.assertFalse(output[1])
+        self.assertFalse(output[2])
+        self.assertFalse(output[3])
+
+    # Tests whether the stable softmax is stable on CPU, with and without XLA
+    def test_xla_stable_softmax(self):
+        large_penalty = -1e9
+        n_tokens = 10
+        batch_size = 8
+
+        def masked_softmax(x, boolean_mask):
+            numerical_mask = (1.0 - tf.cast(boolean_mask, dtype=tf.float32)) * large_penalty
+            masked_x = x + numerical_mask
+            return stable_softmax(masked_x)
+
+        xla_masked_softmax = tf.function(masked_softmax, jit_compile=True)
+        xla_stable_softmax = tf.function(stable_softmax, jit_compile=True)
+        x = tf.random.normal((batch_size, n_tokens))
+
+        # Same outcome regardless of the boolean mask here
+        masked_tokens = random.randint(0, n_tokens)
+        boolean_mask = tf.convert_to_tensor([[1] * (n_tokens - masked_tokens) + [0] * masked_tokens], dtype=tf.int32)
+
+        # We can randomly mask a random numerical input OUTSIDE XLA
+        numerical_mask = (1.0 - tf.cast(boolean_mask, dtype=tf.float32)) * large_penalty
+        masked_x = x + numerical_mask
+        xla_out = xla_stable_softmax(masked_x)
+        out = stable_softmax(masked_x)
+        assert tf.experimental.numpy.allclose(xla_out, out)
+
+        # The stable softmax has the same output as the original softmax
+        unstable_out = tf.nn.softmax(masked_x)
+        assert tf.experimental.numpy.allclose(unstable_out, out)
+
+        # We can randomly mask a random numerical input INSIDE XLA
+        xla_out = xla_masked_softmax(x, boolean_mask)
+        out = masked_softmax(x, boolean_mask)
+        assert tf.experimental.numpy.allclose(xla_out, out)
+
+    def test_checkpoint_sharding_from_hub(self):
+        model = TFBertModel.from_pretrained("ArthurZ/tiny-random-bert-sharded")
+        # the model above is the same as the model below, just a sharded version.
+        ref_model = TFBertModel.from_pretrained("hf-internal-testing/tiny-random-bert")
+        for p1, p2 in zip(model.weights, ref_model.weights):
+            assert np.allclose(p1.numpy(), p2.numpy())
+
+    def test_sharded_checkpoint_with_prefix(self):
+        model = TFBertModel.from_pretrained("hf-internal-testing/tiny-random-bert", load_weight_prefix="a/b")
+        sharded_model = TFBertModel.from_pretrained("ArthurZ/tiny-random-bert-sharded", load_weight_prefix="a/b")
+        for p1, p2 in zip(model.weights, sharded_model.weights):
+            self.assertTrue(np.allclose(p1.numpy(), p2.numpy()))
+            self.assertTrue(p1.name.startswith("a/b/"))
+            self.assertTrue(p2.name.startswith("a/b/"))
+
+    def test_sharded_checkpoint_transfer(self):
+        # If this doesn't throw an error then the test passes
+        TFBertForSequenceClassification.from_pretrained("ArthurZ/tiny-random-bert-sharded")
+
+    def test_shard_checkpoint(self):
+        # This is the model we will use, total size 340,000 bytes.
+        model = keras.Sequential(
+            [
+                keras.layers.Dense(200, use_bias=False),  # size 80,000
+                keras.layers.Dense(200, use_bias=False),  # size 160,000
+                keras.layers.Dense(100, use_bias=False),  # size 80,000
+                keras.layers.Dense(50, use_bias=False),  # size 20,000
+            ]
+        )
+        inputs = tf.zeros((1, 100), dtype=tf.float32)
+        model(inputs)
+        weights = model.weights
+        weights_dict = {w.name: w for w in weights}
+        with self.subTest("No shard when max size is bigger than model size"):
+            shards, index = tf_shard_checkpoint(weights)
+            self.assertIsNone(index)
+            self.assertDictEqual(shards, {TF2_WEIGHTS_NAME: weights})
+
+        with self.subTest("Test sharding, no weights bigger than max size"):
+            shards, index = tf_shard_checkpoint(weights, max_shard_size="300kB")
+            # Split is first two layers then last two.
+            self.assertDictEqual(
+                index,
+                {
+                    "metadata": {"total_size": 340000},
+                    "weight_map": {
+                        "dense/kernel:0": "tf_model-00001-of-00002.h5",
+                        "dense_1/kernel:0": "tf_model-00001-of-00002.h5",
+                        "dense_2/kernel:0": "tf_model-00002-of-00002.h5",
+                        "dense_3/kernel:0": "tf_model-00002-of-00002.h5",
+                    },
+                },
+            )
+
+            shard1 = [weights_dict["dense/kernel:0"], weights_dict["dense_1/kernel:0"]]
+            shard2 = [weights_dict["dense_2/kernel:0"], weights_dict["dense_3/kernel:0"]]
+            self.assertDictEqual(shards, {"tf_model-00001-of-00002.h5": shard1, "tf_model-00002-of-00002.h5": shard2})
+
+        with self.subTest("Test sharding with weights bigger than max size"):
+            shards, index = tf_shard_checkpoint(weights, max_shard_size="100kB")
+            # Split is first layer, second layer then last 2.
+            self.assertDictEqual(
+                index,
+                {
+                    "metadata": {"total_size": 340000},
+                    "weight_map": {
+                        "dense/kernel:0": "tf_model-00001-of-00003.h5",
+                        "dense_1/kernel:0": "tf_model-00002-of-00003.h5",
+                        "dense_2/kernel:0": "tf_model-00003-of-00003.h5",
+                        "dense_3/kernel:0": "tf_model-00003-of-00003.h5",
+                    },
+                },
+            )
+
+            shard1 = [weights_dict["dense/kernel:0"]]
+            shard2 = [weights_dict["dense_1/kernel:0"]]
+            shard3 = [weights_dict["dense_2/kernel:0"], weights_dict["dense_3/kernel:0"]]
+            self.assertDictEqual(
+                shards,
+                {
+                    "tf_model-00001-of-00003.h5": shard1,
+                    "tf_model-00002-of-00003.h5": shard2,
+                    "tf_model-00003-of-00003.h5": shard3,
+                },
+            )
+
+    @slow
+    def test_special_layer_name_sharding(self):
+        retriever = RagRetriever.from_pretrained("facebook/rag-token-nq", index_name="exact", use_dummy_dataset=True)
+        model = TFRagModel.from_pretrained("facebook/rag-token-nq", retriever=retriever)
+
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            for max_size in ["150kB", "150kiB", "200kB", "200kiB"]:
+                model.save_pretrained(tmp_dir, max_shard_size=max_size)
+                ref_model = TFRagModel.from_pretrained(tmp_dir, retriever=retriever)
+                for p1, p2 in zip(model.weights, ref_model.weights):
+                    assert np.allclose(p1.numpy(), p2.numpy())
+
+    @require_safetensors
+    def test_checkpoint_sharding_local(self):
+        model = TFBertModel.from_pretrained("hf-internal-testing/tiny-random-bert")
+
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            # We use the same folder for various sizes to make sure a new save erases the old checkpoint.
+            for max_size in ["150kB", "150kiB", "200kB", "200kiB"]:
+                model.save_pretrained(tmp_dir, max_shard_size=max_size)
+
+                # Get each shard file and its size
+                shard_to_size = {}
+                for shard in os.listdir(tmp_dir):
+                    if shard.endswith(".h5"):
+                        shard_file = os.path.join(tmp_dir, shard)
+                        shard_to_size[shard_file] = os.path.getsize(shard_file)
+
+                index_file = os.path.join(tmp_dir, TF2_WEIGHTS_INDEX_NAME)
+                # Check there is an index but no regular weight file
+                self.assertTrue(os.path.isfile(index_file))
+                self.assertFalse(os.path.isfile(os.path.join(tmp_dir, TF2_WEIGHTS_NAME)))
+
+                # Check a file is bigger than max_size only when it has a single weight
+                for shard_file, size in shard_to_size.items():
+                    if max_size.endswith("kiB"):
+                        max_size_int = int(max_size[:-3]) * 2**10
+                    else:
+                        max_size_int = int(max_size[:-2]) * 10**3
+                    # Note: pickle adds some junk so the weight of the file can end up being slightly bigger than
+                    # the size asked for (since we count parameters)
+                    if size >= max_size_int + 50000:
+                        with h5py.File(shard_file, "r") as state_file:
+                            self.assertEqual(len(state_file), 1)
+
+                # Check the index and the shard files found match
+                with open(index_file, encoding="utf-8") as f:
+                    index = json.loads(f.read())
+
+                all_shards = set(index["weight_map"].values())
+                shards_found = {f for f in os.listdir(tmp_dir) if f.endswith(".h5")}
+                self.assertSetEqual(all_shards, shards_found)
+
+                # Finally, check the model can be reloaded
+                new_model = TFBertModel.from_pretrained(tmp_dir)
+
+                model.build_in_name_scope()
+                new_model.build_in_name_scope()
+
+                for p1, p2 in zip(model.weights, new_model.weights):
+                    self.assertTrue(np.allclose(p1.numpy(), p2.numpy()))
+
+    def test_safetensors_checkpoint_sharding_local(self):
+        model = TFBertModel.from_pretrained("hf-internal-testing/tiny-random-bert")
+
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            # We use the same folder for various sizes to make sure a new save erases the old checkpoint.
+            for max_size in ["150kB", "150kiB", "200kB", "200kiB"]:
+                model.save_pretrained(tmp_dir, max_shard_size=max_size, safe_serialization=True)
+
+                # Get each shard file and its size
+                shard_to_size = {}
+                for shard in os.listdir(tmp_dir):
+                    if shard.endswith(".h5"):
+                        shard_file = os.path.join(tmp_dir, shard)
+                        shard_to_size[shard_file] = os.path.getsize(shard_file)
+
+                index_file = os.path.join(tmp_dir, SAFE_WEIGHTS_INDEX_NAME)
+                # Check there is an index but no regular weight file
+                self.assertTrue(os.path.isfile(index_file))
+                self.assertFalse(os.path.isfile(os.path.join(tmp_dir, SAFE_WEIGHTS_NAME)))
+                self.assertFalse(os.path.isfile(os.path.join(tmp_dir, TF2_WEIGHTS_NAME)))
+                self.assertFalse(os.path.isfile(os.path.join(tmp_dir, TF2_WEIGHTS_INDEX_NAME)))
+
+                # Check the index and the shard files found match
+                with open(index_file, encoding="utf-8") as f:
+                    index = json.loads(f.read())
+
+                all_shards = set(index["weight_map"].values())
+                shards_found = {f for f in os.listdir(tmp_dir) if f.endswith(".safetensors")}
+                self.assertSetEqual(all_shards, shards_found)
+
+                # Finally, check the model can be reloaded
+                new_model = TFBertModel.from_pretrained(tmp_dir)
+
+                model.build_in_name_scope()
+                new_model.build_in_name_scope()
+
+                for p1, p2 in zip(model.weights, new_model.weights):
+                    self.assertTrue(np.allclose(p1.numpy(), p2.numpy()))
+
+    @slow
+    def test_save_pretrained_signatures(self):
+        model = TFBertModel.from_pretrained("hf-internal-testing/tiny-random-bert")
+
+        # Short custom TF signature function.
+        # `input_signature` is specific to BERT.
+        @tf.function(
+            input_signature=[
+                [
+                    tf.TensorSpec([None, None], tf.int32, name="input_ids"),
+                    tf.TensorSpec([None, None], tf.int32, name="token_type_ids"),
+                    tf.TensorSpec([None, None], tf.int32, name="attention_mask"),
+                ]
+            ]
+        )
+        def serving_fn(input):
+            return model(input)
+
+        # Using default signature (default behavior) overrides 'serving_default'
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            model.save_pretrained(tmp_dir, saved_model=True, signatures=None)
+            model_loaded = keras.models.load_model(f"{tmp_dir}/saved_model/1")
+            self.assertTrue("serving_default" in list(model_loaded.signatures.keys()))
+
+        # Providing custom signature function
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            model.save_pretrained(tmp_dir, saved_model=True, signatures={"custom_signature": serving_fn})
+            model_loaded = keras.models.load_model(f"{tmp_dir}/saved_model/1")
+            self.assertTrue("custom_signature" in list(model_loaded.signatures.keys()))
+
+        # Providing multiple custom signature function
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            model.save_pretrained(
+                tmp_dir,
+                saved_model=True,
+                signatures={"custom_signature_1": serving_fn, "custom_signature_2": serving_fn},
+            )
+            model_loaded = keras.models.load_model(f"{tmp_dir}/saved_model/1")
+            self.assertTrue("custom_signature_1" in list(model_loaded.signatures.keys()))
+            self.assertTrue("custom_signature_2" in list(model_loaded.signatures.keys()))
+
+    @require_safetensors
+    def test_safetensors_save_and_load(self):
+        model = TFBertModel.from_pretrained("hf-internal-testing/tiny-random-bert")
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            model.save_pretrained(tmp_dir, safe_serialization=True)
+            # No tf_model.h5 file, only a model.safetensors
+            self.assertTrue(os.path.isfile(os.path.join(tmp_dir, SAFE_WEIGHTS_NAME)))
+            self.assertFalse(os.path.isfile(os.path.join(tmp_dir, SAFE_WEIGHTS_INDEX_NAME)))
+            self.assertFalse(os.path.isfile(os.path.join(tmp_dir, TF2_WEIGHTS_NAME)))
+            self.assertFalse(os.path.isfile(os.path.join(tmp_dir, TF2_WEIGHTS_INDEX_NAME)))
+
+            new_model = TFBertModel.from_pretrained(tmp_dir)
+
+            # Check models are equal
+            for p1, p2 in zip(model.weights, new_model.weights):
+                self.assertTrue(np.allclose(p1.numpy(), p2.numpy()))
+
+    @require_safetensors
+    def test_safetensors_sharded_save_and_load(self):
+        model = TFBertModel.from_pretrained("hf-internal-testing/tiny-random-bert")
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            model.save_pretrained(tmp_dir, safe_serialization=True, max_shard_size="150kB")
+            # No tf weights or index file, only a safetensors index
+            self.assertFalse(os.path.isfile(os.path.join(tmp_dir, SAFE_WEIGHTS_NAME)))
+            self.assertFalse(os.path.isfile(os.path.join(tmp_dir, TF2_WEIGHTS_NAME)))
+            self.assertTrue(os.path.isfile(os.path.join(tmp_dir, SAFE_WEIGHTS_INDEX_NAME)))
+            self.assertFalse(os.path.isfile(os.path.join(tmp_dir, TF2_WEIGHTS_INDEX_NAME)))
+
+            new_model = TFBertModel.from_pretrained(tmp_dir)
+
+            # Check models are equal
+            for p1, p2 in zip(model.weights, new_model.weights):
+                self.assertTrue(np.allclose(p1.numpy(), p2.numpy()))
+
+    @require_safetensors
+    def test_safetensors_load_from_hub(self):
+        tf_model = TFBertModel.from_pretrained("hf-internal-testing/tiny-random-bert")
+
+        # Can load from the TF-formatted checkpoint
+        safetensors_model = TFBertModel.from_pretrained("hf-internal-testing/tiny-random-bert-safetensors-tf")
+
+        # Check models are equal
+        for p1, p2 in zip(safetensors_model.weights, tf_model.weights):
+            self.assertTrue(np.allclose(p1.numpy(), p2.numpy()))
+
+        # Can load from the PyTorch-formatted checkpoint
+        safetensors_model = TFBertModel.from_pretrained("hf-internal-testing/tiny-random-bert-safetensors")
+
+        # Check models are equal
+        for p1, p2 in zip(safetensors_model.weights, tf_model.weights):
+            self.assertTrue(np.allclose(p1.numpy(), p2.numpy()))
+
+    @require_safetensors
+    def test_safetensors_tf_from_tf(self):
+        model = TFBertModel.from_pretrained("hf-internal-testing/tiny-bert-tf-only")
+
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            model.save_pretrained(tmp_dir, safe_serialization=True)
+            new_model = TFBertModel.from_pretrained(tmp_dir)
+
+        for p1, p2 in zip(model.weights, new_model.weights):
+            self.assertTrue(np.allclose(p1.numpy(), p2.numpy()))
+
+    @require_safetensors
+    def test_safetensors_tf_from_sharded_h5_with_sharded_safetensors_local(self):
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            path = snapshot_download("hf-internal-testing/tiny-bert-tf-safetensors-h5-sharded", cache_dir=tmp_dir)
+
+            # This should not raise even if there are two types of sharded weights
+            TFBertModel.from_pretrained(path)
+
+    @require_safetensors
+    def test_safetensors_tf_from_sharded_h5_with_sharded_safetensors_hub(self):
+        # Confirm that we can correctly load the safetensors weights from a sharded hub repo even when TF weights present
+        TFBertModel.from_pretrained("hf-internal-testing/tiny-bert-tf-safetensors-h5-sharded", use_safetensors=True)
+        # Confirm that we can access the TF weights too
+        TFBertModel.from_pretrained("hf-internal-testing/tiny-bert-tf-safetensors-h5-sharded", use_safetensors=False)
+
+    @require_safetensors
+    def test_safetensors_load_from_local(self):
+        """
+        This test checks that we can load safetensors from a checkpoint that only has those on the Hub
+        """
+        with tempfile.TemporaryDirectory() as tmp:
+            location = snapshot_download("hf-internal-testing/tiny-bert-tf-only", cache_dir=tmp)
+            tf_model = TFBertModel.from_pretrained(location)
+
+        with tempfile.TemporaryDirectory() as tmp:
+            location = snapshot_download("hf-internal-testing/tiny-bert-tf-safetensors-only", cache_dir=tmp)
+            safetensors_model = TFBertModel.from_pretrained(location)
+
+        for p1, p2 in zip(tf_model.weights, safetensors_model.weights):
+            self.assertTrue(np.allclose(p1.numpy(), p2.numpy()))
+
+    @require_safetensors
+    def test_safetensors_load_from_hub_from_safetensors_pt(self):
+        """
+        This test checks that we can load safetensors from a checkpoint that only has those on the Hub.
+        saved in the "pt" format.
+        """
+        tf_model = TFBertModel.from_pretrained("hf-internal-testing/tiny-bert-h5")
+
+        # Can load from the PyTorch-formatted checkpoint
+        safetensors_model = TFBertModel.from_pretrained("hf-internal-testing/tiny-bert-pt-safetensors")
+        for p1, p2 in zip(tf_model.weights, safetensors_model.weights):
+            self.assertTrue(np.allclose(p1.numpy(), p2.numpy()))
+
+    @require_safetensors
+    def test_safetensors_load_from_local_from_safetensors_pt(self):
+        """
+        This test checks that we can load safetensors from a local checkpoint that only has those
+        saved in the "pt" format.
+        """
+        with tempfile.TemporaryDirectory() as tmp:
+            location = snapshot_download("hf-internal-testing/tiny-bert-h5", cache_dir=tmp)
+            tf_model = TFBertModel.from_pretrained(location)
+
+        # Can load from the PyTorch-formatted checkpoint
+        with tempfile.TemporaryDirectory() as tmp:
+            location = snapshot_download("hf-internal-testing/tiny-bert-pt-safetensors", cache_dir=tmp)
+            safetensors_model = TFBertModel.from_pretrained(location)
+
+        for p1, p2 in zip(tf_model.weights, safetensors_model.weights):
+            self.assertTrue(np.allclose(p1.numpy(), p2.numpy()))
+
+    @require_safetensors
+    def test_safetensors_load_from_hub_h5_before_safetensors(self):
+        """
+        This test checks that we'll first download h5 weights before safetensors
+        The safetensors file on that repo is a pt safetensors and therefore cannot be loaded without PyTorch
+        """
+        TFBertModel.from_pretrained("hf-internal-testing/tiny-bert-pt-safetensors-msgpack")
+
+    @require_safetensors
+    def test_safetensors_load_from_local_h5_before_safetensors(self):
+        """
+        This test checks that we'll first download h5 weights before safetensors
+        The safetensors file on that repo is a pt safetensors and therefore cannot be loaded without PyTorch
+        """
+        with tempfile.TemporaryDirectory() as tmp:
+            location = snapshot_download("hf-internal-testing/tiny-bert-pt-safetensors-msgpack", cache_dir=tmp)
+            TFBertModel.from_pretrained(location)
+
+
+@require_tf
+@is_staging_test
+class TFModelPushToHubTester(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls._token = TOKEN
+        HfFolder.save_token(TOKEN)
+
+    def test_push_to_hub(self):
+        with TemporaryHubRepo(token=self._token) as tmp_repo:
+            config = BertConfig(
+                vocab_size=99, hidden_size=32, num_hidden_layers=5, num_attention_heads=4, intermediate_size=37
+            )
+            model = TFBertModel(config)
+            # Make sure model is properly initialized
+            model.build_in_name_scope()
+
+            logging.set_verbosity_info()
+            logger = logging.get_logger("transformers.utils.hub")
+            with CaptureLogger(logger) as cl:
+                model.push_to_hub(tmp_repo.repo_id, token=self._token)
+            logging.set_verbosity_warning()
+            # Check the model card was created and uploaded.
+            self.assertIn("Uploading the following files to __DUMMY_TRANSFORMERS_USER__/test-model-tf", cl.out)
+
+            new_model = TFBertModel.from_pretrained(tmp_repo.repo_id)
+            models_equal = True
+            for p1, p2 in zip(model.weights, new_model.weights):
+                if not tf.math.reduce_all(p1 == p2):
+                    models_equal = False
+                    break
+            self.assertTrue(models_equal)
+
+    def test_push_to_hub_via_save_pretrained(self):
+        with TemporaryHubRepo(token=self._token) as tmp_repo:
+            config = BertConfig(
+                vocab_size=99, hidden_size=32, num_hidden_layers=5, num_attention_heads=4, intermediate_size=37
+            )
+            model = TFBertModel(config)
+            # Make sure model is properly initialized
+            model.build_in_name_scope()
+
+            # Push to hub via save_pretrained
+            with tempfile.TemporaryDirectory() as tmp_dir:
+                model.save_pretrained(tmp_dir, repo_id=tmp_repo.repo_id, push_to_hub=True, token=self._token)
+
+            new_model = TFBertModel.from_pretrained(tmp_repo.repo_id)
+            models_equal = True
+            for p1, p2 in zip(model.weights, new_model.weights):
+                if not tf.math.reduce_all(p1 == p2):
+                    models_equal = False
+                    break
+            self.assertTrue(models_equal)
+
+    def test_push_to_hub_in_organization(self):
+        with TemporaryHubRepo(namespace="valid_org", token=self._token) as tmp_repo:
+            config = BertConfig(
+                vocab_size=99, hidden_size=32, num_hidden_layers=5, num_attention_heads=4, intermediate_size=37
+            )
+            model = TFBertModel(config)
+            # Make sure model is properly initialized
+            model.build_in_name_scope()
+
+            model.push_to_hub(tmp_repo.repo_id, token=self._token)
+
+            new_model = TFBertModel.from_pretrained(tmp_repo.repo_id)
+            models_equal = True
+            for p1, p2 in zip(model.weights, new_model.weights):
+                if not tf.math.reduce_all(p1 == p2):
+                    models_equal = False
+                    break
+            self.assertTrue(models_equal)
+
+    def test_push_to_hub_in_organization_via_save_pretrained(self):
+        with TemporaryHubRepo(namespace="valid_org", token=self._token) as tmp_repo:
+            config = BertConfig(
+                vocab_size=99, hidden_size=32, num_hidden_layers=5, num_attention_heads=4, intermediate_size=37
+            )
+            model = TFBertModel(config)
+            # Make sure model is properly initialized
+            model.build_in_name_scope()
+
+            # Push to hub via save_pretrained
+            with tempfile.TemporaryDirectory() as tmp_dir:
+                model.save_pretrained(tmp_dir, push_to_hub=True, token=self._token, repo_id=tmp_repo.repo_id)
+
+            new_model = TFBertModel.from_pretrained(tmp_repo.repo_id)
+            models_equal = True
+            for p1, p2 in zip(model.weights, new_model.weights):
+                if not tf.math.reduce_all(p1 == p2):
+                    models_equal = False
+                    break
+            self.assertTrue(models_equal)

docs/transformers/tests/utils/test_offline.py

@@ -0,0 +1,220 @@
+# Copyright 2020 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import subprocess
+import sys
+import unittest
+
+from transformers import BertConfig, BertModel, BertTokenizer, pipeline
+from transformers.testing_utils import TestCasePlus, require_torch
+
+
+class OfflineTests(TestCasePlus):
+    @require_torch
+    @unittest.skip("This test is failing on main")  # TODO matt/ydshieh, this test needs to be fixed
+    def test_offline_mode(self):
+        # this test is a bit tricky since TRANSFORMERS_OFFLINE can only be changed before
+        # `transformers` is loaded, and it's too late for inside pytest - so we are changing it
+        # while running an external program
+
+        # python one-liner segments
+
+        # this must be loaded before socket.socket is monkey-patched
+        load = """
+from transformers import BertConfig, BertModel, BertTokenizer, pipeline
+        """
+
+        run = """
+mname = "hf-internal-testing/tiny-random-bert"
+BertConfig.from_pretrained(mname)
+BertModel.from_pretrained(mname)
+BertTokenizer.from_pretrained(mname)
+pipe = pipeline(task="fill-mask", model=mname)
+print("success")
+        """
+
+        mock = """
+import socket
+def offline_socket(*args, **kwargs): raise RuntimeError("Offline mode is enabled, we shouldn't access internet")
+socket.socket = offline_socket
+        """
+
+        # Force fetching the files so that we can use the cache
+        mname = "hf-internal-testing/tiny-random-bert"
+        BertConfig.from_pretrained(mname)
+        BertModel.from_pretrained(mname)
+        BertTokenizer.from_pretrained(mname)
+        pipeline(task="fill-mask", model=mname)
+
+        # baseline - just load from_pretrained with normal network
+        # should succeed as TRANSFORMERS_OFFLINE=1 tells it to use local files
+        stdout, _ = self._execute_with_env(load, run, mock, TRANSFORMERS_OFFLINE="1")
+        self.assertIn("success", stdout)
+
+    @require_torch
+    def test_offline_mode_no_internet(self):
+        # python one-liner segments
+        # this must be loaded before socket.socket is monkey-patched
+        load = """
+from transformers import BertConfig, BertModel, BertTokenizer, pipeline
+        """
+
+        run = """
+mname = "hf-internal-testing/tiny-random-bert"
+BertConfig.from_pretrained(mname)
+BertModel.from_pretrained(mname)
+BertTokenizer.from_pretrained(mname)
+pipe = pipeline(task="fill-mask", model=mname)
+print("success")
+        """
+
+        mock = """
+import socket
+def offline_socket(*args, **kwargs): raise socket.error("Faking flaky internet")
+socket.socket = offline_socket
+        """
+
+        # Force fetching the files so that we can use the cache
+        mname = "hf-internal-testing/tiny-random-bert"
+        BertConfig.from_pretrained(mname)
+        BertModel.from_pretrained(mname)
+        BertTokenizer.from_pretrained(mname)
+        pipeline(task="fill-mask", model=mname)
+
+        # baseline - just load from_pretrained with normal network
+        # should succeed
+        stdout, _ = self._execute_with_env(load, run, mock)
+        self.assertIn("success", stdout)
+
+    @require_torch
+    def test_offline_mode_sharded_checkpoint(self):
+        # this test is a bit tricky since TRANSFORMERS_OFFLINE can only be changed before
+        # `transformers` is loaded, and it's too late for inside pytest - so we are changing it
+        # while running an external program
+
+        # python one-liner segments
+
+        # this must be loaded before socket.socket is monkey-patched
+        load = """
+from transformers import BertConfig, BertModel, BertTokenizer
+        """
+
+        run = """
+mname = "hf-internal-testing/tiny-random-bert-sharded"
+BertConfig.from_pretrained(mname)
+BertModel.from_pretrained(mname)
+print("success")
+        """
+
+        mock = """
+import socket
+def offline_socket(*args, **kwargs): raise ValueError("Offline mode is enabled")
+socket.socket = offline_socket
+        """
+
+        # baseline - just load from_pretrained with normal network
+        # should succeed
+        stdout, _ = self._execute_with_env(load, run)
+        self.assertIn("success", stdout)
+
+        # next emulate no network
+        # Doesn't fail anymore since the model is in the cache due to other tests, so commenting this.
+        # self._execute_with_env(load, mock, run, should_fail=True, TRANSFORMERS_OFFLINE="0")
+
+        # should succeed as TRANSFORMERS_OFFLINE=1 tells it to use local files
+        stdout, _ = self._execute_with_env(load, mock, run, TRANSFORMERS_OFFLINE="1")
+        self.assertIn("success", stdout)
+
+    @require_torch
+    def test_offline_mode_pipeline_exception(self):
+        load = """
+from transformers import pipeline
+        """
+        run = """
+mname = "hf-internal-testing/tiny-random-bert"
+pipe = pipeline(model=mname)
+        """
+
+        mock = """
+import socket
+def offline_socket(*args, **kwargs): raise socket.error("Offline mode is enabled")
+socket.socket = offline_socket
+        """
+
+        _, stderr = self._execute_with_env(load, mock, run, should_fail=True, TRANSFORMERS_OFFLINE="1")
+        self.assertIn(
+            "You cannot infer task automatically within `pipeline` when using offline mode",
+            stderr.replace("\n", ""),
+        )
+
+    @require_torch
+    def test_offline_model_dynamic_model(self):
+        load = """
+from transformers import AutoModel
+        """
+        run = """
+mname = "hf-internal-testing/test_dynamic_model"
+AutoModel.from_pretrained(mname, trust_remote_code=True)
+print("success")
+        """
+
+        # baseline - just load from_pretrained with normal network
+        # should succeed
+        stdout, _ = self._execute_with_env(load, run)
+        self.assertIn("success", stdout)
+
+        # should succeed as TRANSFORMERS_OFFLINE=1 tells it to use local files
+        stdout, _ = self._execute_with_env(load, run, TRANSFORMERS_OFFLINE="1")
+        self.assertIn("success", stdout)
+
+    def test_is_offline_mode(self):
+        """
+        Test `_is_offline_mode` helper (should respect both HF_HUB_OFFLINE and legacy TRANSFORMERS_OFFLINE env vars)
+        """
+        load = "from transformers.utils import is_offline_mode"
+        run = "print(is_offline_mode())"
+
+        stdout, _ = self._execute_with_env(load, run)
+        self.assertIn("False", stdout)
+
+        stdout, _ = self._execute_with_env(load, run, TRANSFORMERS_OFFLINE="1")
+        self.assertIn("True", stdout)
+
+        stdout, _ = self._execute_with_env(load, run, HF_HUB_OFFLINE="1")
+        self.assertIn("True", stdout)
+
+    def _execute_with_env(self, *commands: tuple[str, ...], should_fail: bool = False, **env) -> tuple[str, str]:
+        """Execute Python code with a given environment and return the stdout/stderr as strings.
+
+        If `should_fail=True`, the command is expected to fail. Otherwise, it should succeed.
+        Environment variables can be passed as keyword arguments.
+        """
+        # Build command
+        cmd = [sys.executable, "-c", "\n".join(commands)]
+
+        # Configure env
+        new_env = self.get_env()
+        new_env.update(env)
+
+        # Run command
+        result = subprocess.run(cmd, env=new_env, check=False, capture_output=True)
+
+        # Check execution
+        if should_fail:
+            self.assertNotEqual(result.returncode, 0, result.stderr)
+        else:
+            self.assertEqual(result.returncode, 0, result.stderr)
+
+        # Return output
+        return result.stdout.decode(), result.stderr.decode()

docs/transformers/tests/utils/test_processing_utils.py

@@ -0,0 +1,175 @@
+# Copyright 2024 HuggingFace Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+
+from transformers import is_torch_available, is_vision_available
+from transformers.processing_utils import _validate_images_text_input_order
+from transformers.testing_utils import require_torch, require_vision
+
+
+if is_vision_available():
+    import PIL
+
+if is_torch_available():
+    import torch
+
+
+@require_vision
+class ProcessingUtilTester(unittest.TestCase):
+    def test_validate_images_text_input_order(self):
+        # text string and PIL images inputs
+        images = PIL.Image.new("RGB", (224, 224))
+        text = "text"
+        # test correct text and images order
+        valid_images, valid_text = _validate_images_text_input_order(images=images, text=text)
+        self.assertEqual(valid_images, images)
+        self.assertEqual(valid_text, text)
+        # test incorrect text and images order
+        valid_images, valid_text = _validate_images_text_input_order(images=text, text=images)
+        self.assertEqual(valid_images, images)
+        self.assertEqual(valid_text, text)
+
+        # text list of string and numpy images inputs
+        images = np.random.rand(224, 224, 3)
+        text = ["text1", "text2"]
+        # test correct text and images order
+        valid_images, valid_text = _validate_images_text_input_order(images=images, text=text)
+        self.assertTrue(np.array_equal(valid_images, images))
+        self.assertEqual(valid_text, text)
+        # test incorrect text and images order
+        valid_images, valid_text = _validate_images_text_input_order(images=text, text=images)
+        self.assertTrue(np.array_equal(valid_images, images))
+        self.assertEqual(valid_text, text)
+
+        # text nested list of string and list of pil images inputs
+        images = [PIL.Image.new("RGB", (224, 224)), PIL.Image.new("RGB", (224, 224))]
+        text = [["text1", "text2, text3"], ["text3", "text4"]]
+        # test correct text and images order
+        valid_images, valid_text = _validate_images_text_input_order(images=images, text=text)
+        self.assertEqual(valid_images, images)
+        self.assertEqual(valid_text, text)
+        # test incorrect text and images order
+        valid_images, valid_text = _validate_images_text_input_order(images=text, text=images)
+        self.assertEqual(valid_images, images)
+        self.assertEqual(valid_text, text)
+
+        # list of strings and list of numpy images inputs
+        images = [np.random.rand(224, 224, 3), np.random.rand(224, 224, 3)]
+        text = ["text1", "text2"]
+        # test correct text and images order
+        valid_images, valid_text = _validate_images_text_input_order(images=images, text=text)
+        self.assertTrue(np.array_equal(valid_images[0], images[0]))
+        self.assertEqual(valid_text, text)
+        # test incorrect text and images order
+        valid_images, valid_text = _validate_images_text_input_order(images=text, text=images)
+        self.assertTrue(np.array_equal(valid_images[0], images[0]))
+        self.assertEqual(valid_text, text)
+
+        # list of strings and list of url images inputs
+        images = ["https://url1", "https://url2"]
+        text = ["text1", "text2"]
+        # test correct text and images order
+        valid_images, valid_text = _validate_images_text_input_order(images=images, text=text)
+        self.assertEqual(valid_images, images)
+        self.assertEqual(valid_text, text)
+        # test incorrect text and images order
+        valid_images, valid_text = _validate_images_text_input_order(images=text, text=images)
+        self.assertEqual(valid_images, images)
+        self.assertEqual(valid_text, text)
+
+        # list of strings and nested list of numpy images inputs
+        images = [[np.random.rand(224, 224, 3), np.random.rand(224, 224, 3)], [np.random.rand(224, 224, 3)]]
+        text = ["text1", "text2"]
+        # test correct text and images order
+        valid_images, valid_text = _validate_images_text_input_order(images=images, text=text)
+        self.assertTrue(np.array_equal(valid_images[0][0], images[0][0]))
+        self.assertEqual(valid_text, text)
+        # test incorrect text and images order
+        valid_images, valid_text = _validate_images_text_input_order(images=text, text=images)
+        self.assertTrue(np.array_equal(valid_images[0][0], images[0][0]))
+        self.assertEqual(valid_text, text)
+
+        # nested list of strings and nested list of PIL images inputs
+        images = [
+            [PIL.Image.new("RGB", (224, 224)), PIL.Image.new("RGB", (224, 224))],
+            [PIL.Image.new("RGB", (224, 224))],
+        ]
+        text = [["text1", "text2, text3"], ["text3", "text4"]]
+        # test correct text and images order
+        valid_images, valid_text = _validate_images_text_input_order(images=images, text=text)
+        self.assertEqual(valid_images, images)
+        self.assertEqual(valid_text, text)
+        # test incorrect text and images order
+        valid_images, valid_text = _validate_images_text_input_order(images=text, text=images)
+        self.assertEqual(valid_images, images)
+        self.assertEqual(valid_text, text)
+
+        # None images
+        images = None
+        text = "text"
+        # test correct text and images order
+        valid_images, valid_text = _validate_images_text_input_order(images=images, text=text)
+        self.assertEqual(images, None)
+        self.assertEqual(text, text)
+        # test incorrect text and images order
+        valid_images, valid_text = _validate_images_text_input_order(images=text, text=images)
+        self.assertEqual(images, None)
+        self.assertEqual(text, text)
+
+        # None text
+        images = PIL.Image.new("RGB", (224, 224))
+        text = None
+        # test correct text and images order
+        valid_images, valid_text = _validate_images_text_input_order(images=images, text=text)
+        self.assertEqual(images, images)
+        self.assertEqual(text, None)
+        # test incorrect text and images order
+        valid_images, valid_text = _validate_images_text_input_order(images=text, text=images)
+        self.assertEqual(images, images)
+        self.assertEqual(text, None)
+
+        # incorrect inputs
+        images = "text"
+        text = "text"
+        with self.assertRaises(ValueError):
+            _validate_images_text_input_order(images=images, text=text)
+
+    @require_torch
+    def test_validate_images_text_input_order_torch(self):
+        # text string and torch images inputs
+        images = torch.rand(224, 224, 3)
+        text = "text"
+        # test correct text and images order
+        valid_images, valid_text = _validate_images_text_input_order(images=images, text=text)
+        self.assertTrue(torch.equal(valid_images, images))
+        self.assertEqual(valid_text, text)
+        # test incorrect text and images order
+        valid_images, valid_text = _validate_images_text_input_order(images=text, text=images)
+        self.assertTrue(torch.equal(valid_images, images))
+        self.assertEqual(valid_text, text)
+
+        # text list of string and list of torch images inputs
+        images = [torch.rand(224, 224, 3), torch.rand(224, 224, 3)]
+        text = ["text1", "text2"]
+        # test correct text and images order
+        valid_images, valid_text = _validate_images_text_input_order(images=images, text=text)
+        self.assertTrue(torch.equal(valid_images[0], images[0]))
+        self.assertEqual(valid_text, text)
+        # test incorrect text and images order
+        valid_images, valid_text = _validate_images_text_input_order(images=text, text=images)
+        self.assertTrue(torch.equal(valid_images[0], images[0]))
+        self.assertEqual(valid_text, text)

docs/transformers/tests/utils/test_skip_decorators.py

@@ -0,0 +1,119 @@
+# Copyright 2019-present, the HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+#
+#
+# this test validates that we can stack skip decorators in groups and whether
+# they work correctly with other decorators
+#
+# since the decorators have already built their decision params (like checking
+# env[], we can't mock the env and test each of the combinations), so ideally
+# the following 4 should be run. But since we have different CI jobs running
+# different configs, all combinations should get covered
+#
+# RUN_SLOW=1 pytest -rA tests/test_skip_decorators.py
+# RUN_SLOW=1 CUDA_VISIBLE_DEVICES="" pytest -rA tests/test_skip_decorators.py
+# RUN_SLOW=0 pytest -rA tests/test_skip_decorators.py
+# RUN_SLOW=0 CUDA_VISIBLE_DEVICES="" pytest -rA tests/test_skip_decorators.py
+
+import os
+import unittest
+
+import pytest
+from parameterized import parameterized
+
+from transformers.testing_utils import require_torch, require_torch_gpu, slow, torch_device
+
+
+# skipping in unittest tests
+
+params = [(1,)]
+
+
+# test that we can stack our skip decorators with 3rd party decorators
+def check_slow():
+    run_slow = bool(os.getenv("RUN_SLOW", 0))
+    if run_slow:
+        assert True
+    else:
+        assert False, "should have been skipped"
+
+
+# test that we can stack our skip decorators
+def check_slow_torch_cuda():
+    run_slow = bool(os.getenv("RUN_SLOW", 0))
+    if run_slow and torch_device == "cuda":
+        assert True
+    else:
+        assert False, "should have been skipped"
+
+
+@require_torch
+class SkipTester(unittest.TestCase):
+    @slow
+    @require_torch_gpu
+    def test_2_skips_slow_first(self):
+        check_slow_torch_cuda()
+
+    @require_torch_gpu
+    @slow
+    def test_2_skips_slow_last(self):
+        check_slow_torch_cuda()
+
+    # The combination of any skip decorator, followed by parameterized fails to skip the tests
+    # 1. @slow manages to correctly skip `test_param_slow_first`
+    # 2. but then `parameterized` creates new tests, with a unique name for each parameter groups.
+    #    It has no idea that they are to be skipped and so they all run, ignoring @slow
+    # Therefore skip decorators must come after `parameterized`
+    #
+    # @slow
+    # @parameterized.expand(params)
+    # def test_param_slow_first(self, param=None):
+    #     check_slow()
+
+    # This works as expected:
+    # 1. `parameterized` creates new tests with unique names
+    # 2. each of them gets an opportunity to be skipped
+    @parameterized.expand(params)
+    @slow
+    def test_param_slow_last(self, param=None):
+        check_slow()
+
+
+# skipping in non-unittest tests
+# no problem at all here
+
+
+@slow
+@require_torch_gpu
+def test_pytest_2_skips_slow_first():
+    check_slow_torch_cuda()
+
+
+@require_torch_gpu
+@slow
+def test_pytest_2_skips_slow_last():
+    check_slow_torch_cuda()
+
+
+@slow
+@pytest.mark.parametrize("param", [1])
+def test_pytest_param_slow_first(param):
+    check_slow()
+
+
+@pytest.mark.parametrize("param", [1])
+@slow
+def test_pytest_param_slow_last(param):
+    check_slow()

docs/transformers/tests/utils/test_tokenization_utils.py

@@ -0,0 +1,303 @@
+# Copyright 2019 HuggingFace Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import sys
+import tempfile
+import unittest
+import unittest.mock as mock
+from pathlib import Path
+
+from huggingface_hub import HfFolder
+from huggingface_hub.file_download import http_get
+from requests.exceptions import HTTPError
+
+from transformers import (
+    AlbertTokenizer,
+    AutoTokenizer,
+    BertTokenizer,
+    BertTokenizerFast,
+    GPT2TokenizerFast,
+    is_tokenizers_available,
+)
+from transformers.testing_utils import TOKEN, TemporaryHubRepo, is_staging_test, require_tokenizers
+from transformers.tokenization_utils import ExtensionsTrie, Trie
+
+
+sys.path.append(str(Path(__file__).parent.parent.parent / "utils"))
+
+from test_module.custom_tokenization import CustomTokenizer  # noqa E402
+
+
+if is_tokenizers_available():
+    from test_module.custom_tokenization_fast import CustomTokenizerFast
+
+
+class TokenizerUtilTester(unittest.TestCase):
+    def test_cached_files_are_used_when_internet_is_down(self):
+        # A mock response for an HTTP head request to emulate server down
+        response_mock = mock.Mock()
+        response_mock.status_code = 500
+        response_mock.headers = {}
+        response_mock.raise_for_status.side_effect = HTTPError
+        response_mock.json.return_value = {}
+
+        # Download this model to make sure it's in the cache.
+        _ = BertTokenizer.from_pretrained("hf-internal-testing/tiny-random-bert")
+
+        # Under the mock environment we get a 500 error when trying to reach the tokenizer.
+        with mock.patch("requests.Session.request", return_value=response_mock) as mock_head:
+            _ = BertTokenizer.from_pretrained("hf-internal-testing/tiny-random-bert")
+            # This check we did call the fake head request
+            mock_head.assert_called()
+
+    @require_tokenizers
+    def test_cached_files_are_used_when_internet_is_down_missing_files(self):
+        # A mock response for an HTTP head request to emulate server down
+        response_mock = mock.Mock()
+        response_mock.status_code = 500
+        response_mock.headers = {}
+        response_mock.raise_for_status.side_effect = HTTPError
+        response_mock.json.return_value = {}
+
+        # Download this model to make sure it's in the cache.
+        _ = GPT2TokenizerFast.from_pretrained("openai-community/gpt2")
+
+        # Under the mock environment we get a 500 error when trying to reach the tokenizer.
+        with mock.patch("requests.Session.request", return_value=response_mock) as mock_head:
+            _ = GPT2TokenizerFast.from_pretrained("openai-community/gpt2")
+            # This check we did call the fake head request
+            mock_head.assert_called()
+
+    def test_legacy_load_from_one_file(self):
+        # This test is for deprecated behavior and can be removed in v5
+        try:
+            tmp_file = tempfile.NamedTemporaryFile(delete=False).name
+            with open(tmp_file, "wb") as f:
+                http_get("https://huggingface.co/albert/albert-base-v1/resolve/main/spiece.model", f)
+
+            _ = AlbertTokenizer.from_pretrained(tmp_file)
+        finally:
+            os.remove(tmp_file)
+
+        # Supporting this legacy load introduced a weird bug where the tokenizer would load local files if they are in
+        # the current folder and have the right name.
+        if os.path.isfile("tokenizer.json"):
+            # We skip the test if the user has a `tokenizer.json` in this folder to avoid deleting it.
+            self.skipTest(reason="Skipping test as there is a `tokenizer.json` file in the current folder.")
+        try:
+            with open("tokenizer.json", "wb") as f:
+                http_get("https://huggingface.co/hf-internal-testing/tiny-random-bert/blob/main/tokenizer.json", f)
+            tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-gpt2")
+            # The tiny random BERT has a vocab size of 1024, tiny openai-community/gpt2 as a vocab size of 1000
+            self.assertEqual(tokenizer.vocab_size, 1000)
+            # Tokenizer should depend on the remote checkpoint, not the local tokenizer.json file.
+
+        finally:
+            os.remove("tokenizer.json")
+
+
+@is_staging_test
+class TokenizerPushToHubTester(unittest.TestCase):
+    vocab_tokens = ["[UNK]", "[CLS]", "[SEP]", "[PAD]", "[MASK]", "bla", "blou"]
+
+    @classmethod
+    def setUpClass(cls):
+        cls._token = TOKEN
+        HfFolder.save_token(TOKEN)
+
+    def test_push_to_hub(self):
+        with TemporaryHubRepo(token=self._token) as tmp_repo:
+            with tempfile.TemporaryDirectory() as tmp_dir:
+                vocab_file = os.path.join(tmp_dir, "vocab.txt")
+                with open(vocab_file, "w", encoding="utf-8") as vocab_writer:
+                    vocab_writer.write("".join([x + "\n" for x in self.vocab_tokens]))
+                tokenizer = BertTokenizer(vocab_file)
+
+            tokenizer.push_to_hub(tmp_repo.repo_id, token=self._token)
+            new_tokenizer = BertTokenizer.from_pretrained(tmp_repo.repo_id)
+            self.assertDictEqual(new_tokenizer.vocab, tokenizer.vocab)
+
+    def test_push_to_hub_via_save_pretrained(self):
+        with TemporaryHubRepo(token=self._token) as tmp_repo:
+            with tempfile.TemporaryDirectory() as tmp_dir:
+                vocab_file = os.path.join(tmp_dir, "vocab.txt")
+                with open(vocab_file, "w", encoding="utf-8") as vocab_writer:
+                    vocab_writer.write("".join([x + "\n" for x in self.vocab_tokens]))
+                tokenizer = BertTokenizer(vocab_file)
+
+                # Push to hub via save_pretrained
+                tokenizer.save_pretrained(tmp_dir, repo_id=tmp_repo.repo_id, push_to_hub=True, token=self._token)
+
+            new_tokenizer = BertTokenizer.from_pretrained(tmp_repo.repo_id)
+            self.assertDictEqual(new_tokenizer.vocab, tokenizer.vocab)
+
+    def test_push_to_hub_in_organization(self):
+        with TemporaryHubRepo(namespace="valid_org", token=self._token) as tmp_repo:
+            with tempfile.TemporaryDirectory() as tmp_dir:
+                vocab_file = os.path.join(tmp_dir, "vocab.txt")
+                with open(vocab_file, "w", encoding="utf-8") as vocab_writer:
+                    vocab_writer.write("".join([x + "\n" for x in self.vocab_tokens]))
+                tokenizer = BertTokenizer(vocab_file)
+
+            tokenizer.push_to_hub(tmp_repo.repo_id, token=self._token)
+            new_tokenizer = BertTokenizer.from_pretrained(tmp_repo.repo_id)
+            self.assertDictEqual(new_tokenizer.vocab, tokenizer.vocab)
+
+    def test_push_to_hub_in_organization_via_save_pretrained(self):
+        with TemporaryHubRepo(namespace="valid_org", token=self._token) as tmp_repo:
+            with tempfile.TemporaryDirectory() as tmp_dir:
+                vocab_file = os.path.join(tmp_dir, "vocab.txt")
+                with open(vocab_file, "w", encoding="utf-8") as vocab_writer:
+                    vocab_writer.write("".join([x + "\n" for x in self.vocab_tokens]))
+                tokenizer = BertTokenizer(vocab_file)
+
+                # Push to hub via save_pretrained
+                tokenizer.save_pretrained(tmp_dir, repo_id=tmp_repo.repo_id, push_to_hub=True, token=self._token)
+
+            new_tokenizer = BertTokenizer.from_pretrained(tmp_repo.repo_id)
+            self.assertDictEqual(new_tokenizer.vocab, tokenizer.vocab)
+
+    @require_tokenizers
+    def test_push_to_hub_dynamic_tokenizer(self):
+        with TemporaryHubRepo(token=self._token) as tmp_repo:
+            CustomTokenizer.register_for_auto_class()
+            with tempfile.TemporaryDirectory() as tmp_dir:
+                vocab_file = os.path.join(tmp_dir, "vocab.txt")
+                with open(vocab_file, "w", encoding="utf-8") as vocab_writer:
+                    vocab_writer.write("".join([x + "\n" for x in self.vocab_tokens]))
+                tokenizer = CustomTokenizer(vocab_file)
+
+            # No fast custom tokenizer
+            tokenizer.push_to_hub(tmp_repo.repo_id, token=self._token)
+
+            tokenizer = AutoTokenizer.from_pretrained(tmp_repo.repo_id, trust_remote_code=True)
+            # Can't make an isinstance check because the new_model.config is from the CustomTokenizer class of a dynamic module
+            self.assertEqual(tokenizer.__class__.__name__, "CustomTokenizer")
+
+    @require_tokenizers
+    def test_push_to_hub_dynamic_tokenizer_with_both_slow_and_fast_classes(self):
+        with TemporaryHubRepo(token=self._token) as tmp_repo:
+            CustomTokenizer.register_for_auto_class()
+
+            # Fast and slow custom tokenizer
+            CustomTokenizerFast.register_for_auto_class()
+
+            with tempfile.TemporaryDirectory() as tmp_dir:
+                vocab_file = os.path.join(tmp_dir, "vocab.txt")
+                with open(vocab_file, "w", encoding="utf-8") as vocab_writer:
+                    vocab_writer.write("".join([x + "\n" for x in self.vocab_tokens]))
+
+                bert_tokenizer = BertTokenizerFast.from_pretrained(tmp_dir)
+                bert_tokenizer.save_pretrained(tmp_dir)
+                tokenizer = CustomTokenizerFast.from_pretrained(tmp_dir)
+
+            tokenizer.push_to_hub(tmp_repo.repo_id, token=self._token)
+
+            tokenizer = AutoTokenizer.from_pretrained(tmp_repo.repo_id, trust_remote_code=True)
+            # Can't make an isinstance check because the new_model.config is from the FakeConfig class of a dynamic module
+            self.assertEqual(tokenizer.__class__.__name__, "CustomTokenizerFast")
+            tokenizer = AutoTokenizer.from_pretrained(tmp_repo.repo_id, use_fast=False, trust_remote_code=True)
+            # Can't make an isinstance check because the new_model.config is from the FakeConfig class of a dynamic module
+            self.assertEqual(tokenizer.__class__.__name__, "CustomTokenizer")
+
+
+class TrieTest(unittest.TestCase):
+    def test_trie(self):
+        trie = Trie()
+        trie.add("Hello 友達")
+        self.assertEqual(trie.data, {"H": {"e": {"l": {"l": {"o": {" ": {"友": {"達": {"": 1}}}}}}}}})
+        trie.add("Hello")
+        self.assertEqual(trie.data, {"H": {"e": {"l": {"l": {"o": {"": 1, " ": {"友": {"達": {"": 1}}}}}}}}})
+
+    def test_trie_split(self):
+        trie = Trie()
+        self.assertEqual(trie.split("[CLS] This is a extra_id_100"), ["[CLS] This is a extra_id_100"])
+        trie.add("[CLS]")
+        trie.add("extra_id_1")
+        trie.add("extra_id_100")
+        self.assertEqual(trie.split("[CLS] This is a extra_id_100"), ["[CLS]", " This is a ", "extra_id_100"])
+
+    def test_trie_single(self):
+        trie = Trie()
+        trie.add("A")
+        self.assertEqual(trie.split("ABC"), ["A", "BC"])
+        self.assertEqual(trie.split("BCA"), ["BC", "A"])
+
+    def test_trie_final(self):
+        trie = Trie()
+        trie.add("TOKEN]")
+        trie.add("[SPECIAL_TOKEN]")
+        self.assertEqual(trie.split("This is something [SPECIAL_TOKEN]"), ["This is something ", "[SPECIAL_TOKEN]"])
+
+    def test_trie_subtokens(self):
+        trie = Trie()
+        trie.add("A")
+        trie.add("P")
+        trie.add("[SPECIAL_TOKEN]")
+        self.assertEqual(trie.split("This is something [SPECIAL_TOKEN]"), ["This is something ", "[SPECIAL_TOKEN]"])
+
+    def test_trie_suffix_tokens(self):
+        trie = Trie()
+        trie.add("AB")
+        trie.add("B")
+        trie.add("C")
+        self.assertEqual(trie.split("ABC"), ["AB", "C"])
+
+    def test_trie_skip(self):
+        trie = Trie()
+        trie.add("ABC")
+        trie.add("B")
+        trie.add("CD")
+        self.assertEqual(trie.split("ABCD"), ["ABC", "D"])
+
+    def test_cut_text_hardening(self):
+        # Even if the offsets are wrong, we necessarily output correct string
+        # parts.
+        trie = Trie()
+        parts = trie.cut_text("ABC", [0, 0, 2, 1, 2, 3])
+        self.assertEqual(parts, ["AB", "C"])
+
+
+class ExtensionsTrieTest(unittest.TestCase):
+    def test_extensions(self):
+        # Test searching by prefix
+        trie = ExtensionsTrie()
+        trie.add("foo")
+        trie.add("food")
+        trie.add("foodie")
+        trie.add("helium")
+        self.assertEqual(trie.extensions("foo"), ["foo", "food", "foodie"])
+        self.assertEqual(trie.extensions("helium"), ["helium"])
+
+    def test_empty_prefix(self):
+        trie = ExtensionsTrie()
+        # Test searching with an empty prefix returns all values
+        trie.add("hello")
+        trie.add("bye")
+        self.assertEqual(trie.extensions(""), ["hello", "bye"])
+
+    def test_no_extension_match(self):
+        trie = ExtensionsTrie()
+        # Test searching for a prefix that doesn't match any key
+        values = trie.extensions("unknown")
+
+        self.assertEqual(len(values), 0)
+
+    def test_update_value(self):
+        trie = ExtensionsTrie()
+        # Test updating the value of an existing key
+        trie.add("hi")
+        trie.add("hi")
+        self.assertEqual(trie.extensions("hi"), ["hi"])

docs/transformers/tests/utils/test_versions_utils.py

@@ -0,0 +1,97 @@
+# Copyright 2020 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import importlib.metadata
+import sys
+
+from transformers.testing_utils import TestCasePlus
+from transformers.utils.versions import require_version, require_version_core
+
+
+numpy_ver = importlib.metadata.version("numpy")
+python_ver = ".".join([str(x) for x in sys.version_info[:3]])
+
+
+class DependencyVersionCheckTest(TestCasePlus):
+    def test_core(self):
+        # lt + different version strings
+        require_version_core("numpy<1000.4.5")
+        require_version_core("numpy<1000.4")
+        require_version_core("numpy<1000")
+
+        # le
+        require_version_core("numpy<=1000.4.5")
+        require_version_core(f"numpy<={numpy_ver}")
+
+        # eq
+        require_version_core(f"numpy=={numpy_ver}")
+
+        # ne
+        require_version_core("numpy!=1000.4.5")
+
+        # ge
+        require_version_core("numpy>=1.0")
+        require_version_core("numpy>=1.0.0")
+        require_version_core(f"numpy>={numpy_ver}")
+
+        # gt
+        require_version_core("numpy>1.0.0")
+
+        # mix
+        require_version_core("numpy>1.0.0,<1000")
+
+        # requirement w/o version
+        require_version_core("numpy")
+
+        # unmet requirements due to version conflict
+        for req in ["numpy==1.0.0", "numpy>=1000.0.0", f"numpy<{numpy_ver}"]:
+            try:
+                require_version_core(req)
+            except ImportError as e:
+                self.assertIn(f"{req} is required", str(e))
+                self.assertIn("but found", str(e))
+
+        # unmet requirements due to missing module
+        for req in ["numpipypie>1", "numpipypie2"]:
+            try:
+                require_version_core(req)
+            except importlib.metadata.PackageNotFoundError as e:
+                self.assertIn(f"The '{req}' distribution was not found and is required by this application", str(e))
+                self.assertIn("Try: `pip install transformers -U`", str(e))
+
+        # bogus requirements formats:
+        # 1. whole thing
+        for req in ["numpy??1.0.0", "numpy1.0.0"]:
+            try:
+                require_version_core(req)
+            except ValueError as e:
+                self.assertIn("requirement needs to be in the pip package format", str(e))
+        # 2. only operators
+        for req in ["numpy=1.0.0", "numpy == 1.00", "numpy<>1.0.0", "numpy><1.00", "numpy>>1.0.0"]:
+            try:
+                require_version_core(req)
+            except ValueError as e:
+                self.assertIn("need one of ", str(e))
+
+    def test_python(self):
+        # matching requirement
+        require_version("python>=3.9.0")
+
+        # not matching requirements
+        for req in ["python>9.9.9", "python<3.0.0"]:
+            try:
+                require_version_core(req)
+            except ImportError as e:
+                self.assertIn(f"{req} is required", str(e))
+                self.assertIn(f"but found python=={python_ver}", str(e))

docs/transformers/utils/add_pipeline_model_mapping_to_test.py

@@ -0,0 +1,336 @@
+# coding=utf-8
+# Copyright 2023 The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""A script to add and/or update the attribute `pipeline_model_mapping` in model test files.
+
+This script will be (mostly) used in the following 2 situations:
+
+  - run within a (scheduled) CI job to:
+    - check if model test files in the library have updated `pipeline_model_mapping`,
+    - and/or update test files and (possibly) open a GitHub pull request automatically
+  - being run by a `transformers` member to quickly check and update some particular test file(s)
+
+This script is **NOT** intended to be run (manually) by community contributors.
+"""
+
+import argparse
+import glob
+import inspect
+import os
+import re
+import unittest
+
+from get_test_info import get_test_classes
+
+from tests.test_pipeline_mixin import pipeline_test_mapping
+
+
+PIPELINE_TEST_MAPPING = {}
+for task, _ in pipeline_test_mapping.items():
+    PIPELINE_TEST_MAPPING[task] = {"pt": None, "tf": None}
+
+
+# DO **NOT** add item to this set (unless the reason is approved)
+TEST_FILE_TO_IGNORE = {
+    "tests/models/esm/test_modeling_esmfold.py",  # The pipeline test mapping is added to `test_modeling_esm.py`
+}
+
+
+def get_framework(test_class):
+    """Infer the framework from the test class `test_class`."""
+
+    if "ModelTesterMixin" in [x.__name__ for x in test_class.__bases__]:
+        return "pt"
+    elif "TFModelTesterMixin" in [x.__name__ for x in test_class.__bases__]:
+        return "tf"
+    elif "FlaxModelTesterMixin" in [x.__name__ for x in test_class.__bases__]:
+        return "flax"
+    else:
+        return None
+
+
+def get_mapping_for_task(task, framework):
+    """Get mappings defined in `XXXPipelineTests` for the task `task`."""
+    # Use the cached results
+    if PIPELINE_TEST_MAPPING[task].get(framework, None) is not None:
+        return PIPELINE_TEST_MAPPING[task][framework]
+
+    pipeline_test_class = pipeline_test_mapping[task]["test"]
+    mapping = None
+
+    if framework == "pt":
+        mapping = getattr(pipeline_test_class, "model_mapping", None)
+    elif framework == "tf":
+        mapping = getattr(pipeline_test_class, "tf_model_mapping", None)
+
+    if mapping is not None:
+        mapping = dict(mapping.items())
+
+    # cache the results
+    PIPELINE_TEST_MAPPING[task][framework] = mapping
+    return mapping
+
+
+def get_model_for_pipeline_test(test_class, task):
+    """Get the model architecture(s) related to the test class `test_class` for a pipeline `task`."""
+    framework = get_framework(test_class)
+    if framework is None:
+        return None
+    mapping = get_mapping_for_task(task, framework)
+    if mapping is None:
+        return None
+
+    config_classes = list({model_class.config_class for model_class in test_class.all_model_classes})
+    if len(config_classes) != 1:
+        raise ValueError("There should be exactly one configuration class from `test_class.all_model_classes`.")
+
+    # This could be a list/tuple of model classes, but it's rare.
+    model_class = mapping.get(config_classes[0], None)
+    if isinstance(model_class, (tuple, list)):
+        model_class = sorted(model_class, key=lambda x: x.__name__)
+
+    return model_class
+
+
+def get_pipeline_model_mapping(test_class):
+    """Get `pipeline_model_mapping` for `test_class`."""
+    mapping = [(task, get_model_for_pipeline_test(test_class, task)) for task in pipeline_test_mapping]
+    mapping = sorted([(task, model) for task, model in mapping if model is not None], key=lambda x: x[0])
+
+    return dict(mapping)
+
+
+def get_pipeline_model_mapping_string(test_class):
+    """Get `pipeline_model_mapping` for `test_class` as a string (to be added to the test file).
+
+    This will be a 1-line string. After this is added to a test file, `make style` will format it beautifully.
+    """
+    framework = get_framework(test_class)
+    if framework == "pt":
+        framework = "torch"
+    default_value = "{}"
+
+    mapping = get_pipeline_model_mapping(test_class)
+    if len(mapping) == 0:
+        return ""
+
+    texts = []
+    for task, model_classes in mapping.items():
+        if isinstance(model_classes, (tuple, list)):
+            # A list/tuple of model classes
+            value = "(" + ", ".join([x.__name__ for x in model_classes]) + ")"
+        else:
+            # A single model class
+            value = model_classes.__name__
+        texts.append(f'"{task}": {value}')
+    text = "{" + ", ".join(texts) + "}"
+    text = f"pipeline_model_mapping = {text} if is_{framework}_available() else {default_value}"
+
+    return text
+
+
+def is_valid_test_class(test_class):
+    """Restrict to `XXXModelTesterMixin` and should be a subclass of `unittest.TestCase`."""
+    base_class_names = {"ModelTesterMixin", "TFModelTesterMixin", "FlaxModelTesterMixin"}
+    if not issubclass(test_class, unittest.TestCase):
+        return False
+    return len(base_class_names.intersection([x.__name__ for x in test_class.__bases__])) > 0
+
+
+def find_test_class(test_file):
+    """Find a test class in `test_file` to which we will add `pipeline_model_mapping`."""
+    test_classes = [x for x in get_test_classes(test_file) if is_valid_test_class(x)]
+
+    target_test_class = None
+    for test_class in test_classes:
+        # If a test class has defined `pipeline_model_mapping`, let's take it
+        if getattr(test_class, "pipeline_model_mapping", None) is not None:
+            target_test_class = test_class
+            break
+    # Take the test class with the shortest name (just a heuristic)
+    if target_test_class is None and len(test_classes) > 0:
+        target_test_class = sorted(test_classes, key=lambda x: (len(x.__name__), x.__name__))[0]
+
+    return target_test_class
+
+
+def find_block_ending(lines, start_idx, indent_level):
+    end_idx = start_idx
+    for idx, line in enumerate(lines[start_idx:]):
+        indent = len(line) - len(line.lstrip())
+        if idx == 0 or indent > indent_level or (indent == indent_level and line.strip() == ")"):
+            end_idx = start_idx + idx
+        elif idx > 0 and indent <= indent_level:
+            # Outside the definition block of `pipeline_model_mapping`
+            break
+
+    return end_idx
+
+
+def add_pipeline_model_mapping(test_class, overwrite=False):
+    """Add `pipeline_model_mapping` to `test_class`."""
+    if getattr(test_class, "pipeline_model_mapping", None) is not None:
+        if not overwrite:
+            return "", -1
+
+    line_to_add = get_pipeline_model_mapping_string(test_class)
+    if len(line_to_add) == 0:
+        return "", -1
+    line_to_add = line_to_add + "\n"
+
+    # The code defined the class `test_class`
+    class_lines, class_start_line_no = inspect.getsourcelines(test_class)
+    # `inspect` gives the code for an object, including decorator(s) if any.
+    # We (only) need the exact line of the class definition.
+    for idx, line in enumerate(class_lines):
+        if line.lstrip().startswith("class "):
+            class_lines = class_lines[idx:]
+            class_start_line_no += idx
+            break
+    class_end_line_no = class_start_line_no + len(class_lines) - 1
+
+    # The index in `class_lines` that starts the definition of `all_model_classes`, `all_generative_model_classes` or
+    # `pipeline_model_mapping`. This assumes they are defined in such order, and we take the start index of the last
+    # block that appears in a `test_class`.
+    start_idx = None
+    # The indent level of the line at `class_lines[start_idx]` (if defined)
+    indent_level = 0
+    # To record if `pipeline_model_mapping` is found in `test_class`.
+    def_line = None
+    for idx, line in enumerate(class_lines):
+        if line.strip().startswith("all_model_classes = "):
+            indent_level = len(line) - len(line.lstrip())
+            start_idx = idx
+        elif line.strip().startswith("all_generative_model_classes = "):
+            indent_level = len(line) - len(line.lstrip())
+            start_idx = idx
+        elif line.strip().startswith("pipeline_model_mapping = "):
+            indent_level = len(line) - len(line.lstrip())
+            start_idx = idx
+            def_line = line
+            break
+
+    if start_idx is None:
+        return "", -1
+    # Find the ending index (inclusive) of the above found block.
+    end_idx = find_block_ending(class_lines, start_idx, indent_level)
+
+    # Extract `is_xxx_available()` from existing blocks: some models require specific libraries like `timm` and use
+    # `is_timm_available()` instead of `is_torch_available()`.
+    # Keep leading and trailing whitespaces
+    r = re.compile(r"\s(is_\S+?_available\(\))\s")
+    for line in class_lines[start_idx : end_idx + 1]:
+        backend_condition = r.search(line)
+        if backend_condition is not None:
+            # replace the leading and trailing whitespaces to the space character " ".
+            target = " " + backend_condition[0][1:-1] + " "
+            line_to_add = r.sub(target, line_to_add)
+            break
+
+    if def_line is None:
+        # `pipeline_model_mapping` is not defined. The target index is set to the ending index (inclusive) of
+        # `all_model_classes` or `all_generative_model_classes`.
+        target_idx = end_idx
+    else:
+        # `pipeline_model_mapping` is defined. The target index is set to be one **BEFORE** its start index.
+        target_idx = start_idx - 1
+        # mark the lines of the currently existing `pipeline_model_mapping` to be removed.
+        for idx in range(start_idx, end_idx + 1):
+            # These lines are going to be removed before writing to the test file.
+            class_lines[idx] = None  # noqa
+
+    # Make sure the test class is a subclass of `PipelineTesterMixin`.
+    parent_classes = [x.__name__ for x in test_class.__bases__]
+    if "PipelineTesterMixin" not in parent_classes:
+        # Put `PipelineTesterMixin` just before `unittest.TestCase`
+        _parent_classes = [x for x in parent_classes if x != "TestCase"] + ["PipelineTesterMixin"]
+        if "TestCase" in parent_classes:
+            # Here we **assume** the original string is always with `unittest.TestCase`.
+            _parent_classes.append("unittest.TestCase")
+        parent_classes = ", ".join(_parent_classes)
+        for idx, line in enumerate(class_lines):
+            # Find the ending of the declaration of `test_class`
+            if line.strip().endswith("):"):
+                # mark the lines of the declaration of `test_class` to be removed
+                for _idx in range(idx + 1):
+                    class_lines[_idx] = None  # noqa
+                break
+        # Add the new, one-line, class declaration for `test_class`
+        class_lines[0] = f"class {test_class.__name__}({parent_classes}):\n"
+
+    # Add indentation
+    line_to_add = " " * indent_level + line_to_add
+    # Insert `pipeline_model_mapping` to `class_lines`.
+    # (The line at `target_idx` should be kept by definition!)
+    class_lines = class_lines[: target_idx + 1] + [line_to_add] + class_lines[target_idx + 1 :]
+    # Remove the lines that are marked to be removed
+    class_lines = [x for x in class_lines if x is not None]
+
+    # Move from test class to module (in order to write to the test file)
+    module_lines = inspect.getsourcelines(inspect.getmodule(test_class))[0]
+    # Be careful with the 1-off between line numbers and array indices
+    module_lines = module_lines[: class_start_line_no - 1] + class_lines + module_lines[class_end_line_no:]
+    code = "".join(module_lines)
+
+    moddule_file = inspect.getsourcefile(test_class)
+    with open(moddule_file, "w", encoding="UTF-8", newline="\n") as fp:
+        fp.write(code)
+
+    return line_to_add
+
+
+def add_pipeline_model_mapping_to_test_file(test_file, overwrite=False):
+    """Add `pipeline_model_mapping` to `test_file`."""
+    test_class = find_test_class(test_file)
+    if test_class:
+        add_pipeline_model_mapping(test_class, overwrite=overwrite)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--test_file", type=str, help="A path to the test file, starting with the repository's `tests` directory."
+    )
+    parser.add_argument(
+        "--all",
+        action="store_true",
+        help="If to check and modify all test files.",
+    )
+    parser.add_argument(
+        "--overwrite",
+        action="store_true",
+        help="If to overwrite a test class if it has already defined `pipeline_model_mapping`.",
+    )
+    args = parser.parse_args()
+
+    if not args.all and not args.test_file:
+        raise ValueError("Please specify either `test_file` or pass `--all` to check/modify all test files.")
+    elif args.all and args.test_file:
+        raise ValueError("Only one of `--test_file` and `--all` could be specified.")
+
+    test_files = []
+    if args.test_file:
+        test_files = [args.test_file]
+    else:
+        pattern = os.path.join("tests", "models", "**", "test_modeling_*.py")
+        for test_file in glob.glob(pattern):
+            # `Flax` is not concerned at this moment
+            if not test_file.startswith("test_modeling_flax_"):
+                test_files.append(test_file)
+
+    for test_file in test_files:
+        if test_file in TEST_FILE_TO_IGNORE:
+            print(f"[SKIPPED] {test_file} is skipped as it is in `TEST_FILE_TO_IGNORE` in the file {__file__}.")
+            continue
+        add_pipeline_model_mapping_to_test_file(test_file, overwrite=args.overwrite)

docs/transformers/utils/check_bad_commit.py

@@ -0,0 +1,202 @@
+#!/usr/bin/env python
+# coding=utf-8
+
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import argparse
+import json
+import os
+import re
+import subprocess
+
+import requests
+
+
+def create_script(target_test):
+    """Create a python script to be run by `git bisect run` to determine if `target_test` passes or fails.
+    If a test is not found in a commit, the script with exit code `0` (i.e. `Success`).
+
+    Args:
+        target_test (`str`): The test to check.
+
+    Returns:
+        `str`: The script to be run by `git bisect run`.
+    """
+
+    script = f"""
+import os
+import subprocess
+
+result = subprocess.run(
+    ["python3", "-m", "pytest", "-v", f"{target_test}"],
+    capture_output = True,
+    text=True,
+)
+print(result.stdout)
+
+if len(result.stderr) > 0:
+    if "ERROR: file or directory not found: " in result.stderr:
+        print("test file or directory not found in this commit")
+        exit(0)
+    elif "ERROR: not found: " in result.stderr:
+        print("test not found in this commit")
+        exit(0)
+    else:
+        print(f"pytest failed to run: {{result.stderr}}")
+        exit(-1)
+elif f"{target_test} FAILED" in result.stdout:
+    print("test failed")
+    exit(2)
+
+exit(0)
+"""
+
+    with open("target_script.py", "w") as fp:
+        fp.write(script.strip())
+
+
+def find_bad_commit(target_test, start_commit, end_commit):
+    """Find (backward) the earliest commit between `start_commit` and `end_commit` at which `target_test` fails.
+
+    Args:
+        target_test (`str`): The test to check.
+        start_commit (`str`): The latest commit.
+        end_commit (`str`): The earliest commit.
+
+    Returns:
+        `str`: The earliest commit at which `target_test` fails.
+    """
+
+    if start_commit == end_commit:
+        return start_commit
+
+    create_script(target_test=target_test)
+
+    bash = f"""
+git bisect reset
+git bisect start {start_commit} {end_commit}
+git bisect run python3 target_script.py
+"""
+
+    with open("run_git_bisect.sh", "w") as fp:
+        fp.write(bash.strip())
+
+    result = subprocess.run(
+        ["bash", "run_git_bisect.sh"],
+        capture_output=True,
+        text=True,
+    )
+    print(result.stdout)
+
+    if "error: bisect run failed" in result.stderr:
+        index = result.stderr.find("error: bisect run failed")
+        bash_error = result.stderr[index:]
+
+        error_msg = f"Error when running git bisect:\nbash error: {bash_error}"
+
+        pattern = "pytest failed to run: .+"
+        pytest_errors = re.findall(pattern, result.stdout)
+        if len(pytest_errors) > 0:
+            pytest_error = pytest_errors[0]
+            index = pytest_error.find("pytest failed to run: ")
+            index += len("pytest failed to run: ")
+            pytest_error = pytest_error[index:]
+            error_msg += f"pytest error: {pytest_error}"
+
+        raise ValueError(error_msg)
+
+    pattern = r"(.+) is the first bad commit"
+    commits = re.findall(pattern, result.stdout)
+
+    bad_commit = None
+    if len(commits) > 0:
+        bad_commit = commits[0]
+
+    print(f"Between `start_commit` {start_commit} and `end_commit` {end_commit}")
+    print(f"bad_commit: {bad_commit}\n")
+
+    return bad_commit
+
+
+def get_commit_info(commit):
+    """Get information for a commit via `api.github.com`."""
+    pr_number = None
+    author = None
+    merged_author = None
+
+    url = f"https://api.github.com/repos/huggingface/transformers/commits/{commit}/pulls"
+    pr_info_for_commit = requests.get(url).json()
+
+    if len(pr_info_for_commit) > 0:
+        pr_number = pr_info_for_commit[0]["number"]
+
+        url = f"https://api.github.com/repos/huggingface/transformers/pulls/{pr_number}"
+        pr_for_commit = requests.get(url).json()
+        author = pr_for_commit["user"]["login"]
+        merged_author = pr_for_commit["merged_by"]["login"]
+
+    if author is None:
+        url = f"https://api.github.com/repos/huggingface/transformers/commits/{commit}"
+        commit_info = requests.get(url).json()
+        author = commit_info["author"]["login"]
+
+    return {"commit": commit, "pr_number": pr_number, "author": author, "merged_by": merged_author}
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--start_commit", type=str, required=True, help="The latest commit hash to check.")
+    parser.add_argument("--end_commit", type=str, required=True, help="The earliest commit hash to check.")
+    parser.add_argument("--test", type=str, help="The test to check.")
+    parser.add_argument("--file", type=str, help="The report file.")
+    parser.add_argument("--output_file", type=str, required=True, help="The path of the output file.")
+    args = parser.parse_args()
+
+    print(f"start_commit: {args.start_commit}")
+    print(f"end_commit: {args.end_commit}")
+
+    if len({args.test is None, args.file is None}) != 2:
+        raise ValueError("Exactly one argument `test` or `file` must be specified.")
+
+    if args.test is not None:
+        commit = find_bad_commit(target_test=args.test, start_commit=args.start_commit, end_commit=args.end_commit)
+        with open(args.output_file, "w", encoding="UTF-8") as fp:
+            fp.write(f"{args.test}\n{commit}")
+    elif os.path.isfile(args.file):
+        with open(args.file, "r", encoding="UTF-8") as fp:
+            reports = json.load(fp)
+
+        for model in reports:
+            # TODO: make this script able to deal with both `single-gpu` and `multi-gpu` via a new argument.
+            reports[model].pop("multi-gpu", None)
+            failed_tests = reports[model]["single-gpu"]
+
+            failed_tests_with_bad_commits = []
+            for test in failed_tests:
+                commit = find_bad_commit(target_test=test, start_commit=args.start_commit, end_commit=args.end_commit)
+                info = {"test": test, "commit": commit}
+                info.update(get_commit_info(commit))
+                failed_tests_with_bad_commits.append(info)
+
+            # If no single-gpu test failures, remove the key
+            if len(failed_tests_with_bad_commits) > 0:
+                reports[model]["single-gpu"] = failed_tests_with_bad_commits
+            else:
+                reports[model].pop("single-gpu", None)
+
+        # remove the models without any test failure
+        reports = {k: v for k, v in reports.items() if len(v) > 0}
+
+        with open(args.output_file, "w", encoding="UTF-8") as fp:
+            json.dump(reports, fp, ensure_ascii=False, indent=4)

docs/transformers/utils/check_build.py

@@ -0,0 +1,49 @@
+# coding=utf-8
+# Copyright 2023 The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import argparse
+import importlib
+from pathlib import Path
+
+
+# Test all the extensions added in the setup
+FILES_TO_FIND = [
+    "kernels/rwkv/wkv_cuda.cu",
+    "kernels/rwkv/wkv_op.cpp",
+    "kernels/falcon_mamba/selective_scan_with_ln_interface.py",
+    "kernels/falcon_mamba/__init__.py",
+    "kernels/__init__.py",
+    "models/graphormer/algos_graphormer.pyx",
+]
+
+
+def test_custom_files_are_present(transformers_path):
+    # Test all the extensions added in the setup
+    for file in FILES_TO_FIND:
+        if not (transformers_path / file).exists():
+            return False
+    return True
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--check_lib", action="store_true", help="Whether to check the build or the actual package.")
+    args = parser.parse_args()
+    if args.check_lib:
+        transformers_module = importlib.import_module("transformers")
+        transformers_path = Path(transformers_module.__file__).parent
+    else:
+        transformers_path = Path.cwd() / "build/lib/transformers"
+    if not test_custom_files_are_present(transformers_path):
+        raise ValueError("The built release does not contain the custom files. Fix this before going further!")

docs/transformers/utils/check_config_attributes.py

@@ -0,0 +1,470 @@
+# coding=utf-8
+# Copyright 2023 The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import inspect
+import os
+import re
+
+from transformers.configuration_utils import PretrainedConfig
+from transformers.utils import direct_transformers_import
+
+
+# All paths are set with the intent you should run this script from the root of the repo with the command
+# python utils/check_config_docstrings.py
+PATH_TO_TRANSFORMERS = "src/transformers"
+
+
+# This is to make sure the transformers module imported is the one in the repo.
+transformers = direct_transformers_import(PATH_TO_TRANSFORMERS)
+
+CONFIG_MAPPING = transformers.models.auto.configuration_auto.CONFIG_MAPPING
+
+SPECIAL_CASES_TO_ALLOW = {
+    # 'max_position_embeddings' is not used in modeling file, but needed for eval frameworks like Huggingface's lighteval (https://github.com/huggingface/lighteval/blob/af24080ea4f16eaf1683e353042a2dfc9099f038/src/lighteval/models/base_model.py#L264).
+    # periods and offsets are not used in modeling file, but used in the configuration file to define `layers_block_type` and `layers_num_experts`.
+    "BambaConfig": [
+        "attn_layer_indices",
+    ],
+    "JambaConfig": [
+        "max_position_embeddings",
+        "attn_layer_offset",
+        "attn_layer_period",
+        "expert_layer_offset",
+        "expert_layer_period",
+    ],
+    "Qwen2Config": ["use_sliding_window"],
+    "Qwen2MoeConfig": ["use_sliding_window"],
+    "Qwen2VLConfig": ["use_sliding_window"],
+    # `cache_implementation` should be in the default generation config, but we don't yet support per-model
+    # generation configs (TODO joao)
+    "Gemma2Config": ["tie_word_embeddings", "cache_implementation"],
+    "Cohere2Config": ["cache_implementation"],
+    # Dropout with this value was declared but never used
+    "Phi3Config": ["embd_pdrop"],
+    # used to compute the property `self.chunk_length`
+    "EncodecConfig": ["overlap"],
+    # used to compute the property `self.layers_block_type`
+    "RecurrentGemmaConfig": ["block_types"],
+    # used as in the config to define `intermediate_size`
+    "MambaConfig": ["expand"],
+    # used as in the config to define `intermediate_size`
+    "FalconMambaConfig": ["expand"],
+    # used as `self.bert_model = BertModel(config, ...)`
+    "DPRConfig": True,
+    "FuyuConfig": True,
+    # not used in modeling files, but it's an important information
+    "FSMTConfig": ["langs"],
+    # used internally in the configuration class file
+    "GPTNeoConfig": ["attention_types"],
+    # used internally in the configuration class file
+    "EsmConfig": ["is_folding_model"],
+    # used during training (despite we don't have training script for these models yet)
+    "Mask2FormerConfig": ["ignore_value"],
+    # `ignore_value` used during training (despite we don't have training script for these models yet)
+    # `norm` used in conversion script (despite not using in the modeling file)
+    "OneFormerConfig": ["ignore_value", "norm"],
+    # used internally in the configuration class file
+    "T5Config": ["feed_forward_proj"],
+    # used internally in the configuration class file
+    # `tokenizer_class` get default value `T5Tokenizer` intentionally
+    "MT5Config": ["feed_forward_proj", "tokenizer_class"],
+    "UMT5Config": ["feed_forward_proj", "tokenizer_class"],
+    # used internally in the configuration class file
+    "LongT5Config": ["feed_forward_proj"],
+    # used internally in the configuration class file
+    "Pop2PianoConfig": ["feed_forward_proj"],
+    # used internally in the configuration class file
+    "SwitchTransformersConfig": ["feed_forward_proj"],
+    # having default values other than `1e-5` - we can't fix them without breaking
+    "BioGptConfig": ["layer_norm_eps"],
+    # having default values other than `1e-5` - we can't fix them without breaking
+    "GLPNConfig": ["layer_norm_eps"],
+    # having default values other than `1e-5` - we can't fix them without breaking
+    "SegformerConfig": ["layer_norm_eps"],
+    # having default values other than `1e-5` - we can't fix them without breaking
+    "CvtConfig": ["layer_norm_eps"],
+    # having default values other than `1e-5` - we can't fix them without breaking
+    "PerceiverConfig": ["layer_norm_eps"],
+    # used internally to calculate the feature size
+    "InformerConfig": ["num_static_real_features", "num_time_features"],
+    # used internally to calculate the feature size
+    "TimeSeriesTransformerConfig": ["num_static_real_features", "num_time_features"],
+    # used internally to calculate the feature size
+    "AutoformerConfig": ["num_static_real_features", "num_time_features"],
+    # used internally to calculate `mlp_dim`
+    "SamVisionConfig": ["mlp_ratio"],
+    # For (head) training, but so far not implemented
+    "ClapAudioConfig": ["num_classes"],
+    # Not used, but providing useful information to users
+    "SpeechT5HifiGanConfig": ["sampling_rate"],
+    # used internally in the configuration class file
+    "UdopConfig": ["feed_forward_proj"],
+    # Actually used in the config or generation config, in that case necessary for the sub-components generation
+    "SeamlessM4TConfig": [
+        "max_new_tokens",
+        "t2u_max_new_tokens",
+        "t2u_decoder_attention_heads",
+        "t2u_decoder_ffn_dim",
+        "t2u_decoder_layers",
+        "t2u_encoder_attention_heads",
+        "t2u_encoder_ffn_dim",
+        "t2u_encoder_layers",
+        "t2u_max_position_embeddings",
+    ],
+    # Actually used in the config or generation config, in that case necessary for the sub-components generation
+    "SeamlessM4Tv2Config": [
+        "max_new_tokens",
+        "t2u_decoder_attention_heads",
+        "t2u_decoder_ffn_dim",
+        "t2u_decoder_layers",
+        "t2u_encoder_attention_heads",
+        "t2u_encoder_ffn_dim",
+        "t2u_encoder_layers",
+        "t2u_max_position_embeddings",
+        "t2u_variance_pred_dropout",
+        "t2u_variance_predictor_embed_dim",
+        "t2u_variance_predictor_hidden_dim",
+        "t2u_variance_predictor_kernel_size",
+    ],
+    "ZambaConfig": [
+        "tie_word_embeddings",
+        "attn_layer_offset",
+        "attn_layer_period",
+    ],
+    "MllamaTextConfig": [
+        "initializer_range",
+    ],
+    "MllamaVisionConfig": [
+        "initializer_range",
+        "supported_aspect_ratios",
+    ],
+    "ConditionalDetrConfig": [
+        "bbox_cost",
+        "bbox_loss_coefficient",
+        "class_cost",
+        "cls_loss_coefficient",
+        "dice_loss_coefficient",
+        "focal_alpha",
+        "giou_cost",
+        "giou_loss_coefficient",
+        "mask_loss_coefficient",
+    ],
+    "DabDetrConfig": [
+        "dilation",
+        "bbox_cost",
+        "bbox_loss_coefficient",
+        "class_cost",
+        "cls_loss_coefficient",
+        "focal_alpha",
+        "giou_cost",
+        "giou_loss_coefficient",
+    ],
+    "DetrConfig": [
+        "bbox_cost",
+        "bbox_loss_coefficient",
+        "class_cost",
+        "dice_loss_coefficient",
+        "eos_coefficient",
+        "giou_cost",
+        "giou_loss_coefficient",
+        "mask_loss_coefficient",
+    ],
+    "GroundingDinoConfig": [
+        "bbox_cost",
+        "bbox_loss_coefficient",
+        "class_cost",
+        "focal_alpha",
+        "giou_cost",
+        "giou_loss_coefficient",
+    ],
+    "RTDetrConfig": [
+        "eos_coefficient",
+        "focal_loss_alpha",
+        "focal_loss_gamma",
+        "matcher_alpha",
+        "matcher_bbox_cost",
+        "matcher_class_cost",
+        "matcher_gamma",
+        "matcher_giou_cost",
+        "use_focal_loss",
+        "weight_loss_bbox",
+        "weight_loss_giou",
+        "weight_loss_vfl",
+    ],
+    "RTDetrV2Config": [
+        "eos_coefficient",
+        "focal_loss_alpha",
+        "focal_loss_gamma",
+        "matcher_alpha",
+        "matcher_bbox_cost",
+        "matcher_class_cost",
+        "matcher_gamma",
+        "matcher_giou_cost",
+        "use_focal_loss",
+        "weight_loss_bbox",
+        "weight_loss_giou",
+        "weight_loss_vfl",
+    ],
+    "YolosConfig": [
+        "bbox_cost",
+        "bbox_loss_coefficient",
+        "class_cost",
+        "eos_coefficient",
+        "giou_cost",
+        "giou_loss_coefficient",
+    ],
+    "GPTNeoXConfig": ["rotary_emb_base"],
+    "Gemma3Config": ["boi_token_index", "eoi_token_index"],
+    "Gemma3TextConfig": ["cache_implementation", "tie_word_embeddings"],
+    "ShieldGemma2Config": [
+        "boi_token_index",
+        "eoi_token_index",
+        "initializer_range",
+        "mm_tokens_per_image",
+        "text_config",
+        "vision_config",
+    ],
+    "Llama4Config": ["boi_token_index", "eoi_token_index"],
+    "Llama4TextConfig": [
+        "interleave_moe_layer_step",
+        "no_rope_layer_interval",
+        "no_rope_layers",
+        "output_router_logits",
+        "router_aux_loss_coef",
+        "router_jitter_noise",
+        "cache_implementation",
+    ],
+    "Llama4VisionConfig": ["multi_modal_projector_bias", "norm_eps"],
+}
+
+
+# TODO (ydshieh): Check the failing cases, try to fix them or move some cases to the above block once we are sure
+SPECIAL_CASES_TO_ALLOW.update(
+    {
+        "CLIPSegConfig": True,
+        "DeformableDetrConfig": True,
+        "DinatConfig": True,
+        "DonutSwinConfig": True,
+        "FastSpeech2ConformerConfig": True,
+        "FSMTConfig": True,
+        "LayoutLMv2Config": True,
+        "MaskFormerSwinConfig": True,
+        "MT5Config": True,
+        # For backward compatibility with trust remote code models
+        "MptConfig": True,
+        "MptAttentionConfig": True,
+        "OneFormerConfig": True,
+        "PerceiverConfig": True,
+        "RagConfig": True,
+        "SpeechT5Config": True,
+        "SwinConfig": True,
+        "Swin2SRConfig": True,
+        "Swinv2Config": True,
+        "SwitchTransformersConfig": True,
+        "TableTransformerConfig": True,
+        "TapasConfig": True,
+        "UniSpeechConfig": True,
+        "UniSpeechSatConfig": True,
+        "WavLMConfig": True,
+        "WhisperConfig": True,
+        # TODO: @Arthur (for `alignment_head` and `alignment_layer`)
+        "JukeboxPriorConfig": True,
+        # TODO: @Younes (for `is_decoder`)
+        "Pix2StructTextConfig": True,
+        "IdeficsConfig": True,
+        "IdeficsVisionConfig": True,
+        "IdeficsPerceiverConfig": True,
+    }
+)
+
+
+def check_attribute_being_used(config_class, attributes, default_value, source_strings):
+    """Check if any name in `attributes` is used in one of the strings in `source_strings`
+
+    Args:
+        config_class (`type`):
+            The configuration class for which the arguments in its `__init__` will be checked.
+        attributes (`List[str]`):
+            The name of an argument (or attribute) and its variant names if any.
+        default_value (`Any`):
+            A default value for the attribute in `attributes` assigned in the `__init__` of `config_class`.
+        source_strings (`List[str]`):
+            The python source code strings in the same modeling directory where `config_class` is defined. The file
+            containing the definition of `config_class` should be excluded.
+    """
+    attribute_used = False
+    for attribute in attributes:
+        for modeling_source in source_strings:
+            # check if we can find `config.xxx`, `getattr(config, "xxx", ...)` or `getattr(self.config, "xxx", ...)`
+            if (
+                f"config.{attribute}" in modeling_source
+                or f'getattr(config, "{attribute}"' in modeling_source
+                or f'getattr(self.config, "{attribute}"' in modeling_source
+                or (
+                    "TextConfig" in config_class.__name__
+                    and f"config.get_text_config().{attribute}" in modeling_source
+                )
+            ):
+                attribute_used = True
+            # Deal with multi-line cases
+            elif (
+                re.search(
+                    rf'getattr[ \t\v\n\r\f]*\([ \t\v\n\r\f]*(self\.)?config,[ \t\v\n\r\f]*"{attribute}"',
+                    modeling_source,
+                )
+                is not None
+            ):
+                attribute_used = True
+            if attribute_used:
+                break
+        if attribute_used:
+            break
+
+    # common and important attributes, even if they do not always appear in the modeling files
+    attributes_to_allow = [
+        "initializer_range",
+        "bos_index",
+        "eos_index",
+        "pad_index",
+        "unk_index",
+        "mask_index",
+        "image_token_id",  # for VLMs
+        "video_token_id",
+        "image_seq_length",
+        "video_seq_length",
+        "image_size",
+        "text_config",  # may appear as `get_text_config()`
+        "use_cache",
+        "out_features",
+        "out_indices",
+        "sampling_rate",
+        # backbone related arguments passed to load_backbone
+        "use_pretrained_backbone",
+        "backbone",
+        "backbone_config",
+        "use_timm_backbone",
+        "backbone_kwargs",
+        # rope attributes may not appear directly in the modeling but are used
+        "rope_theta",
+        "partial_rotary_factor",
+        "pretraining_tp",
+        "boi_token_index",
+        "eoi_token_index",
+    ]
+    attributes_used_in_generation = ["encoder_no_repeat_ngram_size"]
+
+    # Special cases to be allowed
+    case_allowed = True
+    if not attribute_used:
+        case_allowed = False
+        for attribute in attributes:
+            # Allow if the default value in the configuration class is different from the one in `PretrainedConfig`
+            if attribute in ["is_encoder_decoder"] and default_value is True:
+                case_allowed = True
+            elif attribute in ["tie_word_embeddings"] and default_value is False:
+                case_allowed = True
+
+            # Allow cases without checking the default value in the configuration class
+            elif attribute in attributes_to_allow + attributes_used_in_generation:
+                case_allowed = True
+            elif attribute.endswith("_token_id"):
+                case_allowed = True
+
+            # configuration class specific cases
+            if not case_allowed:
+                allowed_cases = SPECIAL_CASES_TO_ALLOW.get(config_class.__name__, [])
+                case_allowed = allowed_cases is True or attribute in allowed_cases
+
+    return attribute_used or case_allowed
+
+
+def check_config_attributes_being_used(config_class):
+    """Check the arguments in `__init__` of `config_class` are used in the modeling files in the same directory
+
+    Args:
+        config_class (`type`):
+            The configuration class for which the arguments in its `__init__` will be checked.
+    """
+    # Get the parameters in `__init__` of the configuration class, and the default values if any
+    signature = dict(inspect.signature(config_class.__init__).parameters)
+    parameter_names = [x for x in list(signature.keys()) if x not in ["self", "kwargs"]]
+    parameter_defaults = [signature[param].default for param in parameter_names]
+
+    # If `attribute_map` exists, an attribute can have different names to be used in the modeling files, and as long
+    # as one variant is used, the test should pass
+    reversed_attribute_map = {}
+    if len(config_class.attribute_map) > 0:
+        reversed_attribute_map = {v: k for k, v in config_class.attribute_map.items()}
+
+    # Get the path to modeling source files
+    config_source_file = inspect.getsourcefile(config_class)
+    model_dir = os.path.dirname(config_source_file)
+    # Let's check against all frameworks: as long as one framework uses an attribute, we are good.
+    modeling_paths = [os.path.join(model_dir, fn) for fn in os.listdir(model_dir) if fn.startswith("modeling_")]
+
+    # Get the source code strings
+    modeling_sources = []
+    for path in modeling_paths:
+        if os.path.isfile(path):
+            with open(path, encoding="utf8") as fp:
+                modeling_sources.append(fp.read())
+
+    unused_attributes = []
+    for config_param, default_value in zip(parameter_names, parameter_defaults):
+        # `attributes` here is all the variant names for `config_param`
+        attributes = [config_param]
+        # some configuration classes have non-empty `attribute_map`, and both names could be used in the
+        # corresponding modeling files. As long as one of them appears, it is fine.
+        if config_param in reversed_attribute_map:
+            attributes.append(reversed_attribute_map[config_param])
+
+        if not check_attribute_being_used(config_class, attributes, default_value, modeling_sources):
+            unused_attributes.append(attributes[0])
+
+    return sorted(unused_attributes)
+
+
+def check_config_attributes():
+    """Check the arguments in `__init__` of all configuration classes are used in python files"""
+    configs_with_unused_attributes = {}
+    for _config_class in list(CONFIG_MAPPING.values()):
+        # Skip deprecated models
+        if "models.deprecated" in _config_class.__module__:
+            continue
+        # Some config classes are not in `CONFIG_MAPPING` (e.g. `CLIPVisionConfig`, `Blip2VisionConfig`, etc.)
+        config_classes_in_module = [
+            cls
+            for name, cls in inspect.getmembers(
+                inspect.getmodule(_config_class),
+                lambda x: inspect.isclass(x)
+                and issubclass(x, PretrainedConfig)
+                and inspect.getmodule(x) == inspect.getmodule(_config_class),
+            )
+        ]
+        for config_class in config_classes_in_module:
+            unused_attributes = check_config_attributes_being_used(config_class)
+            if len(unused_attributes) > 0:
+                configs_with_unused_attributes[config_class.__name__] = unused_attributes
+
+    if len(configs_with_unused_attributes) > 0:
+        error = "The following configuration classes contain unused attributes in the corresponding modeling files:\n"
+        for name, attributes in configs_with_unused_attributes.items():
+            error += f"{name}: {attributes}\n"
+
+        raise ValueError(error)
+
+
+if __name__ == "__main__":
+    check_config_attributes()

docs/transformers/utils/check_config_docstrings.py

@@ -0,0 +1,102 @@
+# coding=utf-8
+# Copyright 2022 The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import inspect
+import re
+
+from transformers.utils import direct_transformers_import
+
+
+# All paths are set with the intent you should run this script from the root of the repo with the command
+# python utils/check_config_docstrings.py
+PATH_TO_TRANSFORMERS = "src/transformers"
+
+
+# This is to make sure the transformers module imported is the one in the repo.
+transformers = direct_transformers_import(PATH_TO_TRANSFORMERS)
+
+CONFIG_MAPPING = transformers.models.auto.configuration_auto.CONFIG_MAPPING
+
+# Regex pattern used to find the checkpoint mentioned in the docstring of `config_class`.
+# For example, `[google-bert/bert-base-uncased](https://huggingface.co/google-bert/bert-base-uncased)`
+_re_checkpoint = re.compile(r"\[(.+?)\]\((https://huggingface\.co/.+?)\)")
+
+
+CONFIG_CLASSES_TO_IGNORE_FOR_DOCSTRING_CHECKPOINT_CHECK = {
+    "DecisionTransformerConfig",
+    "EncoderDecoderConfig",
+    "MusicgenConfig",
+    "RagConfig",
+    "SpeechEncoderDecoderConfig",
+    "TimmBackboneConfig",
+    "TimmWrapperConfig",
+    "VisionEncoderDecoderConfig",
+    "VisionTextDualEncoderConfig",
+    "LlamaConfig",
+    "GraniteConfig",
+    "GraniteMoeConfig",
+    "Qwen3MoeConfig",
+    "GraniteSpeechConfig",
+}
+
+
+def get_checkpoint_from_config_class(config_class):
+    checkpoint = None
+
+    # source code of `config_class`
+    config_source = inspect.getsource(config_class)
+    checkpoints = _re_checkpoint.findall(config_source)
+
+    # Each `checkpoint` is a tuple of a checkpoint name and a checkpoint link.
+    # For example, `('google-bert/bert-base-uncased', 'https://huggingface.co/google-bert/bert-base-uncased')`
+    for ckpt_name, ckpt_link in checkpoints:
+        # allow the link to end with `/`
+        if ckpt_link.endswith("/"):
+            ckpt_link = ckpt_link[:-1]
+
+        # verify the checkpoint name corresponds to the checkpoint link
+        ckpt_link_from_name = f"https://huggingface.co/{ckpt_name}"
+        if ckpt_link == ckpt_link_from_name:
+            checkpoint = ckpt_name
+            break
+
+    return checkpoint
+
+
+def check_config_docstrings_have_checkpoints():
+    configs_without_checkpoint = []
+
+    for config_class in list(CONFIG_MAPPING.values()):
+        # Skip deprecated models
+        if "models.deprecated" in config_class.__module__:
+            continue
+        checkpoint = get_checkpoint_from_config_class(config_class)
+
+        name = config_class.__name__
+        if checkpoint is None and name not in CONFIG_CLASSES_TO_IGNORE_FOR_DOCSTRING_CHECKPOINT_CHECK:
+            configs_without_checkpoint.append(name)
+
+    if len(configs_without_checkpoint) > 0:
+        message = "\n".join(sorted(configs_without_checkpoint))
+        raise ValueError(
+            f"The following configurations don't contain any valid checkpoint:\n{message}\n\n"
+            "The requirement is to include a link pointing to one of the models of this architecture in the "
+            "docstring of the config classes listed above. The link should have be a markdown format like "
+            "[myorg/mymodel](https://huggingface.co/myorg/mymodel)."
+        )
+
+
+if __name__ == "__main__":
+    check_config_docstrings_have_checkpoints()

docs/transformers/utils/check_copies.py

@@ -0,0 +1,1078 @@
+# coding=utf-8
+# Copyright 2020 The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Utility that checks whether the copies defined in the library match the original or not. This includes:
+- All code commented with `# Copied from` comments,
+- The list of models in the main README.md matches the ones in the localized READMEs,
+- Files that are registered as full copies of one another in the `FULL_COPIES` constant of this script.
+
+This also checks the list of models in the README is complete (has all models) and add a line to complete if there is
+a model missing.
+
+Use from the root of the repo with:
+
+```bash
+python utils/check_copies.py
+```
+
+for a check that will error in case of inconsistencies (used by `make repo-consistency`) or
+
+```bash
+python utils/check_copies.py --fix_and_overwrite
+```
+
+for a check that will fix all inconsistencies automatically (used by `make fix-copies`).
+"""
+
+import argparse
+import glob
+import os
+import re
+import subprocess
+from collections import OrderedDict
+from typing import List, Optional, Tuple, Union
+
+from transformers.utils import direct_transformers_import
+
+
+# All paths are set with the intent you should run this script from the root of the repo with the command
+# python utils/check_copies.py
+TRANSFORMERS_PATH = "src/transformers"
+MODEL_TEST_PATH = "tests/models"
+PATH_TO_DOCS = "docs/source/en"
+REPO_PATH = "."
+
+# Mapping for files that are full copies of others (keys are copies, values the file to keep them up to data with)
+FULL_COPIES = {
+    "examples/tensorflow/question-answering/utils_qa.py": "examples/pytorch/question-answering/utils_qa.py",
+    "examples/flax/question-answering/utils_qa.py": "examples/pytorch/question-answering/utils_qa.py",
+}
+
+
+LOCALIZED_READMES = {
+    # If the introduction or the conclusion of the list change, the prompts may need to be updated.
+    "README.md": {
+        "start_prompt": "🤗 Transformers currently provides the following architectures",
+        "end_prompt": "1. Want to contribute a new model?",
+        "format_model_list": (
+            "**[{title}]({model_link})** (from {paper_affiliations}) released with the paper {paper_title_link} by"
+            " {paper_authors}.{supplements}"
+        ),
+    },
+    "README_zh-hans.md": {
+        "start_prompt": "🤗 Transformers 目前支持如下的架构",
+        "end_prompt": "1. 想要贡献新的模型？",
+        "format_model_list": (
+            "**[{title}]({model_link})** (来自 {paper_affiliations}) 伴随论文 {paper_title_link} 由 {paper_authors}"
+            " 发布。{supplements}"
+        ),
+    },
+    "README_zh-hant.md": {
+        "start_prompt": "🤗 Transformers 目前支援以下的架構",
+        "end_prompt": "1. 想要貢獻新的模型？",
+        "format_model_list": (
+            "**[{title}]({model_link})** (from {paper_affiliations}) released with the paper {paper_title_link} by"
+            " {paper_authors}.{supplements}"
+        ),
+    },
+    "README_ko.md": {
+        "start_prompt": "🤗 Transformers는 다음 모델들을 제공합니다",
+        "end_prompt": "1. 새로운 모델을 올리고 싶나요?",
+        "format_model_list": (
+            "**[{title}]({model_link})** ({paper_affiliations} 에서 제공)은 {paper_authors}.{supplements}의"
+            " {paper_title_link}논문과 함께 발표했습니다."
+        ),
+    },
+    "README_es.md": {
+        "start_prompt": "🤗 Transformers actualmente proporciona las siguientes arquitecturas",
+        "end_prompt": "1. ¿Quieres aportar un nuevo modelo?",
+        "format_model_list": (
+            "**[{title}]({model_link})** (from {paper_affiliations}) released with the paper {paper_title_link} by"
+            " {paper_authors}.{supplements}"
+        ),
+    },
+    "README_ja.md": {
+        "start_prompt": "🤗Transformersは現在、以下のアーキテクチャを提供しています",
+        "end_prompt": "1. 新しいモデルを投稿したいですか？",
+        "format_model_list": (
+            "**[{title}]({model_link})** ({paper_affiliations} から) {paper_authors}.{supplements} から公開された研究論文"
+            " {paper_title_link}"
+        ),
+    },
+    "README_hd.md": {
+        "start_prompt": "🤗 ट्रांसफॉर्मर वर्तमान में निम्नलिखित आर्किटेक्चर का समर्थन करते हैं",
+        "end_prompt": "1. एक नए मॉडल में योगदान देना चाहते हैं?",
+        "format_model_list": (
+            "**[{title}]({model_link})** ({paper_affiliations} से) {paper_authors}.{supplements} द्वारा"
+            "अनुसंधान पत्र {paper_title_link} के साथ जारी किया गया"
+        ),
+    },
+    "README_ru.md": {
+        "start_prompt": "🤗 В настоящее время Transformers предоставляет следующие архитектуры",
+        "end_prompt": "1. Хотите внести новую модель?",
+        "format_model_list": (
+            "**[{title}]({model_link})** (from {paper_affiliations}) released with the paper {paper_title_link} by"
+            " {paper_authors}.{supplements}"
+        ),
+    },
+    "README_pt-br.md": {
+        "start_prompt": "🤗 Transformers atualmente fornece as seguintes arquiteturas",
+        "end_prompt": "1. Quer contribuir com um novo modelo?",
+        "format_model_list": (
+            "**[{title}]({model_link})** (from {paper_affiliations}) released with the paper {paper_title_link} by"
+            " {paper_authors}.{supplements}"
+        ),
+    },
+    "README_te.md": {
+        "start_prompt": "🤗 ట్రాన్స్‌ఫార్మర్లు ప్రస్తుతం కింది ఆర్కిటెక్చర్‌లను అందజేస్తున్నాయి",
+        "end_prompt": "1. కొత్త మోడల్‌ను అందించాలనుకుంటున్నారా?",
+        "format_model_list": (
+            "**[{title}]({model_link})** (from {paper_affiliations}) released with the paper {paper_title_link} by"
+            " {paper_authors}.{supplements}"
+        ),
+    },
+    "README_fr.md": {
+        "start_prompt": "🤗 Transformers fournit actuellement les architectures suivantes",
+        "end_prompt": "1. Vous souhaitez contribuer avec un nouveau modèle ?",
+        "format_model_list": (
+            "**[{title}]({model_link})** (de {paper_affiliations}) publié dans l'article {paper_title_link} par"
+            "{paper_authors}.{supplements}"
+        ),
+    },
+    "README_de.md": {
+        "start_prompt": "🤗 Transformers bietet derzeit die folgenden Architekturen an",
+        "end_prompt": "1. Möchten Sie ein neues Modell beitragen?",
+        "format_model_list": (
+            "**[{title}]({model_link})** (from {paper_affiliations}) released with the paper {paper_title_link} by"
+            " {paper_authors}.{supplements}"
+        ),
+    },
+    "README_vi.md": {
+        "start_prompt": "🤗 Transformers hiện đang cung cấp các kiến trúc sau đây",
+        "end_prompt": "1. Muốn đóng góp một mô hình mới?",
+        "format_model_list": (
+            "**[{title}]({model_link})** (từ {paper_affiliations}) được phát hành với bài báo {paper_title_link} by"
+            " {paper_authors}.{supplements}"
+        ),
+    },
+}
+
+# This is to make sure the transformers module imported is the one in the repo.
+transformers_module = direct_transformers_import(TRANSFORMERS_PATH)
+
+
+def _is_definition_header_ending_line(line: str) -> bool:
+    # Helper function. Returns `True` if `line` is the end parenthesis of a class/function definition
+    return re.search(r"^\s*\)(\s*->.*:|:)\s*$", line) is not None
+
+
+def _should_continue(line: str, indent: str) -> bool:
+    # Helper function. Returns `True` if `line` is empty, starts with the `indent` or is the end parenthesis of a
+    # class/function definition
+    return line.startswith(indent) or len(line.strip()) == 0 or _is_definition_header_ending_line(line)
+
+
+def _sanity_check_splits(splits_1, splits_2, is_class, filename):
+    """Check the two (inner) block structures of the corresponding code block given by `split_code_into_blocks` match.
+
+    For the case of `class`, they must be of one of the following 3 cases:
+
+        - a single block without name:
+
+            class foo:
+                a = 1
+
+        - a consecutive sequence of (1 or more) blocks with name
+
+            class foo:
+
+                def f(x):
+                    return x
+
+        - a block without name, followed by a consecutive sequence of (1 or more) blocks with name
+
+            class foo:
+                a = 1
+
+                def f(x):
+                    return x
+
+                def g(x):
+                    return None
+
+    The 2 code snippets that give `splits_1` and `splits_2` have to be in the same case to pass this check, but the
+    number of blocks with name in the consecutive sequence is not taken into account.
+
+    For the case of `function or method`, we don't require it to be in one of the above 3 cases. However, the structure
+    of`splits_1` and `splits_2` have to match exactly. In particular, the number of blocks with name in a consecutive
+    sequence is taken into account.
+    """
+    block_names_1 = []
+    block_names_2 = []
+
+    for block in splits_1[1:]:
+        if block[0].startswith("_block_without_name_"):
+            block_names_1.append("block_without_name")
+        elif not block[0].startswith("_empty_block_") and (
+            not is_class or len(block_names_1) == 0 or block_names_1[-1].startswith("block_without_name")
+        ):
+            block_names_1.append("block_with_name")
+
+    for block in splits_2[1:]:
+        if block[0].startswith("_block_without_name_"):
+            block_names_2.append("block_without_name")
+        elif not block[0].startswith("_empty_block_") and (
+            not is_class or len(block_names_2) == 0 or block_names_2[-1].startswith("block_without_name")
+        ):
+            block_names_2.append("block_with_name")
+
+    if is_class:
+        if block_names_1 not in [
+            ["block_without_name"],
+            ["block_with_name"],
+            ["block_without_name", "block_with_name"],
+        ]:
+            raise ValueError(
+                f"""Class defined in {filename} doesn't have the expected structure.
+                See the docstring of `_sanity_check_splits` in the file `utils/check_copies.py`""",
+            )
+
+    if block_names_1 != block_names_2:
+        raise ValueError(f"In {filename}, two code blocks expected to be copies have different structures.")
+
+
+def find_block_end(lines: List[str], start_index: int, indent: int) -> int:
+    """
+    Find the end of the class/func block starting at `start_index` in a source code (defined by `lines`).
+
+    Args:
+        lines (`List[str]`):
+            The source code, represented by a list of lines.
+        start_index (`int`):
+            The starting index of the target class/func block.
+        indent (`int`):
+            The indent of the class/func body.
+
+    Returns:
+        `int`: The index of the block's ending line plus by 1 (i.e. exclusive).
+    """
+    indent = " " * indent
+    # enter the block body
+    line_index = start_index + 1
+
+    while line_index < len(lines) and _should_continue(lines[line_index], indent):
+        line_index += 1
+    # Clean up empty lines at the end (if any).
+    while len(lines[line_index - 1]) <= 1:
+        line_index -= 1
+
+    return line_index
+
+
+def split_code_into_blocks(
+    lines: List[str], start_index: int, end_index: int, indent: int, backtrace: bool = False
+) -> List[Tuple[str, int, int]]:
+    """
+    Split the class/func block starting at `start_index` in a source code (defined by `lines`) into *inner blocks*.
+
+    The block's header is included as the first element. The contiguous regions (without empty lines) that are not
+    inside any inner block are included as blocks. The contiguous regions of empty lines that are not inside any inner
+    block are also included as (dummy) blocks.
+
+    Args:
+        lines (`List[str]`):
+            The source code, represented by a list of lines.
+        start_index (`int`):
+            The starting index of the target class/func block.
+        end_index (`int`):
+            The ending index of the target class/func block.
+        indent (`int`):
+            The indent of the class/func body.
+        backtrace (`bool`, *optional*, defaults to `False`):
+            Whether or not to include the lines before the inner class/func block's header (e.g. comments, decorators,
+            etc.) until an empty line is encountered.
+
+    Returns:
+        `List[Tuple[str, int, int]]`: A list of elements with the form `(block_name, start_index, end_index)`.
+    """
+    splits = []
+    # `indent - 4` is the indent level of the target class/func header
+    try:
+        target_block_name = re.search(
+            rf"^{' ' * (indent - 4)}((class|def)\s+\S+)(\(|\:)", lines[start_index]
+        ).groups()[0]
+    except Exception:
+        start_context = min(start_index - 10, 0)
+        end_context = min(end_index + 10, len(lines))
+        raise ValueError(
+            f"Tried to split a class or function. It did not work. Error comes from line {start_index}: \n```\n"
+            + "".join(lines[start_context:end_context])
+            + "```\n"
+        )
+
+    # from now on, the `block` means inner blocks unless explicitly specified
+    indent_str = " " * indent
+    block_without_name_idx = 0
+    empty_block_idx = 0
+
+    # Find the lines for the definition header
+    index = start_index
+    if "(" in lines[start_index] and "):" not in lines[start_index] in lines[start_index]:
+        while index < end_index:
+            if _is_definition_header_ending_line(lines[index]):
+                break
+            index += 1
+
+    # the first line outside the definition header
+    index += 1
+    splits.append((target_block_name, start_index, index))
+
+    block_start_index, prev_block_end_index = index, index
+    while index < end_index:
+        # if found, it will be an inner block
+        block_found = re.search(rf"^{indent_str}((class|def)\s+\S+)(\(|\:)", lines[index])
+        if block_found:
+            name = block_found.groups()[0]
+
+            block_end_index = find_block_end(lines, index, indent + 4)
+
+            # backtrace to include the lines before the found block's definition header (e.g. comments, decorators,
+            # etc.) until an empty line is encountered.
+            block_start_index = index
+            if index > prev_block_end_index and backtrace:
+                idx = index - 1
+                for idx in range(index - 1, prev_block_end_index - 2, -1):
+                    if not (len(lines[idx].strip()) > 0 and lines[idx].startswith(indent_str)):
+                        break
+                idx += 1
+                if idx < index:
+                    block_start_index = idx
+
+            # between the current found block and the previous found block
+            if block_start_index > prev_block_end_index:
+                # give it a dummy name
+                if len("".join(lines[prev_block_end_index:block_start_index]).strip()) == 0:
+                    prev_block_name = f"_empty_block_{empty_block_idx}"
+                    empty_block_idx += 1
+                else:
+                    prev_block_name = f"_block_without_name_{block_without_name_idx}"
+                    block_without_name_idx += 1
+                # Add it as a block
+                splits.append((prev_block_name, prev_block_end_index, block_start_index))
+
+            # Add the current found block
+            splits.append((name, block_start_index, block_end_index))
+            prev_block_end_index = block_end_index
+            index = block_end_index - 1
+
+        index += 1
+
+    if index > prev_block_end_index:
+        if len("".join(lines[prev_block_end_index:index]).strip()) == 0:
+            prev_block_name = f"_empty_block_{empty_block_idx}"
+        else:
+            prev_block_name = f"_block_without_name_{block_without_name_idx}"
+        splits.append((prev_block_name, prev_block_end_index, index))
+
+    return splits
+
+
+def find_code_in_transformers(
+    object_name: str, base_path: str = None, return_indices: bool = False
+) -> Union[str, Tuple[List[str], int, int]]:
+    """
+    Find and return the source code of an object.
+
+    Args:
+        object_name (`str`):
+            The name of the object we want the source code of.
+        base_path (`str`, *optional*):
+            The path to the base folder where files are checked. If not set, it will be set to `TRANSFORMERS_PATH`.
+        return_indices(`bool`, *optional*, defaults to `False`):
+            If `False`, will only return the code (as a string), otherwise it will also return the whole lines of the
+            file where the object specified by `object_name` is defined, together the start/end indices of the block in
+            the file that defines the object.
+
+    Returns:
+        `Union[str, Tuple[List[str], int, int]]`: If `return_indices=False`, only the source code of the object will be
+        returned. Otherwise, it also returns the whole lines of the file where the object specified by `object_name` is
+        defined, together the start/end indices of the block in the file that defines the object.
+    """
+    parts = object_name.split(".")
+    i = 0
+
+    # We can't set this as the default value in the argument, otherwise `CopyCheckTester` will fail, as it uses a
+    # patched temp directory.
+    if base_path is None:
+        base_path = TRANSFORMERS_PATH
+
+    # Detail: the `Copied from` statement is originally designed to work with the last part of `TRANSFORMERS_PATH`,
+    # (which is `transformers`). The same should be applied for `MODEL_TEST_PATH`. However, its last part is `models`
+    # (to only check and search in it) which is a bit confusing. So we keep the copied statement staring with
+    # `tests.models.` and change it to `tests` here.
+    if base_path == MODEL_TEST_PATH:
+        base_path = "tests"
+
+    # First let's find the module where our object lives.
+    module = parts[i]
+    while i < len(parts) and not os.path.isfile(os.path.join(base_path, f"{module}.py")):
+        i += 1
+        if i < len(parts):
+            module = os.path.join(module, parts[i])
+    if i >= len(parts):
+        raise ValueError(
+            f"`object_name` should begin with the name of a module of transformers but got {object_name}."
+        )
+
+    with open(os.path.join(base_path, f"{module}.py"), "r", encoding="utf-8", newline="\n") as f:
+        lines = f.readlines()
+
+    # Now let's find the class / func in the code!
+    indent = ""
+    line_index = 0
+    for name in parts[i + 1 :]:
+        while (
+            line_index < len(lines) and re.search(rf"^{indent}(class|def)\s+{name}(\(|\:)", lines[line_index]) is None
+        ):
+            line_index += 1
+        # find the target specified in the current level in `parts` -> increase `indent` so we can search the next
+        indent += "    "
+        # the index of the first line in the (currently found) block *body*
+        line_index += 1
+
+    if line_index >= len(lines):
+        raise ValueError(f" {object_name} does not match any function or class in {module}.")
+
+    # `indent` is already one level deeper than the (found) class/func block's definition header
+
+    # We found the beginning of the class / func, now let's find the end (when the indent diminishes).
+    # `start_index` is the index of the class/func block's definition header
+    start_index = line_index - 1
+    end_index = find_block_end(lines, start_index, len(indent))
+
+    code = "".join(lines[start_index:end_index])
+    return (code, (lines, start_index, end_index)) if return_indices else code
+
+
+def replace_code(code: str, replace_pattern: str) -> str:
+    """Replace `code` by a pattern of the form `with X1->X2,Y1->Y2,Z1->Z2`.
+
+    Args:
+        code (`str`): The code to be modified.
+        replace_pattern (`str`): The pattern used to modify `code`.
+
+    Returns:
+        `str`: The modified code.
+    """
+    if len(replace_pattern) > 0:
+        patterns = replace_pattern.replace("with", "").split(",")
+        patterns = [_re_replace_pattern.search(p) for p in patterns]
+        for pattern in patterns:
+            if pattern is None:
+                continue
+            obj1, obj2, option = pattern.groups()
+            code = re.sub(obj1, obj2, code)
+            if option.strip() == "all-casing":
+                code = re.sub(obj1.lower(), obj2.lower(), code)
+                code = re.sub(obj1.upper(), obj2.upper(), code)
+
+    return code
+
+
+def find_code_and_splits(object_name: str, base_path: str, buffer: dict = None):
+    """Find the code of an object (specified by `object_name`) and split it into blocks.
+
+    Args:
+        object_name (`str`):
+            The name of the object, e.g. `transformers.models.bert.modeling_bert.BertAttention` or
+            `tests.models.llama.test_modeling_llama.LlamaModelTest.test_config`.
+        base_path (`str`):
+            The path to the base directory within which the search will be performed. It could be either
+            `TRANSFORMERS_PATH` or `MODEL_TEST_PATH`.
+        buffer (`dict`, *optional*):
+            The buffer used to store the previous results in order to speed up the process.
+
+    Returns:
+        lines (`List[str]`):
+            The lines of the whole file where the object is defined.
+        code (`str`):
+            The object's code.
+        code_splits (`List[Tuple[str, int, int]]`):
+            `code` splitted into blocks. See `split_code_into_blocks`.
+    """
+    if buffer is None:
+        buffer = {}
+
+    if (object_name, base_path) in buffer:
+        lines, code, code_splits = buffer[(object_name, base_path)]
+    else:
+        code, (lines, target_start_index, target_end_index) = find_code_in_transformers(
+            object_name, base_path=base_path, return_indices=True
+        )
+        indent = get_indent(code)
+
+        # Split the code into blocks
+        # `indent` is the indent of the class/func definition header, but `code_splits` expects the indent level of the
+        # block body.
+        code_splits = split_code_into_blocks(
+            lines, target_start_index, target_end_index, len(indent) + 4, backtrace=True
+        )
+        buffer[(object_name, base_path)] = lines, code, code_splits
+
+    return lines, code, code_splits
+
+
+_re_copy_warning = re.compile(r"^(\s*)#\s*Copied from\s+transformers\.(\S+\.\S+)\s*($|\S.*$)")
+_re_copy_warning_for_test_file = re.compile(r"^(\s*)#\s*Copied from\s+tests\.(\S+\.\S+)\s*($|\S.*$)")
+_re_replace_pattern = re.compile(r"^\s*(\S+)->(\S+)(\s+.*|$)")
+_re_fill_pattern = re.compile(r"<FILL\s+[^>]*>")
+
+
+def get_indent(code: str) -> str:
+    """
+    Find the indent in the first non empty line in a code sample.
+
+    Args:
+        code (`str`): The code to inspect.
+
+    Returns:
+        `str`: The indent looked at (as string).
+    """
+    lines = code.split("\n")
+    idx = 0
+    while idx < len(lines) and len(lines[idx]) == 0:
+        idx += 1
+    if idx < len(lines):
+        return re.search(r"^(\s*)\S", lines[idx]).groups()[0]
+    return ""
+
+
+def run_ruff(code, check=False):
+    if check:
+        command = ["ruff", "check", "-", "--fix", "--exit-zero"]
+    else:
+        command = ["ruff", "format", "-", "--config", "pyproject.toml", "--silent"]
+    process = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, stdin=subprocess.PIPE)
+    stdout, _ = process.communicate(input=code.encode())
+    return stdout.decode()
+
+
+def stylify(code: str) -> str:
+    """
+    Applies the ruff part of our `make style` command to some code. This formats the code using `ruff format`.
+    As `ruff` does not provide a python api this cannot be done on the fly.
+
+    Args:
+        code (`str`): The code to format.
+
+    Returns:
+        `str`: The formatted code.
+    """
+    has_indent = len(get_indent(code)) > 0
+    if has_indent:
+        code = f"class Bla:\n{code}"
+    formatted_code = run_ruff(code)
+    return formatted_code[len("class Bla:\n") :] if has_indent else formatted_code
+
+
+def check_codes_match(observed_code: str, theoretical_code: str) -> Optional[int]:
+    """
+    Checks if two version of a code match with the exception of the class/function name.
+
+    Args:
+        observed_code (`str`): The code found.
+        theoretical_code (`str`): The code to match.
+
+    Returns:
+        `Optional[int]`: The index of the first line where there is a difference (if any) and `None` if the codes
+        match.
+    """
+    observed_code_header = observed_code.split("\n")[0]
+    theoretical_code_header = theoretical_code.split("\n")[0]
+
+    # Catch the function/class name: it is expected that those do not match.
+    _re_class_match = re.compile(r"class\s+([^\(:]+)(?:\(|:)")
+    _re_func_match = re.compile(r"def\s+([^\(]+)\(")
+    for re_pattern in [_re_class_match, _re_func_match]:
+        if re_pattern.match(observed_code_header) is not None:
+            try:
+                observed_obj_name = re_pattern.search(observed_code_header).groups()[0]
+            except Exception:
+                raise ValueError(
+                    "Tried to split a class or function. It did not work. Error comes from: \n```\n"
+                    + observed_code_header
+                    + "\n```\n"
+                )
+
+            try:
+                theoretical_name = re_pattern.search(theoretical_code_header).groups()[0]
+            except Exception:
+                raise ValueError(
+                    "Tried to split a class or function. It did not work. Error comes from: \n```\n"
+                    + theoretical_code_header
+                    + "\n```\n"
+                )
+            theoretical_code_header = theoretical_code_header.replace(theoretical_name, observed_obj_name)
+
+    # Find the first diff. Line 0 is special since we need to compare with the function/class names ignored.
+    diff_index = 0
+    if theoretical_code_header != observed_code_header:
+        return 0
+
+    diff_index = 1
+    for observed_line, theoretical_line in zip(observed_code.split("\n")[1:], theoretical_code.split("\n")[1:]):
+        if observed_line != theoretical_line:
+            return diff_index
+        diff_index += 1
+
+
+def is_copy_consistent(filename: str, overwrite: bool = False, buffer: dict = None) -> Optional[List[Tuple[str, int]]]:
+    """
+    Check if the code commented as a copy in a file matches the original.
+
+    Args:
+        filename (`str`):
+            The name of the file to check.
+        overwrite (`bool`, *optional*, defaults to `False`):
+            Whether or not to overwrite the copies when they don't match.
+        buffer (`dict`, *optional*):
+            The buffer used to store the previous results in order to speed up the process.
+
+    Returns:
+        `Optional[List[Tuple[str, int]]]`: If `overwrite=False`, returns the list of differences as tuples `(str, int)`
+        with the name of the object having a diff and the line number where there is the first diff.
+    """
+    base_path = TRANSFORMERS_PATH if not filename.startswith("tests") else MODEL_TEST_PATH
+
+    with open(filename, "r", encoding="utf-8", newline="\n") as f:
+        lines = f.readlines()
+    diffs = []
+    line_index = 0
+    # Not a for loop cause `lines` is going to change (if `overwrite=True`).
+    search_re = _re_copy_warning_for_test_file if filename.startswith("tests") else _re_copy_warning
+    while line_index < len(lines):
+        search = search_re.search(lines[line_index])
+        if search is None:
+            line_index += 1
+            continue
+
+        # There is some copied code here, let's retrieve the original.
+        indent, object_name, replace_pattern = search.groups()
+
+        # Find the file lines, the object's code, and its blocks
+        try:
+            target_lines, theoretical_code, theoretical_code_splits = find_code_and_splits(
+                object_name, base_path, buffer=buffer
+            )
+        except Exception as exc:
+            exc.args = (f"Error while trying to find source code for {filename}.\n\n" + str(exc),)
+            raise
+
+        # code replaced by the patterns
+        theoretical_code_blocks = OrderedDict()
+        for name, start, end in theoretical_code_splits:
+            name = replace_code(name, replace_pattern)
+            code = "".join(target_lines[start:end])
+            code = replace_code(code, replace_pattern)
+            theoretical_code_blocks[name] = code
+
+        theoretical_indent = get_indent(theoretical_code)
+
+        # `start_index` is the index of the first line (the definition header) after `# Copied from`.
+        # (`indent != theoretical_indent` doesn't seem to occur so far, not sure what this case is for.)
+        start_index = line_index + 1 if indent == theoretical_indent else line_index
+        # enter the block body
+        line_index = start_index + 1
+
+        subcode = "\n".join(theoretical_code.split("\n")[1:])
+        indent = get_indent(subcode)
+        # Loop to check the observed code, stop when indentation diminishes or if we see a End copy comment.
+        # We can't call `find_block_end` directly as there is sth. special `# End copy"` here.
+        should_continue = True
+        while line_index < len(lines) and should_continue:
+            line_index += 1
+            if line_index >= len(lines):
+                break
+            line = lines[line_index]
+            # There is a special pattern `# End copy` to stop early. It's not documented cause it shouldn't really be
+            # used.
+            should_continue = _should_continue(line, indent) and re.search(f"^{indent}# End copy", line) is None
+        # `line_index` is outside the block
+        # Clean up empty lines at the end (if any).
+        while len(lines[line_index - 1]) <= 1:
+            line_index -= 1
+
+        # Split the observed code into blocks
+        observed_code_splits = split_code_into_blocks(lines, start_index, line_index, len(indent), backtrace=True)
+
+        is_class = lines[start_index].startswith(f"{' ' * (len(indent) - 4)}class ")
+        # sanity check
+        _sanity_check_splits(theoretical_code_splits, observed_code_splits, is_class=is_class, filename=filename)
+
+        # observed code in a structured way (a dict mapping block names to blocks' code)
+        observed_code_blocks = OrderedDict()
+        for name, start, end in observed_code_splits:
+            code = "".join(lines[start:end])
+            observed_code_blocks[name] = code
+
+        # Below, we change some names in `theoretical_code_blocks` and `observed_code_blocks`. These mappings map the
+        # original names to the modified names: this is used to restore the original order of the code blocks.
+        name_mappings_1 = {k: k for k in theoretical_code_blocks.keys()}
+        name_mappings_2 = {k: k for k in observed_code_blocks.keys()}
+
+        # Update code blocks' name and content:
+        #   If `"# Ignore copy"` is found in a block of the observed code:
+        #     1. if it's a block only in the observed code --> add it to the theoretical code.
+        #     2. if it's also in the theoretical code () --> put its content (body) to the corresponding block under the
+        #        same name in the theoretical code.
+        #   In both cases, we change the name to have a prefix `_ignored_` so we know if we can discard them during the
+        #   comparison.
+        ignored_existing_block_index = 0
+        ignored_new_block_index = 0
+        for name in list(observed_code_blocks.keys()):
+            code = observed_code_blocks[name]
+            if "# Ignore copy" in code:
+                if name in theoretical_code_blocks:
+                    # in the target --> just copy the content
+                    del theoretical_code_blocks[name]
+                    theoretical_code_blocks[f"_ignored_existing_block_{ignored_existing_block_index}"] = code
+                    name_mappings_1[name] = f"_ignored_existing_block_{ignored_existing_block_index}"
+
+                    del observed_code_blocks[name]
+                    observed_code_blocks[f"_ignored_existing_block_{ignored_existing_block_index}"] = code
+                    name_mappings_2[name] = f"_ignored_existing_block_{ignored_existing_block_index}"
+                    ignored_existing_block_index += 1
+                else:
+                    # not in the target --> add it
+                    theoretical_code_blocks[f"_ignored_new_block_{ignored_new_block_index}"] = code
+                    name_mappings_1[f"_ignored_new_block_{ignored_new_block_index}"] = (
+                        f"_ignored_new_block_{ignored_new_block_index}"
+                    )
+
+                    del observed_code_blocks[name]
+                    observed_code_blocks[f"_ignored_new_block_{ignored_new_block_index}"] = code
+                    name_mappings_2[name] = f"_ignored_new_block_{ignored_new_block_index}"
+                    ignored_new_block_index += 1
+
+        # Respect the original block order:
+        #   1. in `theoretical_code_blocks`: the new blocks will follow the existing ones
+        #   2. in `observed_code_blocks`: the original order are kept with names modified potentially. This is necessary
+        #      to compute the correct `diff_index` if `overwrite=True` and there is a diff.
+        theoretical_code_blocks = {
+            name_mappings_1[orig_name]: theoretical_code_blocks[name_mappings_1[orig_name]]
+            for orig_name in name_mappings_1
+        }
+        observed_code_blocks = {
+            name_mappings_2[orig_name]: observed_code_blocks[name_mappings_2[orig_name]]
+            for orig_name in name_mappings_2
+        }
+
+        # Ignore the blocks specified to be ignored. This is the version used to check if there is a mismatch
+        theoretical_code_blocks_clean = {
+            k: v
+            for k, v in theoretical_code_blocks.items()
+            if not (k.startswith(("_ignored_existing_block_", "_ignored_new_block_")))
+        }
+        theoretical_code = "".join(list(theoretical_code_blocks_clean.values()))
+
+        # stylify `theoretical_code` before compare (this is needed only when `replace_pattern` is not empty)
+        if replace_pattern:
+            theoretical_code = stylify(theoretical_code)
+        # Remove `\n\n` in `theoretical_code` before compare (so no empty line)
+        while "\n\n" in theoretical_code:
+            theoretical_code = theoretical_code.replace("\n\n", "\n")
+
+        # Compute `observed_code` where we don't include any empty line + keep track the line index between the
+        # original/processed `observed_code` so we can have the correct `diff_index`.
+        idx_to_orig_idx_mapping_for_observed_code_lines = {}
+        idx = -1
+        orig_idx = -1
+        observed_code = ""
+        for name, code in observed_code_blocks.items():
+            if code.endswith("\n"):
+                code = code[:-1]
+            for code_line in code.split("\n"):
+                orig_idx += 1
+                if code_line.strip() and not name.startswith(("_ignored_existing_block_", "_ignored_new_block_")):
+                    idx += 1
+                    observed_code += code_line + "\n"
+                    idx_to_orig_idx_mapping_for_observed_code_lines[idx] = orig_idx
+
+        # Test for a diff and act accordingly.
+        diff_index = check_codes_match(observed_code, theoretical_code)
+        if diff_index is not None:
+            # switch to the index in the original `observed_code` (i.e. before removing empty lines)
+            diff_index = idx_to_orig_idx_mapping_for_observed_code_lines[diff_index]
+            diffs.append([object_name, diff_index + start_index + 1])
+            if overwrite:
+                # `theoretical_code_to_write` is a single string but may have several lines.
+                theoretical_code_to_write = stylify("".join(list(theoretical_code_blocks.values())))
+                lines = lines[:start_index] + [theoretical_code_to_write] + lines[line_index:]
+                # Here we treat it as a single entry in `lines`.
+                line_index = start_index + 1
+
+    if overwrite and len(diffs) > 0:
+        # Warn the user a file has been modified.
+        print(f"Detected changes, rewriting {filename}.")
+        with open(filename, "w", encoding="utf-8", newline="\n") as f:
+            f.writelines(lines)
+    return diffs
+
+
+def check_copies(overwrite: bool = False, file: str = None):
+    """
+    Check every file is copy-consistent with the original. Also check the model list in the main README and other
+    READMEs are consistent.
+
+    Args:
+        overwrite (`bool`, *optional*, defaults to `False`):
+            Whether or not to overwrite the copies when they don't match.
+        file (`bool`, *optional*):
+            The path to a specific file to check and/or fix.
+    """
+    buffer = {}
+
+    if file is None:
+        all_files = glob.glob(os.path.join(TRANSFORMERS_PATH, "**/*.py"), recursive=True)
+        all_test_files = glob.glob(os.path.join(MODEL_TEST_PATH, "**/*.py"), recursive=True)
+        all_files = list(all_files) + list(all_test_files)
+    else:
+        all_files = [file]
+
+    diffs = []
+    for filename in all_files:
+        new_diffs = is_copy_consistent(filename, overwrite, buffer)
+        diffs += [f"- {filename}: copy does not match {d[0]} at line {d[1]}" for d in new_diffs]
+    if not overwrite and len(diffs) > 0:
+        diff = "\n".join(diffs)
+        raise Exception(
+            "Found the following copy inconsistencies:\n"
+            + diff
+            + "\nRun `make fix-copies` or `python utils/check_copies.py --fix_and_overwrite` to fix them."
+        )
+
+
+def check_full_copies(overwrite: bool = False):
+    """
+    Check the files that are full copies of others (as indicated in `FULL_COPIES`) are copy-consistent.
+
+    Args:
+        overwrite (`bool`, *optional*, defaults to `False`):
+            Whether or not to overwrite the copies when they don't match.
+    """
+    diffs = []
+    for target, source in FULL_COPIES.items():
+        with open(source, "r", encoding="utf-8") as f:
+            source_code = f.read()
+        with open(target, "r", encoding="utf-8") as f:
+            target_code = f.read()
+        if source_code != target_code:
+            if overwrite:
+                with open(target, "w", encoding="utf-8") as f:
+                    print(f"Replacing the content of {target} by the one of {source}.")
+                    f.write(source_code)
+            else:
+                diffs.append(f"- {target}: copy does not match {source}.")
+
+    if not overwrite and len(diffs) > 0:
+        diff = "\n".join(diffs)
+        raise Exception(
+            "Found the following copy inconsistencies:\n"
+            + diff
+            + "\nRun `make fix-copies` or `python utils/check_copies.py --fix_and_overwrite` to fix them."
+        )
+
+
+def get_model_list(filename: str, start_prompt: str, end_prompt: str) -> str:
+    """
+    Extracts the model list from a README.
+
+    Args:
+        filename (`str`): The name of the README file to check.
+        start_prompt (`str`): The string to look for that introduces the model list.
+        end_prompt (`str`): The string to look for that ends the model list.
+
+    Returns:
+        `str`: The model list.
+    """
+    with open(os.path.join(REPO_PATH, filename), "r", encoding="utf-8", newline="\n") as f:
+        lines = f.readlines()
+    # Find the start of the list.
+    start_index = 0
+    while not lines[start_index].startswith(start_prompt):
+        start_index += 1
+    start_index += 1
+
+    result = []
+    current_line = ""
+    end_index = start_index
+
+    # Keep going until the end of the list.
+    while not lines[end_index].startswith(end_prompt):
+        if lines[end_index].startswith("1."):
+            if len(current_line) > 1:
+                result.append(current_line)
+            current_line = lines[end_index]
+        elif len(lines[end_index]) > 1:
+            current_line = f"{current_line[:-1]} {lines[end_index].lstrip()}"
+        end_index += 1
+    if len(current_line) > 1:
+        result.append(current_line)
+
+    return "".join(result)
+
+
+def convert_to_localized_md(model_list: str, localized_model_list: str, format_str: str) -> Tuple[bool, str]:
+    """
+    Compare the model list from the main README to the one in a localized README.
+
+    Args:
+        model_list (`str`): The model list in the main README.
+        localized_model_list (`str`): The model list in one of the localized README.
+        format_str (`str`):
+            The template for a model entry in the localized README (look at the `format_model_list` in the entries of
+            `LOCALIZED_READMES` for examples).
+
+    Returns:
+        `Tuple[bool, str]`: A tuple where the first value indicates if the READMEs match or not, and the second value
+        is the correct localized README.
+    """
+
+    def _rep(match):
+        title, model_link, paper_affiliations, paper_title_link, paper_authors, supplements = match.groups()
+        return format_str.format(
+            title=title,
+            model_link=model_link,
+            paper_affiliations=paper_affiliations,
+            paper_title_link=paper_title_link,
+            paper_authors=paper_authors,
+            supplements=" " + supplements.strip() if len(supplements) != 0 else "",
+        )
+
+    # This regex captures metadata from an English model description, including model title, model link,
+    # affiliations of the paper, title of the paper, authors of the paper, and supplemental data (see DistilBERT for
+    # example).
+    _re_capture_meta = re.compile(
+        r"\*\*\[([^\]]*)\]\(([^\)]*)\)\*\* \(from ([^)]*)\)[^\[]*([^\)]*\)).*?by (.*?[A-Za-z\*]{2,}?)\. (.*)$"
+    )
+    # This regex is used to synchronize title link.
+    _re_capture_title_link = re.compile(r"\*\*\[([^\]]*)\]\(([^\)]*)\)\*\*")
+    # This regex is used to synchronize paper title and link.
+    _re_capture_paper_link = re.compile(r" \[([^\]]*)\]\(([^\)]*)\)")
+
+    if len(localized_model_list) == 0:
+        localized_model_index = {}
+    else:
+        try:
+            localized_model_index = {
+                re.search(r"\*\*\[([^\]]*)", line).groups()[0]: line
+                for line in localized_model_list.strip().split("\n")
+            }
+        except AttributeError:
+            raise AttributeError("A model name in localized READMEs cannot be recognized.")
+
+    model_keys = [re.search(r"\*\*\[([^\]]*)", line).groups()[0] for line in model_list.strip().split("\n")]
+
+    # We exclude keys in localized README not in the main one.
+    readmes_match = not any(k not in model_keys for k in localized_model_index)
+    localized_model_index = {k: v for k, v in localized_model_index.items() if k in model_keys}
+
+    for model in model_list.strip().split("\n"):
+        title, model_link = _re_capture_title_link.search(model).groups()
+        if title not in localized_model_index:
+            readmes_match = False
+            # Add an anchor white space behind a model description string for regex.
+            # If metadata cannot be captured, the English version will be directly copied.
+            localized_model_index[title] = _re_capture_meta.sub(_rep, model + " ")
+        elif _re_fill_pattern.search(localized_model_index[title]) is not None:
+            update = _re_capture_meta.sub(_rep, model + " ")
+            if update != localized_model_index[title]:
+                readmes_match = False
+                localized_model_index[title] = update
+        else:
+            # Synchronize title link
+            converted_model = _re_capture_title_link.sub(
+                f"**[{title}]({model_link})**", localized_model_index[title], count=1
+            )
+
+            # Synchronize paper title and its link (if found)
+            paper_title_link = _re_capture_paper_link.search(model)
+            if paper_title_link is not None:
+                paper_title, paper_link = paper_title_link.groups()
+                converted_model = _re_capture_paper_link.sub(
+                    f" [{paper_title}]({paper_link})", converted_model, count=1
+                )
+
+            if converted_model != localized_model_index[title]:
+                readmes_match = False
+                localized_model_index[title] = converted_model
+
+    sorted_index = sorted(localized_model_index.items(), key=lambda x: x[0].lower())
+
+    return readmes_match, "\n".join((x[1] for x in sorted_index)) + "\n"
+
+
+# Map a model name with the name it has in the README for the check_readme check
+SPECIAL_MODEL_NAMES = {
+    "Bert Generation": "BERT For Sequence Generation",
+    "BigBird": "BigBird-RoBERTa",
+    "Data2VecAudio": "Data2Vec",
+    "Data2VecText": "Data2Vec",
+    "Data2VecVision": "Data2Vec",
+    "DonutSwin": "Swin Transformer",
+    "Marian": "MarianMT",
+    "MaskFormerSwin": "Swin Transformer",
+    "OpenAI GPT-2": "GPT-2",
+    "OpenAI GPT": "GPT",
+    "Perceiver": "Perceiver IO",
+    "SAM": "Segment Anything",
+    "ViT": "Vision Transformer (ViT)",
+}
+
+# Update this list with the models that shouldn't be in the README. This only concerns modular models or those who do
+# not have an associated paper.
+MODELS_NOT_IN_README = [
+    "BertJapanese",
+    "Encoder decoder",
+    "FairSeq Machine-Translation",
+    "HerBERT",
+    "RetriBERT",
+    "Speech Encoder decoder",
+    "Speech2Text",
+    "Speech2Text2",
+    "TimmBackbone",
+    "Vision Encoder decoder",
+    "VisionTextDualEncoder",
+    "CLIPVisionModel",
+    "SiglipVisionModel",
+    "ChineseCLIPVisionModel",
+    "VitPoseBackbone",
+]
+
+# Template for new entries to add in the main README when we have missing models.
+README_TEMPLATE = (
+    "1. **[{model_name}](https://huggingface.co/docs/main/transformers/model_doc/{model_type})** (from "
+    "<FILL INSTITUTION>) released with the paper [<FILL PAPER TITLE>](<FILL ARKIV LINK>) by <FILL AUTHORS>."
+)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--file", type=str, default=None, help="A specific file to check and/or fix")
+    parser.add_argument("--fix_and_overwrite", action="store_true", help="Whether to fix inconsistencies.")
+    args = parser.parse_args()
+
+    check_copies(args.fix_and_overwrite, args.file)
+    check_full_copies(args.fix_and_overwrite)

docs/transformers/utils/check_doc_toc.py

@@ -0,0 +1,134 @@
+# coding=utf-8
+# Copyright 2022 The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+This script is responsible for cleaning the model section of the table of content by removing duplicates and sorting
+the entries in alphabetical order.
+
+Usage (from the root of the repo):
+
+Check that the table of content is properly sorted (used in `make quality`):
+
+```bash
+python utils/check_doc_toc.py
+```
+
+Auto-sort the table of content if it is not properly sorted (used in `make style`):
+
+```bash
+python utils/check_doc_toc.py --fix_and_overwrite
+```
+"""
+
+import argparse
+from collections import defaultdict
+from typing import List
+
+import yaml
+
+
+PATH_TO_TOC = "docs/source/en/_toctree.yml"
+
+
+def clean_model_doc_toc(model_doc: List[dict]) -> List[dict]:
+    """
+    Cleans a section of the table of content of the model documentation (one specific modality) by removing duplicates
+    and sorting models alphabetically.
+
+    Args:
+        model_doc (`List[dict]`):
+            The list of dictionaries extracted from the `_toctree.yml` file for this specific modality.
+
+    Returns:
+        `List[dict]`: List of dictionaries like the input, but cleaned up and sorted.
+    """
+    counts = defaultdict(int)
+    for doc in model_doc:
+        counts[doc["local"]] += 1
+    duplicates = [key for key, value in counts.items() if value > 1]
+
+    new_doc = []
+    for duplicate_key in duplicates:
+        titles = list({doc["title"] for doc in model_doc if doc["local"] == duplicate_key})
+        if len(titles) > 1:
+            raise ValueError(
+                f"{duplicate_key} is present several times in the documentation table of content at "
+                "`docs/source/en/_toctree.yml` with different *Title* values. Choose one of those and remove the "
+                "others."
+            )
+        # Only add this once
+        new_doc.append({"local": duplicate_key, "title": titles[0]})
+
+    # Add none duplicate-keys
+    new_doc.extend([doc for doc in model_doc if counts[doc["local"]] == 1])
+
+    # Sort
+    return sorted(new_doc, key=lambda s: s["title"].lower())
+
+
+def check_model_doc(overwrite: bool = False):
+    """
+    Check that the content of the table of content in `_toctree.yml` is clean (no duplicates and sorted for the model
+    API doc) and potentially auto-cleans it.
+
+    Args:
+        overwrite (`bool`, *optional*, defaults to `False`):
+            Whether to just check if the TOC is clean or to auto-clean it (when `overwrite=True`).
+    """
+    with open(PATH_TO_TOC, encoding="utf-8") as f:
+        content = yaml.safe_load(f.read())
+
+    # Get to the API doc
+    api_idx = 0
+    while content[api_idx]["title"] != "API":
+        api_idx += 1
+    api_doc = content[api_idx]["sections"]
+
+    # Then to the model doc
+    model_idx = 0
+    while api_doc[model_idx]["title"] != "Models":
+        model_idx += 1
+
+    model_doc = api_doc[model_idx]["sections"]
+
+    # Extract the modalities and clean them one by one.
+    modalities_docs = [(idx, section) for idx, section in enumerate(model_doc) if "sections" in section]
+    diff = False
+    for idx, modality_doc in modalities_docs:
+        old_modality_doc = modality_doc["sections"]
+        new_modality_doc = clean_model_doc_toc(old_modality_doc)
+
+        if old_modality_doc != new_modality_doc:
+            diff = True
+            if overwrite:
+                model_doc[idx]["sections"] = new_modality_doc
+
+    if diff:
+        if overwrite:
+            api_doc[model_idx]["sections"] = model_doc
+            content[api_idx]["sections"] = api_doc
+            with open(PATH_TO_TOC, "w", encoding="utf-8") as f:
+                f.write(yaml.dump(content, allow_unicode=True))
+        else:
+            raise ValueError(
+                "The model doc part of the table of content is not properly sorted, run `make style` to fix this."
+            )
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--fix_and_overwrite", action="store_true", help="Whether to fix inconsistencies.")
+    args = parser.parse_args()
+
+    check_model_doc(args.fix_and_overwrite)

docs/transformers/utils/check_docstrings.py

@@ -0,0 +1,1061 @@
+# coding=utf-8
+# Copyright 2023 The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Utility that checks all docstrings of public objects have an argument section matching their signature.
+
+Use from the root of the repo with:
+
+```bash
+python utils/check_docstrings.py
+```
+
+for a check that will error in case of inconsistencies (used by `make repo-consistency`).
+
+To auto-fix issues run:
+
+```bash
+python utils/check_docstrings.py --fix_and_overwrite
+```
+
+which is used by `make fix-copies` (note that this fills what it cans, you might have to manually fill information
+like argument descriptions).
+"""
+
+import argparse
+import ast
+import enum
+import inspect
+import operator as op
+import re
+from pathlib import Path
+from typing import Any, Optional, Tuple, Union
+
+from check_repo import ignore_undocumented
+from git import Repo
+
+from transformers.utils import direct_transformers_import
+
+
+PATH_TO_REPO = Path(__file__).parent.parent.resolve()
+PATH_TO_TRANSFORMERS = Path("src").resolve() / "transformers"
+
+# This is to make sure the transformers module imported is the one in the repo.
+transformers = direct_transformers_import(PATH_TO_TRANSFORMERS)
+
+OPTIONAL_KEYWORD = "*optional*"
+# Re pattern that catches args blocks in docstrings (with all variation around the name supported).
+_re_args = re.compile(r"^\s*(Args?|Arguments?|Attributes?|Params?|Parameters?):\s*$")
+# Re pattern that parses the start of an arg block: catches <name> (<description>) in those lines.
+_re_parse_arg = re.compile(r"^(\s*)(\S+)\s+\((.+)\)(?:\:|$)")
+# Re pattern that parses the end of a description of an arg (catches the default in *optional*, defaults to xxx).
+_re_parse_description = re.compile(r"\*optional\*, defaults to (.*)$")
+
+
+# This is a temporary list of objects to ignore while we progressively fix them. Do not add anything here, fix the
+# docstrings instead. If formatting should be ignored for the docstring, you can put a comment # no-format on the
+# line before the docstring.
+OBJECTS_TO_IGNORE = [
+    "Llama4Processor",
+    # Deprecated
+    "InputExample",
+    "InputFeatures",
+    # Signature is *args/**kwargs
+    "TFSequenceSummary",
+    "TFBertTokenizer",
+    "TFGPT2Tokenizer",
+    # Missing arguments in the docstring
+    "ASTFeatureExtractor",
+    "AlbertModel",
+    "AlbertTokenizerFast",
+    "AlignTextModel",
+    "AlignVisionConfig",
+    "AudioClassificationPipeline",
+    "AutoformerConfig",
+    "AutomaticSpeechRecognitionPipeline",
+    "BarkCoarseConfig",
+    "BarkConfig",
+    "BarkFineConfig",
+    "BarkSemanticConfig",
+    "BartConfig",
+    "BartTokenizerFast",
+    "BarthezTokenizerFast",
+    "BeitModel",
+    "BertConfig",
+    "BertJapaneseTokenizer",
+    "BertModel",
+    "BertTokenizerFast",
+    "BigBirdConfig",
+    "BigBirdForQuestionAnswering",
+    "BigBirdModel",
+    "BigBirdPegasusConfig",
+    "BigBirdTokenizerFast",
+    "BitImageProcessor",
+    "BlenderbotConfig",
+    "BlenderbotSmallConfig",
+    "BlenderbotSmallTokenizerFast",
+    "BlenderbotTokenizerFast",
+    "Blip2VisionConfig",
+    "BlipTextConfig",
+    "BlipVisionConfig",
+    "BloomConfig",
+    "BloomTokenizerFast",
+    "BridgeTowerTextConfig",
+    "BridgeTowerVisionConfig",
+    "BrosModel",
+    "CamembertConfig",
+    "CamembertModel",
+    "CamembertTokenizerFast",
+    "CanineModel",
+    "CanineTokenizer",
+    "ChineseCLIPTextModel",
+    "ClapTextConfig",
+    "ConditionalDetrConfig",
+    "ConditionalDetrImageProcessor",
+    "ConvBertConfig",
+    "ConvBertTokenizerFast",
+    "ConvNextConfig",
+    "ConvNextV2Config",
+    "CpmAntTokenizer",
+    "CvtConfig",
+    "CvtModel",
+    "DeiTImageProcessor",
+    "DPRReaderTokenizer",
+    "DPRReaderTokenizerFast",
+    "DPTModel",
+    "Data2VecAudioConfig",
+    "Data2VecTextConfig",
+    "Data2VecTextModel",
+    "Data2VecVisionModel",
+    "DataCollatorForLanguageModeling",
+    "DebertaConfig",
+    "DebertaV2Config",
+    "DebertaV2Tokenizer",
+    "DebertaV2TokenizerFast",
+    "DecisionTransformerConfig",
+    "DeformableDetrConfig",
+    "DeformableDetrImageProcessor",
+    "DeiTModel",
+    "DepthEstimationPipeline",
+    "DetaConfig",
+    "DetaImageProcessor",
+    "DetrConfig",
+    "DetrImageProcessor",
+    "DinatModel",
+    "DistilBertConfig",
+    "DistilBertTokenizerFast",
+    "DocumentQuestionAnsweringPipeline",
+    "DonutSwinModel",
+    "EarlyStoppingCallback",
+    "EfficientFormerConfig",
+    "EfficientFormerImageProcessor",
+    "EfficientNetConfig",
+    "ElectraConfig",
+    "ElectraTokenizerFast",
+    "EncoderDecoderModel",
+    "ErnieMModel",
+    "ErnieModel",
+    "ErnieMTokenizer",
+    "EsmConfig",
+    "EsmModel",
+    "FlaxAlbertForMaskedLM",
+    "FlaxAlbertForMultipleChoice",
+    "FlaxAlbertForPreTraining",
+    "FlaxAlbertForQuestionAnswering",
+    "FlaxAlbertForSequenceClassification",
+    "FlaxAlbertForTokenClassification",
+    "FlaxAlbertModel",
+    "FlaxBartForCausalLM",
+    "FlaxBartForConditionalGeneration",
+    "FlaxBartForQuestionAnswering",
+    "FlaxBartForSequenceClassification",
+    "FlaxBartModel",
+    "FlaxBeitForImageClassification",
+    "FlaxBeitForMaskedImageModeling",
+    "FlaxBeitModel",
+    "FlaxBertForCausalLM",
+    "FlaxBertForMaskedLM",
+    "FlaxBertForMultipleChoice",
+    "FlaxBertForNextSentencePrediction",
+    "FlaxBertForPreTraining",
+    "FlaxBertForQuestionAnswering",
+    "FlaxBertForSequenceClassification",
+    "FlaxBertForTokenClassification",
+    "FlaxBertModel",
+    "FlaxBigBirdForCausalLM",
+    "FlaxBigBirdForMaskedLM",
+    "FlaxBigBirdForMultipleChoice",
+    "FlaxBigBirdForPreTraining",
+    "FlaxBigBirdForQuestionAnswering",
+    "FlaxBigBirdForSequenceClassification",
+    "FlaxBigBirdForTokenClassification",
+    "FlaxBigBirdModel",
+    "FlaxBlenderbotForConditionalGeneration",
+    "FlaxBlenderbotModel",
+    "FlaxBlenderbotSmallForConditionalGeneration",
+    "FlaxBlenderbotSmallModel",
+    "FlaxBloomForCausalLM",
+    "FlaxBloomModel",
+    "FlaxCLIPModel",
+    "FlaxDinov2ForImageClassification",
+    "FlaxDinov2Model",
+    "FlaxDistilBertForMaskedLM",
+    "FlaxDistilBertForMultipleChoice",
+    "FlaxDistilBertForQuestionAnswering",
+    "FlaxDistilBertForSequenceClassification",
+    "FlaxDistilBertForTokenClassification",
+    "FlaxDistilBertModel",
+    "FlaxElectraForCausalLM",
+    "FlaxElectraForMaskedLM",
+    "FlaxElectraForMultipleChoice",
+    "FlaxElectraForPreTraining",
+    "FlaxElectraForQuestionAnswering",
+    "FlaxElectraForSequenceClassification",
+    "FlaxElectraForTokenClassification",
+    "FlaxElectraModel",
+    "FlaxEncoderDecoderModel",
+    "FlaxGPT2LMHeadModel",
+    "FlaxGPT2Model",
+    "FlaxGPTJForCausalLM",
+    "FlaxGPTJModel",
+    "FlaxGPTNeoForCausalLM",
+    "FlaxGPTNeoModel",
+    "FlaxLlamaForCausalLM",
+    "FlaxLlamaModel",
+    "FlaxGemmaForCausalLM",
+    "FlaxGemmaModel",
+    "FlaxMBartForConditionalGeneration",
+    "FlaxMBartForQuestionAnswering",
+    "FlaxMBartForSequenceClassification",
+    "FlaxMBartModel",
+    "FlaxMarianMTModel",
+    "FlaxMarianModel",
+    "FlaxMistralForCausalLM",
+    "FlaxMistralModel",
+    "FlaxOPTForCausalLM",
+    "FlaxPegasusForConditionalGeneration",
+    "FlaxPegasusModel",
+    "FlaxRegNetForImageClassification",
+    "FlaxRegNetModel",
+    "FlaxResNetForImageClassification",
+    "FlaxResNetModel",
+    "FlaxRoFormerForMaskedLM",
+    "FlaxRoFormerForMultipleChoice",
+    "FlaxRoFormerForQuestionAnswering",
+    "FlaxRoFormerForSequenceClassification",
+    "FlaxRoFormerForTokenClassification",
+    "FlaxRoFormerModel",
+    "FlaxRobertaForCausalLM",
+    "FlaxRobertaForMaskedLM",
+    "FlaxRobertaForMultipleChoice",
+    "FlaxRobertaForQuestionAnswering",
+    "FlaxRobertaForSequenceClassification",
+    "FlaxRobertaForTokenClassification",
+    "FlaxRobertaModel",
+    "FlaxRobertaPreLayerNormForCausalLM",
+    "FlaxRobertaPreLayerNormForMaskedLM",
+    "FlaxRobertaPreLayerNormForMultipleChoice",
+    "FlaxRobertaPreLayerNormForQuestionAnswering",
+    "FlaxRobertaPreLayerNormForSequenceClassification",
+    "FlaxRobertaPreLayerNormForTokenClassification",
+    "FlaxRobertaPreLayerNormModel",
+    "FlaxSpeechEncoderDecoderModel",
+    "FlaxViTForImageClassification",
+    "FlaxViTModel",
+    "FlaxVisionEncoderDecoderModel",
+    "FlaxVisionTextDualEncoderModel",
+    "FlaxWav2Vec2ForCTC",
+    "FlaxWav2Vec2ForPreTraining",
+    "FlaxWav2Vec2Model",
+    "FlaxWhisperForAudioClassification",
+    "FlaxWhisperForConditionalGeneration",
+    "FlaxWhisperModel",
+    "FlaxWhisperTimeStampLogitsProcessor",
+    "FlaxXGLMForCausalLM",
+    "FlaxXGLMModel",
+    "FlaxXLMRobertaForCausalLM",
+    "FlaxXLMRobertaForMaskedLM",
+    "FlaxXLMRobertaForMultipleChoice",
+    "FlaxXLMRobertaForQuestionAnswering",
+    "FlaxXLMRobertaForSequenceClassification",
+    "FlaxXLMRobertaForTokenClassification",
+    "FlaxXLMRobertaModel",
+    "FNetConfig",
+    "FNetModel",
+    "FNetTokenizerFast",
+    "FSMTConfig",
+    "FeatureExtractionPipeline",
+    "FillMaskPipeline",
+    "FlaubertConfig",
+    "FlavaConfig",
+    "FlavaForPreTraining",
+    "FlavaImageModel",
+    "FlavaImageProcessor",
+    "FlavaMultimodalModel",
+    "FlavaTextConfig",
+    "FlavaTextModel",
+    "FocalNetModel",
+    "FunnelTokenizerFast",
+    "GPTBigCodeConfig",
+    "GPTJConfig",
+    "GPTNeoXConfig",
+    "GPTNeoXJapaneseConfig",
+    "GPTNeoXTokenizerFast",
+    "GPTSanJapaneseConfig",
+    "GitConfig",
+    "GitVisionConfig",
+    "GraphormerConfig",
+    "GroupViTTextConfig",
+    "GroupViTVisionConfig",
+    "HerbertTokenizerFast",
+    "HubertConfig",
+    "HubertForCTC",
+    "IBertConfig",
+    "IBertModel",
+    "IdeficsConfig",
+    "IdeficsProcessor",
+    "IJepaModel",
+    "ImageClassificationPipeline",
+    "ImageFeatureExtractionPipeline",
+    "ImageGPTConfig",
+    "ImageSegmentationPipeline",
+    "ImageTextToTextPipeline",
+    "ImageToImagePipeline",
+    "ImageToTextPipeline",
+    "InformerConfig",
+    "JukeboxPriorConfig",
+    "JukeboxTokenizer",
+    "LEDConfig",
+    "LEDTokenizerFast",
+    "LayoutLMForQuestionAnswering",
+    "LayoutLMTokenizerFast",
+    "LayoutLMv2Config",
+    "LayoutLMv2ForQuestionAnswering",
+    "LayoutLMv2TokenizerFast",
+    "LayoutLMv3Config",
+    "LayoutLMv3ImageProcessor",
+    "LayoutLMv3TokenizerFast",
+    "LayoutXLMTokenizerFast",
+    "LevitConfig",
+    "LiltConfig",
+    "LiltModel",
+    "LongT5Config",
+    "LongformerConfig",
+    "LongformerModel",
+    "LongformerTokenizerFast",
+    "LukeModel",
+    "LukeTokenizer",
+    "LxmertTokenizerFast",
+    "M2M100Config",
+    "M2M100Tokenizer",
+    "MarkupLMProcessor",
+    "MaskGenerationPipeline",
+    "MBart50TokenizerFast",
+    "MBartConfig",
+    "MCTCTFeatureExtractor",
+    "MPNetConfig",
+    "MPNetModel",
+    "MPNetTokenizerFast",
+    "MT5Config",
+    "MT5TokenizerFast",
+    "MarianConfig",
+    "MarianTokenizer",
+    "MarkupLMConfig",
+    "MarkupLMModel",
+    "MarkupLMTokenizer",
+    "MarkupLMTokenizerFast",
+    "Mask2FormerConfig",
+    "MaskFormerConfig",
+    "MaxTimeCriteria",
+    "MegaConfig",
+    "MegaModel",
+    "MegatronBertConfig",
+    "MegatronBertForPreTraining",
+    "MegatronBertModel",
+    "MLCDVisionConfig",
+    "MobileBertConfig",
+    "MobileBertModel",
+    "MobileBertTokenizerFast",
+    "MobileNetV1ImageProcessor",
+    "MobileNetV1Model",
+    "MobileNetV2ImageProcessor",
+    "MobileNetV2Model",
+    "MobileViTModel",
+    "MobileViTV2Model",
+    "MLukeTokenizer",
+    "MraConfig",
+    "MusicgenDecoderConfig",
+    "MusicgenForConditionalGeneration",
+    "MusicgenMelodyForConditionalGeneration",
+    "MvpConfig",
+    "MvpTokenizerFast",
+    "MT5Tokenizer",
+    "NatModel",
+    "NerPipeline",
+    "NezhaConfig",
+    "NezhaModel",
+    "NllbMoeConfig",
+    "NllbTokenizer",
+    "NllbTokenizerFast",
+    "NystromformerConfig",
+    "OPTConfig",
+    "ObjectDetectionPipeline",
+    "OneFormerProcessor",
+    "OpenAIGPTTokenizerFast",
+    "OpenLlamaConfig",
+    "PLBartConfig",
+    "PegasusConfig",
+    "PegasusTokenizer",
+    "PegasusTokenizerFast",
+    "PegasusXConfig",
+    "PerceiverImageProcessor",
+    "PerceiverModel",
+    "PerceiverTokenizer",
+    "PersimmonConfig",
+    "Pipeline",
+    "Pix2StructConfig",
+    "Pix2StructTextConfig",
+    "PLBartTokenizer",
+    "Pop2PianoConfig",
+    "PreTrainedTokenizer",
+    "PreTrainedTokenizerBase",
+    "PreTrainedTokenizerFast",
+    "PrefixConstrainedLogitsProcessor",
+    "ProphetNetConfig",
+    "QDQBertConfig",
+    "QDQBertModel",
+    "QuestionAnsweringPipeline",
+    "RagConfig",
+    "RagModel",
+    "RagRetriever",
+    "RagSequenceForGeneration",
+    "RagTokenForGeneration",
+    "RealmConfig",
+    "RealmForOpenQA",
+    "RealmScorer",
+    "RealmTokenizerFast",
+    "ReformerConfig",
+    "ReformerTokenizerFast",
+    "RegNetConfig",
+    "RemBertConfig",
+    "RemBertModel",
+    "RemBertTokenizer",
+    "RemBertTokenizerFast",
+    "RetriBertConfig",
+    "RetriBertTokenizerFast",
+    "RoCBertConfig",
+    "RoCBertModel",
+    "RoCBertTokenizer",
+    "RoFormerConfig",
+    "RobertaConfig",
+    "RobertaModel",
+    "RobertaPreLayerNormConfig",
+    "RobertaPreLayerNormModel",
+    "RobertaTokenizerFast",
+    "SEWConfig",
+    "SEWDConfig",
+    "SEWDForCTC",
+    "SEWForCTC",
+    "SamConfig",
+    "SamPromptEncoderConfig",
+    "SeamlessM4TConfig",  # use of unconventional markdown
+    "SeamlessM4Tv2Config",  # use of unconventional markdown
+    "Seq2SeqTrainingArguments",
+    "SpecialTokensMixin",
+    "Speech2Text2Config",
+    "Speech2Text2Tokenizer",
+    "Speech2TextTokenizer",
+    "SpeechEncoderDecoderModel",
+    "SpeechT5Config",
+    "SpeechT5Model",
+    "SplinterConfig",
+    "SplinterTokenizerFast",
+    "SqueezeBertTokenizerFast",
+    "SummarizationPipeline",
+    "Swin2SRImageProcessor",
+    "Swinv2Model",
+    "SwitchTransformersConfig",
+    "T5Config",
+    "T5Tokenizer",
+    "T5TokenizerFast",
+    "TableQuestionAnsweringPipeline",
+    "TableTransformerConfig",
+    "TapasConfig",
+    "TapasModel",
+    "TapasTokenizer",
+    "Text2TextGenerationPipeline",
+    "TextClassificationPipeline",
+    "TextGenerationPipeline",
+    "TFBartForConditionalGeneration",
+    "TFBartForSequenceClassification",
+    "TFBartModel",
+    "TFBertModel",
+    "TFConvNextModel",
+    "TFData2VecVisionModel",
+    "TFDeiTModel",
+    "TFEncoderDecoderModel",
+    "TFEsmModel",
+    "TFMobileViTModel",
+    "TFRagModel",
+    "TFRagSequenceForGeneration",
+    "TFRagTokenForGeneration",
+    "TFRepetitionPenaltyLogitsProcessor",
+    "TFSwinModel",
+    "TFViTModel",
+    "TFVisionEncoderDecoderModel",
+    "TFVisionTextDualEncoderModel",
+    "TFXGLMForCausalLM",
+    "TFXGLMModel",
+    "TimeSeriesTransformerConfig",
+    "TokenClassificationPipeline",
+    "TrOCRConfig",
+    "Phi4MultimodalProcessor",
+    "TrainerState",
+    "TrainingArguments",
+    "TrajectoryTransformerConfig",
+    "TranslationPipeline",
+    "TvltImageProcessor",
+    "UMT5Config",
+    "UperNetConfig",
+    "UperNetForSemanticSegmentation",
+    "ViTHybridImageProcessor",
+    "ViTHybridModel",
+    "ViTMSNModel",
+    "ViTModel",
+    "VideoClassificationPipeline",
+    "ViltConfig",
+    "ViltForImagesAndTextClassification",
+    "ViltModel",
+    "VisionEncoderDecoderModel",
+    "VisionTextDualEncoderModel",
+    "VisualBertConfig",
+    "VisualBertModel",
+    "VisualQuestionAnsweringPipeline",
+    "VitMatteForImageMatting",
+    "VitsTokenizer",
+    "VivitModel",
+    "Wav2Vec2BertForCTC",
+    "Wav2Vec2CTCTokenizer",
+    "Wav2Vec2Config",
+    "Wav2Vec2ConformerConfig",
+    "Wav2Vec2ConformerForCTC",
+    "Wav2Vec2FeatureExtractor",
+    "Wav2Vec2PhonemeCTCTokenizer",
+    "WavLMConfig",
+    "WavLMForCTC",
+    "WhisperConfig",
+    "WhisperFeatureExtractor",
+    "WhisperForAudioClassification",
+    "XCLIPTextConfig",
+    "XCLIPVisionConfig",
+    "XGLMConfig",
+    "XGLMModel",
+    "XGLMTokenizerFast",
+    "XLMConfig",
+    "XLMProphetNetConfig",
+    "XLMRobertaConfig",
+    "XLMRobertaModel",
+    "XLMRobertaTokenizerFast",
+    "XLMRobertaXLConfig",
+    "XLMRobertaXLModel",
+    "XLNetConfig",
+    "XLNetTokenizerFast",
+    "XmodConfig",
+    "XmodModel",
+    "YolosImageProcessor",
+    "YolosModel",
+    "YosoConfig",
+    "ZeroShotAudioClassificationPipeline",
+    "ZeroShotClassificationPipeline",
+    "ZeroShotImageClassificationPipeline",
+    "ZeroShotObjectDetectionPipeline",
+    "Llama4TextConfig",
+]
+
+# Supported math operations when interpreting the value of defaults.
+MATH_OPERATORS = {
+    ast.Add: op.add,
+    ast.Sub: op.sub,
+    ast.Mult: op.mul,
+    ast.Div: op.truediv,
+    ast.Pow: op.pow,
+    ast.BitXor: op.xor,
+    ast.USub: op.neg,
+}
+
+
+def find_indent(line: str) -> int:
+    """
+    Returns the number of spaces that start a line indent.
+    """
+    search = re.search(r"^(\s*)(?:\S|$)", line)
+    if search is None:
+        return 0
+    return len(search.groups()[0])
+
+
+def stringify_default(default: Any) -> str:
+    """
+    Returns the string representation of a default value, as used in docstring: numbers are left as is, all other
+    objects are in backtiks.
+
+    Args:
+        default (`Any`): The default value to process
+
+    Returns:
+        `str`: The string representation of that default.
+    """
+    if isinstance(default, bool):
+        # We need to test for bool first as a bool passes isinstance(xxx, (int, float))
+        return f"`{default}`"
+    elif isinstance(default, enum.Enum):
+        # We need to test for enum first as an enum with int values will pass isinstance(xxx, (int, float))
+        return f"`{str(default)}`"
+    elif isinstance(default, int):
+        return str(default)
+    elif isinstance(default, float):
+        result = str(default)
+        return str(round(default, 2)) if len(result) > 6 else result
+    elif isinstance(default, str):
+        return str(default) if default.isnumeric() else f'`"{default}"`'
+    elif isinstance(default, type):
+        return f"`{default.__name__}`"
+    else:
+        return f"`{default}`"
+
+
+def eval_math_expression(expression: str) -> Optional[Union[float, int]]:
+    # Mainly taken from the excellent https://stackoverflow.com/a/9558001
+    """
+    Evaluate (safely) a mathematial expression and returns its value.
+
+    Args:
+        expression (`str`): The expression to evaluate.
+
+    Returns:
+        `Optional[Union[float, int]]`: Returns `None` if the evaluation fails in any way and the value computed
+        otherwise.
+
+    Example:
+
+    ```py
+    >>> eval_expr('2^6')
+    4
+    >>> eval_expr('2**6')
+    64
+    >>> eval_expr('1 + 2*3**(4^5) / (6 + -7)')
+    -5.0
+    ```
+    """
+    try:
+        return eval_node(ast.parse(expression, mode="eval").body)
+    except TypeError:
+        return
+
+
+def eval_node(node):
+    if isinstance(node, ast.Num):  # <number>
+        return node.n
+    elif isinstance(node, ast.BinOp):  # <left> <operator> <right>
+        return MATH_OPERATORS[type(node.op)](eval_node(node.left), eval_node(node.right))
+    elif isinstance(node, ast.UnaryOp):  # <operator> <operand> e.g., -1
+        return MATH_OPERATORS[type(node.op)](eval_node(node.operand))
+    else:
+        raise TypeError(node)
+
+
+def replace_default_in_arg_description(description: str, default: Any) -> str:
+    """
+    Catches the default value in the description of an argument inside a docstring and replaces it by the value passed.
+
+    Args:
+        description (`str`): The description of an argument in a docstring to process.
+        default (`Any`): The default value that would be in the docstring of that argument.
+
+    Returns:
+       `str`: The description updated with the new default value.
+    """
+    # Lots of docstrings have `optional` or **opational** instead of *optional* so we do this fix here.
+    description = description.replace("`optional`", OPTIONAL_KEYWORD)
+    description = description.replace("**optional**", OPTIONAL_KEYWORD)
+    if default is inspect._empty:
+        # No default, make sure the description doesn't have any either
+        idx = description.find(OPTIONAL_KEYWORD)
+        if idx != -1:
+            description = description[:idx].rstrip()
+            if description.endswith(","):
+                description = description[:-1].rstrip()
+    elif default is None:
+        # Default None are not written, we just set `*optional*`. If there is default that is not None specified in the
+        # description, we do not erase it (as sometimes we set the default to `None` because the default is a mutable
+        # object).
+        idx = description.find(OPTIONAL_KEYWORD)
+        if idx == -1:
+            description = f"{description}, {OPTIONAL_KEYWORD}"
+        elif re.search(r"defaults to `?None`?", description) is not None:
+            len_optional = len(OPTIONAL_KEYWORD)
+            description = description[: idx + len_optional]
+    else:
+        str_default = None
+        # For numbers we may have a default that is given by a math operation (1/255 is really popular). We don't
+        # want to replace those by their actual values.
+        if isinstance(default, (int, float)) and re.search("defaults to `?(.*?)(?:`|$)", description) is not None:
+            # Grab the default and evaluate it.
+            current_default = re.search("defaults to `?(.*?)(?:`|$)", description).groups()[0]
+            if default == eval_math_expression(current_default):
+                try:
+                    # If it can be directly converted to the type of the default, it's a simple value
+                    str_default = str(type(default)(current_default))
+                except Exception:
+                    # Otherwise there is a math operator so we add a code block.
+                    str_default = f"`{current_default}`"
+            elif isinstance(default, enum.Enum) and default.name == current_default.split(".")[-1]:
+                # When the default is an Enum (this is often the case for PIL.Image.Resampling), and the docstring
+                # matches the enum name, keep the existing docstring rather than clobbering it with the enum value.
+                str_default = f"`{current_default}`"
+
+        if str_default is None:
+            str_default = stringify_default(default)
+        # Make sure default match
+        if OPTIONAL_KEYWORD not in description:
+            description = f"{description}, {OPTIONAL_KEYWORD}, defaults to {str_default}"
+        elif _re_parse_description.search(description) is None:
+            idx = description.find(OPTIONAL_KEYWORD)
+            len_optional = len(OPTIONAL_KEYWORD)
+            description = f"{description[: idx + len_optional]}, defaults to {str_default}"
+        else:
+            description = _re_parse_description.sub(rf"*optional*, defaults to {str_default}", description)
+
+    return description
+
+
+def get_default_description(arg: inspect.Parameter) -> str:
+    """
+    Builds a default description for a parameter that was not documented.
+
+    Args:
+        arg (`inspect.Parameter`): The argument in the signature to generate a description for.
+
+    Returns:
+        `str`: The description.
+    """
+    if arg.annotation is inspect._empty:
+        arg_type = "<fill_type>"
+    elif hasattr(arg.annotation, "__name__"):
+        arg_type = arg.annotation.__name__
+    else:
+        arg_type = str(arg.annotation)
+
+    if arg.default is inspect._empty:
+        return f"`{arg_type}`"
+    elif arg.default is None:
+        return f"`{arg_type}`, {OPTIONAL_KEYWORD}"
+    else:
+        str_default = stringify_default(arg.default)
+        return f"`{arg_type}`, {OPTIONAL_KEYWORD}, defaults to {str_default}"
+
+
+def find_source_file(obj: Any) -> Path:
+    """
+    Finds the source file of an object.
+
+    Args:
+        obj (`Any`): The object whose source file we are looking for.
+
+    Returns:
+        `Path`: The source file.
+    """
+    module = obj.__module__
+    obj_file = PATH_TO_TRANSFORMERS
+    for part in module.split(".")[1:]:
+        obj_file = obj_file / part
+    return obj_file.with_suffix(".py")
+
+
+def match_docstring_with_signature(obj: Any) -> Optional[Tuple[str, str]]:
+    """
+    Matches the docstring of an object with its signature.
+
+    Args:
+        obj (`Any`): The object to process.
+
+    Returns:
+        `Optional[Tuple[str, str]]`: Returns `None` if there is no docstring or no parameters documented in the
+        docstring, otherwise returns a tuple of two strings: the current documentation of the arguments in the
+        docstring and the one matched with the signature.
+    """
+    if len(getattr(obj, "__doc__", "")) == 0:
+        # Nothing to do, there is no docstring.
+        return
+
+    # Read the docstring in the source code to see if there is a special command to ignore this object.
+    try:
+        source, _ = inspect.getsourcelines(obj)
+    except OSError:
+        source = []
+
+    idx = 0
+    while idx < len(source) and '"""' not in source[idx]:
+        idx += 1
+
+    ignore_order = False
+    if idx < len(source):
+        line_before_docstring = source[idx - 1]
+        if re.search(r"^\s*#\s*no-format\s*$", line_before_docstring):
+            # This object is ignored
+            return
+        elif re.search(r"^\s*#\s*ignore-order\s*$", line_before_docstring):
+            ignore_order = True
+
+    # Read the signature
+    signature = inspect.signature(obj).parameters
+
+    obj_doc_lines = obj.__doc__.split("\n")
+    # Get to the line where we start documenting arguments
+    idx = 0
+    while idx < len(obj_doc_lines) and _re_args.search(obj_doc_lines[idx]) is None:
+        idx += 1
+
+    if idx == len(obj_doc_lines):
+        # Nothing to do, no parameters are documented.
+        return
+
+    if "kwargs" in signature and signature["kwargs"].annotation != inspect._empty:
+        # Inspecting signature with typed kwargs is not supported yet.
+        return
+
+    indent = find_indent(obj_doc_lines[idx])
+    arguments = {}
+    current_arg = None
+    idx += 1
+    start_idx = idx
+    # Keep going until the arg section is finished (nonempty line at the same indent level) or the end of the docstring.
+    while idx < len(obj_doc_lines) and (
+        len(obj_doc_lines[idx].strip()) == 0 or find_indent(obj_doc_lines[idx]) > indent
+    ):
+        if find_indent(obj_doc_lines[idx]) == indent + 4:
+            # New argument -> let's generate the proper doc for it
+            re_search_arg = _re_parse_arg.search(obj_doc_lines[idx])
+            if re_search_arg is not None:
+                _, name, description = re_search_arg.groups()
+                current_arg = name
+                if name in signature:
+                    default = signature[name].default
+                    if signature[name].kind is inspect._ParameterKind.VAR_KEYWORD:
+                        default = None
+                    new_description = replace_default_in_arg_description(description, default)
+                else:
+                    new_description = description
+                init_doc = _re_parse_arg.sub(rf"\1\2 ({new_description}):", obj_doc_lines[idx])
+                arguments[current_arg] = [init_doc]
+        elif current_arg is not None:
+            arguments[current_arg].append(obj_doc_lines[idx])
+
+        idx += 1
+
+    # We went too far by one (perhaps more if there are a lot of new lines)
+    idx -= 1
+    if current_arg:
+        while len(obj_doc_lines[idx].strip()) == 0:
+            arguments[current_arg] = arguments[current_arg][:-1]
+            idx -= 1
+    # And we went too far by one again.
+    idx += 1
+
+    old_doc_arg = "\n".join(obj_doc_lines[start_idx:idx])
+
+    old_arguments = list(arguments.keys())
+    arguments = {name: "\n".join(doc) for name, doc in arguments.items()}
+    # Add missing arguments with a template
+    for name in set(signature.keys()) - set(arguments.keys()):
+        arg = signature[name]
+        # We ignore private arguments or *args/**kwargs (unless they are documented by the user)
+        if name.startswith("_") or arg.kind in [
+            inspect._ParameterKind.VAR_KEYWORD,
+            inspect._ParameterKind.VAR_POSITIONAL,
+        ]:
+            arguments[name] = ""
+        else:
+            arg_desc = get_default_description(arg)
+            arguments[name] = " " * (indent + 4) + f"{name} ({arg_desc}): <fill_docstring>"
+
+    # Arguments are sorted by the order in the signature unless a special comment is put.
+    if ignore_order:
+        new_param_docs = [arguments[name] for name in old_arguments if name in signature]
+        missing = set(signature.keys()) - set(old_arguments)
+        new_param_docs.extend([arguments[name] for name in missing if len(arguments[name]) > 0])
+    else:
+        new_param_docs = [arguments[name] for name in signature.keys() if len(arguments[name]) > 0]
+    new_doc_arg = "\n".join(new_param_docs)
+
+    return old_doc_arg, new_doc_arg
+
+
+def fix_docstring(obj: Any, old_doc_args: str, new_doc_args: str):
+    """
+    Fixes the docstring of an object by replacing its arguments documentation by the one matched with the signature.
+
+    Args:
+        obj (`Any`):
+            The object whose dostring we are fixing.
+        old_doc_args (`str`):
+            The current documentation of the parameters of `obj` in the docstring (as returned by
+            `match_docstring_with_signature`).
+        new_doc_args (`str`):
+            The documentation of the parameters of `obj` matched with its signature (as returned by
+            `match_docstring_with_signature`).
+    """
+    # Read the docstring in the source code and make sure we have the right part of the docstring
+    source, line_number = inspect.getsourcelines(obj)
+
+    # Get to the line where we start documenting arguments
+    idx = 0
+    while idx < len(source) and _re_args.search(source[idx]) is None:
+        idx += 1
+
+    if idx == len(source):
+        # Args are not defined in the docstring of this object
+        return
+
+    # Get to the line where we stop documenting arguments
+    indent = find_indent(source[idx])
+    idx += 1
+    start_idx = idx
+    while idx < len(source) and (len(source[idx].strip()) == 0 or find_indent(source[idx]) > indent):
+        idx += 1
+
+    idx -= 1
+    while len(source[idx].strip()) == 0:
+        idx -= 1
+    idx += 1
+
+    if "".join(source[start_idx:idx])[:-1] != old_doc_args:
+        # Args are not fully defined in the docstring of this object
+        return
+
+    obj_file = find_source_file(obj)
+    with open(obj_file, "r", encoding="utf-8") as f:
+        content = f.read()
+
+    # Replace content
+    lines = content.split("\n")
+    lines = lines[: line_number + start_idx - 1] + [new_doc_args] + lines[line_number + idx - 1 :]
+
+    print(f"Fixing the docstring of {obj.__name__} in {obj_file}.")
+    with open(obj_file, "w", encoding="utf-8") as f:
+        f.write("\n".join(lines))
+
+
+def check_docstrings(overwrite: bool = False, check_all: bool = False):
+    """
+    Check docstrings of all public objects that are callables and are documented. By default, only checks the diff.
+
+    Args:
+        overwrite (`bool`, *optional*, defaults to `False`):
+            Whether to fix inconsistencies or not.
+        check_all (`bool`, *optional*, defaults to `False`):
+            Whether to check all files.
+    """
+    module_diff_files = None
+    if not check_all:
+        module_diff_files = set()
+        repo = Repo(PATH_TO_REPO)
+        # Diff from index to unstaged files
+        for modified_file_diff in repo.index.diff(None):
+            if modified_file_diff.a_path.startswith("src/transformers"):
+                module_diff_files.add(modified_file_diff.a_path)
+        # Diff from index to `main`
+        for modified_file_diff in repo.index.diff(repo.refs.main.commit):
+            if modified_file_diff.a_path.startswith("src/transformers"):
+                module_diff_files.add(modified_file_diff.a_path)
+        # quick escape route: if there are no module files in the diff, skip this check
+        if len(module_diff_files) == 0:
+            return
+        print("    Checking docstrings in the following files:" + "\n    - " + "\n    - ".join(module_diff_files))
+
+    failures = []
+    hard_failures = []
+    to_clean = []
+    for name in dir(transformers):
+        # Skip objects that are private or not documented.
+        if name.startswith("_") or ignore_undocumented(name) or name in OBJECTS_TO_IGNORE:
+            continue
+
+        obj = getattr(transformers, name)
+        if not callable(obj) or not isinstance(obj, type) or getattr(obj, "__doc__", None) is None:
+            continue
+
+        # If we are checking against the diff, we skip objects that are not part of the diff.
+        if module_diff_files is not None:
+            object_file = find_source_file(getattr(transformers, name))
+            object_file_relative_path = "src/" + str(object_file).split("/src/")[1]
+            if object_file_relative_path not in module_diff_files:
+                continue
+
+        # Check docstring
+        try:
+            result = match_docstring_with_signature(obj)
+            if result is not None:
+                old_doc, new_doc = result
+            else:
+                old_doc, new_doc = None, None
+        except Exception as e:
+            print(e)
+            hard_failures.append(name)
+            continue
+        if old_doc != new_doc:
+            if overwrite:
+                fix_docstring(obj, old_doc, new_doc)
+            else:
+                failures.append(name)
+        elif not overwrite and new_doc is not None and ("<fill_type>" in new_doc or "<fill_docstring>" in new_doc):
+            to_clean.append(name)
+
+    # Deal with errors
+    error_message = ""
+    if len(hard_failures) > 0:
+        error_message += (
+            "The argument part of the docstrings of the following objects could not be processed, check they are "
+            "properly formatted."
+        )
+        error_message += "\n" + "\n".join([f"- {name}" for name in hard_failures])
+    if len(failures) > 0:
+        error_message += (
+            "The following objects docstrings do not match their signature. Run `make fix-copies` to fix this. "
+            "In some cases, this error may be raised incorrectly by the docstring checker. If you think this is the "
+            "case, you can manually check the docstrings and then add the object name to `OBJECTS_TO_IGNORE` in "
+            "`utils/check_docstrings.py`."
+        )
+        error_message += "\n" + "\n".join([f"- {name}" for name in failures])
+    if len(to_clean) > 0:
+        error_message += (
+            "The following objects docstrings contain templates you need to fix: search for `<fill_type>` or "
+            "`<fill_docstring>`."
+        )
+        error_message += "\n" + "\n".join([f"- {name}" for name in to_clean])
+
+    if len(error_message) > 0:
+        error_message = "There was at least one problem when checking docstrings of public objects.\n" + error_message
+        raise ValueError(error_message)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--fix_and_overwrite", action="store_true", help="Whether to fix inconsistencies.")
+    parser.add_argument(
+        "--check_all", action="store_true", help="Whether to check all files. By default, only checks the diff"
+    )
+    args = parser.parse_args()
+
+    check_docstrings(overwrite=args.fix_and_overwrite, check_all=args.check_all)