repo_description: 🤗 Transformers: the model-definition framework for state-of-the-art machine learning models in text, vision, audio, and multimodal models, for both inference and training. 
repo_topics: audio deep-learning deepseek gemma glm hacktoberfest llm machine-learning model-hub natural-language-processing nlp pretrained-models python pytorch pytorch-transformers qwen speech-recognition transformer vlm
code: # 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 ast
  import collections
  import contextlib
  import copy
  import doctest
  import functools
  import gc
  import importlib
  import inspect
  import json
  import logging
  import multiprocessing
  import os
  import re
  import shlex
  import shutil
  import subprocess
  import sys
  import tempfile
  import threading
  import time
  import traceback
  import types
  import unittest
  from collections import UserDict, defaultdict
  from collections.abc import Callable, Generator, Iterable, Iterator, Mapping
  from contextlib import contextmanager
  from dataclasses import MISSING, fields
  from functools import cache, wraps
  from io import StringIO
  from pathlib import Path
  from typing import TYPE_CHECKING, Any
  from unittest import mock
  from unittest.mock import patch
  
  import httpx
  from huggingface_hub import create_repo, delete_repo
  from packaging import version
  
  from transformers import logging as transformers_logging
  
  
  if TYPE_CHECKING:
      from .trainer import Trainer
  else:
      Trainer = Any  # type: ignore
  
  from .integrations import (
      is_clearml_available,
      is_optuna_available,
      is_ray_available,
      is_swanlab_available,
      is_tensorboard_available,
      is_trackio_available,
      is_wandb_available,
  )
  from .integrations.deepspeed import is_deepspeed_available
  from .utils import (
      ACCELERATE_MIN_VERSION,
      GGUF_MIN_VERSION,
      SAFE_WEIGHTS_INDEX_NAME,
      TRITON_MIN_VERSION,
      WEIGHTS_INDEX_NAME,
      is_accelerate_available,
      is_apex_available,
      is_apollo_torch_available,
      is_aqlm_available,
      is_auto_round_available,
      is_av_available,
      is_bitsandbytes_available,
      is_bs4_available,
      is_compressed_tensors_available,
      is_cv2_available,
      is_cython_available,
      is_decord_available,
      is_detectron2_available,
      is_essentia_available,
      is_faiss_available,
      is_fbgemm_gpu_available,
      is_flash_attn_2_available,
      is_flash_attn_3_available,
      is_flash_attn_4_available,
      is_flute_available,
      is_fouroversix_available,
      is_fp_quant_available,
      is_fsdp_available,
      is_g2p_en_available,
      is_galore_torch_available,
      is_gguf_available,
      is_gptqmodel_available,
      is_grokadamw_available,
      is_hadamard_available,
      is_hqq_available,
      is_huggingface_hub_greater_or_equal,
      is_ipython_available,
      is_jinja_available,
      is_jmespath_available,
      is_jumanpp_available,
      is_kernels_available,
      is_levenshtein_available,
      is_librosa_available,
      is_liger_kernel_available,
      is_lomo_available,
      is_mistral_common_available,
      is_multipart_available,
      is_natten_available,
      is_nltk_available,
      is_numba_available,
      is_onnx_available,
      is_openai_available,
      is_optimum_available,
      is_optimum_quanto_available,
      is_pandas_available,
      is_peft_available,
      is_phonemizer_available,
      is_pretty_midi_available,
      is_psutil_available,
      is_pyctcdecode_available,
      is_pytesseract_available,
      is_pytest_available,
      is_pytest_order_available,
      is_pytorch_quantization_available,
      is_quark_available,
      is_qutlass_available,
      is_rjieba_available,
      is_sacremoses_available,
      is_schedulefree_available,
      is_scipy_available,
      is_sentencepiece_available,
      is_seqio_available,
      is_serve_available,
      is_spacy_available,
      is_speech_available,
      is_spqr_available,
      is_sudachi_available,
      is_sudachi_projection_available,
      is_tiktoken_available,
      is_timm_available,
      is_tokenizers_available,
      is_torch_available,
      is_torch_bf16_available_on_device,
      is_torch_fp16_available_on_device,
      is_torch_greater_or_equal,
      is_torch_hpu_available,
      is_torch_mlu_available,
      is_torch_neuroncore_available,
      is_torch_npu_available,
      is_torch_optimi_available,
      is_torch_tensorrt_fx_available,
      is_torch_tf32_available,
      is_torch_xla_available,
      is_torch_xpu_available,
      is_torchao_available,
      is_torchaudio_available,
      is_torchcodec_available,
      is_torchvision_available,
      is_triton_available,
      is_vision_available,
      is_vptq_available,
      strtobool,
  )
  
  
  if is_accelerate_available():
      from accelerate.state import AcceleratorState, PartialState
      from accelerate.utils.imports import is_fp8_available
  
  
  if is_pytest_available():
      from _pytest.doctest import (
          Module,
          _get_checker,
          _get_continue_on_failure,
          _get_runner,
          _is_mocked,
          _patch_unwrap_mock_aware,
          get_optionflags,
      )
      from _pytest.outcomes import skip
      from _pytest.pathlib import import_path
      from pytest import DoctestItem
  else:
      Module = object
      DoctestItem = object
  
  
  SMALL_MODEL_IDENTIFIER = "julien-c/bert-xsmall-dummy"
  DUMMY_UNKNOWN_IDENTIFIER = "julien-c/dummy-unknown"
  DUMMY_DIFF_TOKENIZER_IDENTIFIER = "julien-c/dummy-diff-tokenizer"
  # Used to test Auto{Config, Model, Tokenizer} model_type detection.
  
  # Used to test the hub
  USER = "__DUMMY_TRANSFORMERS_USER__"
  ENDPOINT_STAGING = "https://hub-ci.huggingface.co"
  
  # Not critical, only usable on the sandboxed CI instance.
  TOKEN = "[REDACTED_API_KEY]"
  
  
  # Used in CausalLMModelTester (and related classes/methods) to infer the common model classes from the base model class
  _COMMON_MODEL_NAMES_MAP = {
      "config_class": "Config",
      "causal_lm_class": "ForCausalLM",
      "question_answering_class": "ForQuestionAnswering",
      "sequence_classification_class": "ForSequenceClassification",
      "token_classification_class": "ForTokenClassification",
  }
  
  # Used in VLMModelTester (and related classes/methods) to infer the common model classes from the base model class
  _VLM_COMMON_MODEL_NAMES_MAP = {
      "config_class": "Config",
      "text_config_class": "TextConfig",
      "vision_config_class": "VisionConfig",
      "conditional_generation_class": "ForConditionalGeneration",
  }
  
  
  if is_torch_available():
      import torch
      from safetensors.torch import load_file
  
      from .modeling_utils import FLASH_ATTN_KERNEL_FALLBACK, PreTrainedModel
  
      IS_ROCM_SYSTEM = torch.version.hip is not None
      IS_CUDA_SYSTEM = torch.version.cuda is not None
      IS_XPU_SYSTEM = getattr(torch.version, "xpu", None) is not None
      IS_NPU_SYSTEM = getattr(torch, "npu", None) is not None
  else:
      IS_ROCM_SYSTEM = False
      IS_CUDA_SYSTEM = False
      IS_XPU_SYSTEM = False
      IS_NPU_SYSTEM = False
  
  logger = transformers_logging.get_logger(__name__)
  
  
  def parse_flag_from_env(key, default=False):
      try:
          value = os.environ[key]
      except KeyError:
          # KEY isn't set, default to `default`.
          _value = default
      else:
          # KEY is set, convert it to True or False.
          try:
              _value = strtobool(value)
          except ValueError:
              # More values are supported, but let's keep the message simple.
              raise ValueError(f"If set, {key} must be yes or no.")
      return _value
  
  
  def parse_int_from_env(key, default=None):
      try:
          value = os.environ[key]
      except KeyError:
          _value = default
      else:
          try:
              _value = int(value)
          except ValueError:
              raise ValueError(f"If set, {key} must be a int.")
      return _value
  
  
  _run_slow_tests = parse_flag_from_env("RUN_SLOW", default=False)
  _run_flaky_tests = parse_flag_from_env("RUN_FLAKY", default=True)
  _run_custom_tokenizers = parse_flag_from_env("RUN_CUSTOM_TOKENIZERS", default=False)
  _run_staging = parse_flag_from_env("HUGGINGFACE_CO_STAGING", default=False)
  _run_pipeline_tests = parse_flag_from_env("RUN_PIPELINE_TESTS", default=True)
  _run_agent_tests = parse_flag_from_env("RUN_AGENT_TESTS", default=False)
  _run_training_tests = parse_flag_from_env("RUN_TRAINING_TESTS", default=True)
  _run_tensor_parallel_tests = parse_flag_from_env("RUN_TENSOR_PARALLEL_TESTS", default=True)
  
  
  def is_staging_test(test_case):
      """
      Decorator marking a test as a staging test.
  
      Those tests will run using the staging environment of huggingface.co instead of the real model hub.
      """
      if not _run_staging:
          return unittest.skip(reason="test is staging test")(test_case)
      else:
          try:
              import pytest  # We don't need a hard dependency on pytest in the main library
          except ImportError:
              return test_case
          else:
              return pytest.mark.is_staging_test()(test_case)
  
  
  def is_pipeline_test(test_case):
      """
      Decorator marking a test as a pipeline test. If RUN_PIPELINE_TESTS is set to a falsy value, those tests will be
      skipped.
      """
      if not _run_pipeline_tests:
          return unittest.skip(reason="test is pipeline test")(test_case)
      else:
          try:
              import pytest  # We don't need a hard dependency on pytest in the main library
          except ImportError:
              return test_case
          else:
              return pytest.mark.is_pipeline_test()(test_case)
  
  
  def is_agent_test(test_case):
      """
      Decorator marking a test as an agent test. If RUN_TOOL_TESTS is set to a falsy value, those tests will be skipped.
      """
      if not _run_agent_tests:
          return unittest.skip(reason="test is an agent test")(test_case)
      else:
          try:
              import pytest  # We don't need a hard dependency on pytest in the main library
          except ImportError:
              return test_case
          else:
              return pytest.mark.is_agent_test()(test_case)
  
  
  def is_training_test(test_case):
      """
      Decorator marking a test as a training test. If RUN_TRAINING_TESTS is set to a falsy value, those tests will be
      skipped.
      """
      if not _run_training_tests:
          return unittest.skip(reason="test is training test")(test_case)
      else:
          try:
              import pytest  # We don't need a hard dependency on pytest in the main library
          except ImportError:
              return test_case
          else:
              return pytest.mark.is_training_test()(test_case)
  
  
  def is_tensor_parallel_test(test_case):
      """
      Decorator marking a test as a tensor parallel test. If RUN_TENSOR_PARALLEL_TESTS is set to a falsy value, those
      tests will be skipped.
      """
      if not _run_tensor_parallel_tests:
          return unittest.skip(reason="test is tensor parallel test")(test_case)
      else:
          try:
              import pytest  # We don't need a hard dependency on pytest in the main library
          except ImportError:
              return test_case
          else:
              return pytest.mark.is_tensor_parallel_test()(test_case)
  
  
  def slow(test_case):
      """
      Decorator marking a test as slow.
  
      Slow tests are skipped by default. Set the RUN_SLOW environment variable to a truthy value to run them.
  
      """
      return unittest.skipUnless(_run_slow_tests, "test is slow")(test_case)
  
  
  def tooslow(test_case):
      """
      Decorator marking a test as too slow.
  
      Slow tests are skipped while they're in the process of being fixed. No test should stay tagged as "tooslow" as
      these will not be tested by the CI.
  
      """
      return unittest.skip(reason="test is too slow")(test_case)
  
  
  def skip_if_not_implemented(test_func):
      @functools.wraps(test_func)
      def wrapper(*args, **kwargs):
          try:
              return test_func(*args, **kwargs)
          except NotImplementedError as e:
              raise unittest.SkipTest(f"Test skipped due to NotImplementedError: {e}")
  
      return wrapper
  
  
  def apply_skip_if_not_implemented(cls):
      """
      Class decorator to apply @skip_if_not_implemented to all test methods.
      """
      for attr_name in dir(cls):
          if attr_name.startswith("test_"):
              attr = getattr(cls, attr_name)
              if callable(attr):
                  setattr(cls, attr_name, skip_if_not_implemented(attr))
      return cls
  
  
  def custom_tokenizers(test_case):
      """
      Decorator marking a test for a custom tokenizer.
  
      Custom tokenizers require additional dependencies, and are skipped by default. Set the RUN_CUSTOM_TOKENIZERS
      environment variable to a truthy value to run them.
      """
      return unittest.skipUnless(_run_custom_tokenizers, "test of custom tokenizers")(test_case)
  
  
  def require_bs4(test_case):
      """
      Decorator marking a test that requires BeautifulSoup4. These tests are skipped when BeautifulSoup4 isn't installed.
      """
      return unittest.skipUnless(is_bs4_available(), "test requires BeautifulSoup4")(test_case)
  
  
  def require_galore_torch(test_case):
      """
      Decorator marking a test that requires GaLore. These tests are skipped when GaLore isn't installed.
      https://github.com/jiaweizzhao/GaLore
      """
      return unittest.skipUnless(is_galore_torch_available(), "test requires GaLore")(test_case)
  
  
  def require_apollo_torch(test_case):
      """
      Decorator marking a test that requires GaLore. These tests are skipped when APOLLO isn't installed.
      https://github.com/zhuhanqing/APOLLO
      """
      return unittest.skipUnless(is_apollo_torch_available(), "test requires APOLLO")(test_case)
  
  
  def require_torch_optimi(test_case):
      """
      Decorator marking a test that requires torch-optimi. These tests are skipped when torch-optimi isn't installed.
      https://github.com/jxnl/torch-optimi
      """
      return unittest.skipUnless(is_torch_optimi_available(), "test requires torch-optimi")(test_case)
  
  
  def require_lomo(test_case):
      """
      Decorator marking a test that requires LOMO. These tests are skipped when LOMO-optim isn't installed.
      https://github.com/OpenLMLab/LOMO
      """
      return unittest.skipUnless(is_lomo_available(), "test requires LOMO")(test_case)
  
  
  def require_grokadamw(test_case):
      """
      Decorator marking a test that requires GrokAdamW. These tests are skipped when GrokAdamW isn't installed.
      """
      return unittest.skipUnless(is_grokadamw_available(), "test requires GrokAdamW")(test_case)
  
  
  def require_schedulefree(test_case):
      """
      Decorator marking a test that requires schedulefree. These tests are skipped when schedulefree isn't installed.
      https://github.com/facebookresearch/schedule_free
      """
      return unittest.skipUnless(is_schedulefree_available(), "test requires schedulefree")(test_case)
  
  
  def require_cv2(test_case):
      """
      Decorator marking a test that requires OpenCV.
  
      These tests are skipped when OpenCV isn't installed.
  
      """
      return unittest.skipUnless(is_cv2_available(), "test requires OpenCV")(test_case)
  
  
  def require_levenshtein(test_case):
      """
      Decorator marking a test that requires Levenshtein.
  
      These tests are skipped when Levenshtein isn't installed.
  
      """
      return unittest.skipUnless(is_levenshtein_available(), "test requires Levenshtein")(test_case)
  
  
  def require_nltk(test_case):
      """
      Decorator marking a test that requires NLTK.
  
      These tests are skipped when NLTK isn't installed.
  
      """
      return unittest.skipUnless(is_nltk_available(), "test requires NLTK")(test_case)
  
  
  def require_accelerate(test_case, min_version: str = ACCELERATE_MIN_VERSION):
      """
      Decorator marking a test that requires accelerate. These tests are skipped when accelerate isn't installed.
      """
      return unittest.skipUnless(
          is_accelerate_available(min_version), f"test requires accelerate version >= {min_version}"
      )(test_case)
  
  
  def require_triton(min_version: str = TRITON_MIN_VERSION):
      """
      Decorator marking a test that requires triton. These tests are skipped when triton isn't installed.
      """
  
      def decorator(test_case):
          return unittest.skipUnless(is_triton_available(min_version), f"test requires triton version >= {min_version}")(
              test_case
          )
  
      return decorator
  
  
  def require_gguf(test_case, min_version: str = GGUF_MIN_VERSION):
      """
      Decorator marking a test that requires ggguf. These tests are skipped when gguf isn't installed.
      """
      return unittest.skipUnless(is_gguf_available(min_version), f"test requires gguf version >= {min_version}")(
          test_case
      )
  
  
  def require_fsdp(test_case, min_version: str = "1.12.0"):
      """
      Decorator marking a test that requires fsdp. These tests are skipped when fsdp isn't installed.
      """
      return unittest.skipUnless(is_fsdp_available(min_version), f"test requires torch version >= {min_version}")(
          test_case
      )
  
  
  def require_g2p_en(test_case):
      """
      Decorator marking a test that requires g2p_en. These tests are skipped when SentencePiece isn't installed.
      """
      return unittest.skipUnless(is_g2p_en_available(), "test requires g2p_en")(test_case)
  
  
  def require_rjieba(test_case):
      """
      Decorator marking a test that requires rjieba. These tests are skipped when rjieba isn't installed.
      """
      return unittest.skipUnless(is_rjieba_available(), "test requires rjieba")(test_case)
  
  
  def require_jinja(test_case):
      """
      Decorator marking a test that requires jinja. These tests are skipped when jinja isn't installed.
      """
      return unittest.skipUnless(is_jinja_available(), "test requires jinja")(test_case)
  
  
  def require_jmespath(test_case):
      """
      Decorator marking a test that requires jmespath. These tests are skipped when jmespath isn't installed.
      """
      return unittest.skipUnless(is_jmespath_available(), "test requires jmespath")(test_case)
  
  
  def require_onnx(test_case):
      return unittest.skipUnless(is_onnx_available(), "test requires ONNX")(test_case)
  
  
  def require_timm(test_case):
      """
      Decorator marking a test that requires Timm.
  
      These tests are skipped when Timm isn't installed.
  
      """
      return unittest.skipUnless(is_timm_available(), "test requires Timm")(test_case)
  
  
  def require_natten(test_case):
      """
      Decorator marking a test that requires NATTEN.
  
      These tests are skipped when NATTEN isn't installed.
  
      """
      return unittest.skipUnless(is_natten_available(), "test requires natten")(test_case)
  
  
  def require_torch(test_case):
      """
      Decorator marking a test that requires PyTorch.
  
      These tests are skipped when PyTorch isn't installed.
  
      """
      return unittest.skipUnless(is_torch_available(), "test requires PyTorch")(test_case)
  
  
  def require_torch_greater_or_equal(version: str):
      """
      Decorator marking a test that requires PyTorch version >= `version`.
  
      These tests are skipped when PyTorch version is less than `version`.
      """
  
      def decorator(test_case):
          return unittest.skipUnless(is_torch_greater_or_equal(version), f"test requires PyTorch version >= {version}")(
              test_case
          )
  
      return decorator
  
  
  def require_huggingface_hub_greater_or_equal(version: str):
      """
      Decorator marking a test that requires huggingface_hub version >= `version`.
  
      These tests are skipped when huggingface_hub version is less than `version`.
      """
  
      def decorator(test_case):
          return unittest.skipUnless(
              is_huggingface_hub_greater_or_equal(version), f"test requires huggingface_hub version >= {version}"
          )(test_case)
  
      return decorator
  
  
  def require_flash_attn(test_case):
      """
      Decorator marking a test that requires Flash Attention.
  
      These tests are skipped when Flash Attention isn't installed.
  
      """
      flash_attn_available = is_flash_attn_2_available()
      kernels_available = is_kernels_available()
      try:
          from kernels import get_kernel
  
          get_kernel(FLASH_ATTN_KERNEL_FALLBACK["flash_attention_2"])
      except Exception as _:
          kernels_available = False
  
      return unittest.skipUnless(kernels_available | flash_attn_available, "test requires Flash Attention")(test_case)
  
  
  def require_kernels(test_case):
      """
      Decorator marking a test that requires the kernels library.
  
      These tests are skipped when the kernels library isn't installed.
  
      """
      return unittest.skipUnless(is_kernels_available(), "test requires the kernels library")(test_case)
  
  
  def require_flash_attn_3(test_case):
      """
      Decorator marking a test that requires Flash Attention 3.
  
      These tests are skipped when Flash Attention 3 isn't installed.
      """
      return unittest.skipUnless(is_flash_attn_3_available(), "test requires Flash Attention 3")(test_case)
  
  
  def require_flash_attn_4(test_case):
      """
      Decorator marking a test that requires Flash Attention 4.
  
      These tests are skipped when Flash Attention 4 isn't installed.
      """
      return unittest.skipUnless(is_flash_attn_4_available(), "test requires Flash Attention 4")(test_case)
  
  
  def require_all_flash_attn(test_case):
      flash_attn_available = is_flash_attn_2_available()
      kernels_available = is_kernels_available()
      try:
          from kernels import get_kernel
  
          get_kernel(FLASH_ATTN_KERNEL_FALLBACK["flash_attention_2"])
      except Exception as _:
          kernels_available = False
  
      return unittest.skipUnless(
          all(
              (
                  flash_attn_available | kernels_available,
                  is_flash_attn_3_available(),
                  is_flash_attn_4_available(),
              )
          ),
          "test requires all mainline Flash Attention packages",
      )(test_case)
  
  
  def require_peft(test_case):
      """
      Decorator marking a test that requires PEFT.
  
      These tests are skipped when PEFT isn't installed.
  
      """
      return unittest.skipUnless(is_peft_available(), "test requires PEFT")(test_case)
  
  
  def require_torchvision(test_case):
      """
      Decorator marking a test that requires Torchvision.
  
      These tests are skipped when Torchvision isn't installed.
  
      """
      return unittest.skipUnless(is_torchvision_available(), "test requires Torchvision")(test_case)
  
  
  def require_torchcodec(test_case):
      """
      Decorator marking a test that requires Torchcodec.
  
      These tests are skipped when Torchcodec isn't installed.
  
      """
      return unittest.skipUnless(is_torchcodec_available(), "test requires Torchcodec")(test_case)
  
  
  def require_torchaudio(test_case):
      """
      Decorator marking a test that requires torchaudio. These tests are skipped when torchaudio isn't installed.
      """
      return unittest.skipUnless(is_torchaudio_available(), "test requires torchaudio")(test_case)
  
  
  def require_sentencepiece(test_case):
      """
      Decorator marking a test that requires SentencePiece. These tests are skipped when SentencePiece isn't installed.
      """
      return unittest.skipUnless(is_sentencepiece_available(), "test requires SentencePiece")(test_case)
  
  
  def require_sacremoses(test_case):
      """
      Decorator marking a test that requires Sacremoses. These tests are skipped when Sacremoses isn't installed.
      """
      return unittest.skipUnless(is_sacremoses_available(), "test requires Sacremoses")(test_case)
  
  
  def require_seqio(test_case):
      """
      Decorator marking a test that requires SentencePiece. These tests are skipped when SentencePiece isn't installed.
      """
      return unittest.skipUnless(is_seqio_available(), "test requires Seqio")(test_case)
  
  
  def require_scipy(test_case):
      """
      Decorator marking a test that requires Scipy. These tests are skipped when SentencePiece isn't installed.
      """
      return unittest.skipUnless(is_scipy_available(), "test requires Scipy")(test_case)
  
  
  def require_tokenizers(test_case):
      """
      Decorator marking a test that requires 🤗 Tokenizers. These tests are skipped when 🤗 Tokenizers isn't installed.
      """
      return unittest.skipUnless(is_tokenizers_available(), "test requires tokenizers")(test_case)
  
  
  def require_pandas(test_case):
      """
      Decorator marking a test that requires pandas. These tests are skipped when pandas isn't installed.
      """
      return unittest.skipUnless(is_pandas_available(), "test requires pandas")(test_case)
  
  
  def require_pytesseract(test_case):
      """
      Decorator marking a test that requires PyTesseract. These tests are skipped when PyTesseract isn't installed.
      """
      return unittest.skipUnless(is_pytesseract_available(), "test requires PyTesseract")(test_case)
  
  
  def require_pytorch_quantization(test_case):
      """
      Decorator marking a test that requires PyTorch Quantization Toolkit. These tests are skipped when PyTorch
      Quantization Toolkit isn't installed.
      """
      return unittest.skipUnless(is_pytorch_quantization_available(), "test requires PyTorch Quantization Toolkit")(
          test_case
      )
  
  
  def require_vision(test_case):
      """
      Decorator marking a test that requires the vision dependencies. These tests are skipped when torchaudio isn't
      installed.
      """
      return unittest.skipUnless(is_vision_available(), "test requires vision")(test_case)
  
  
  def require_spacy(test_case):
      """
      Decorator marking a test that requires SpaCy. These tests are skipped when SpaCy isn't installed.
      """
      return unittest.skipUnless(is_spacy_available(), "test requires spacy")(test_case)
  
  
  def require_torch_multi_gpu(test_case):
      """
      Decorator marking a test that requires a multi-GPU CUDA setup (in PyTorch). These tests are skipped on a machine without
      multiple CUDA GPUs.
  
      To run *only* the multi_gpu tests, assuming all test names contain multi_gpu: $ pytest -sv ./tests -k "multi_gpu"
      """
      if not is_torch_available():
          return unittest.skip(reason="test requires PyTorch")(test_case)
  
      import torch
  
      return unittest.skipUnless(torch.cuda.device_count() > 1, "test requires multiple CUDA GPUs")(test_case)
  
  
  def require_torch_multi_accelerator(test_case):
      """
      Decorator marking a test that requires a multi-accelerator (in PyTorch). These tests are skipped on a machine
      without multiple accelerators. To run *only* the multi_accelerator tests, assuming all test names contain
      multi_accelerator: $ pytest -sv ./tests -k "multi_accelerator"
      """
      if not is_torch_available():
          return unittest.skip(reason="test requires PyTorch")(test_case)
  
      return unittest.skipUnless(backend_device_count(torch_device) > 1, "test requires multiple accelerators")(
          test_case
      )
  
  
  def require_torch_non_multi_gpu(test_case):
      """
      Decorator marking a test that requires 0 or 1 GPU setup (in PyTorch).
      """
      if not is_torch_available():
          return unittest.skip(reason="test requires PyTorch")(test_case)
  
      import torch
  
      return unittest.skipUnless(torch.cuda.device_count() < 2, "test requires 0 or 1 GPU")(test_case)
  
  
  def require_torch_non_multi_accelerator(test_case):
      """
      Decorator marking a test that requires 0 or 1 accelerator setup (in PyTorch).
      """
      if not is_torch_available():
          return unittest.skip(reason="test requires PyTorch")(test_case)
  
      return unittest.skipUnless(backend_device_count(torch_device) < 2, "test requires 0 or 1 accelerator")(test_case)
  
  
  def require_torch_up_to_2_gpus(test_case):
      """
      Decorator marking a test that requires 0 or 1 or 2 GPU setup (in PyTorch).
      """
      if not is_torch_available():
          return unittest.skip(reason="test requires PyTorch")(test_case)
  
      import torch
  
      return unittest.skipUnless(torch.cuda.device_count() < 3, "test requires 0 or 1 or 2 GPUs")(test_case)
  
  
  def require_torch_up_to_2_accelerators(test_case):
      """
      Decorator marking a test that requires 0 or 1 or 2 accelerator setup (in PyTorch).
      """
      if not is_torch_available():
          return unittest.skip(reason="test requires PyTorch")(test_case)
  
      return unittest.skipUnless(backend_device_count(torch_device) < 3, "test requires 0 or 1 or 2 accelerators")(
          test_case
      )
  
  
  def require_torch_xla(test_case):
      """
      Decorator marking a test that requires TorchXLA (in PyTorch).
      """
      return unittest.skipUnless(is_torch_xla_available(), "test requires TorchXLA")(test_case)
  
  
  def require_torch_neuroncore(test_case):
      """
      Decorator marking a test that requires NeuronCore (in PyTorch).
      """
      return unittest.skipUnless(is_torch_neuroncore_available(check_device=False), "test requires PyTorch NeuronCore")(
          test_case
      )
  
  
  def require_torch_npu(test_case):
      """
      Decorator marking a test that requires NPU (in PyTorch).
      """
      return unittest.skipUnless(is_torch_npu_available(), "test requires PyTorch NPU")(test_case)
  
  
  def require_torch_multi_npu(test_case):
      """
      Decorator marking a test that requires a multi-NPU setup (in PyTorch). These tests are skipped on a machine without
      multiple NPUs.
  
      To run *only* the multi_npu tests, assuming all test names contain multi_npu: $ pytest -sv ./tests -k "multi_npu"
      """
      if not is_torch_npu_available():
          return unittest.skip(reason="test requires PyTorch NPU")(test_case)
  
      return unittest.skipUnless(torch.npu.device_count() > 1, "test requires multiple NPUs")(test_case)
  
  
  def require_non_hpu(test_case):
      """
      Decorator marking a test that should be skipped for HPU.
      """
      return unittest.skipUnless(torch_device != "hpu", "test requires a non-HPU")(test_case)
  
  
  def require_torch_xpu(test_case):
      """
      Decorator marking a test that requires XPU (in PyTorch).
  
      These tests are skipped when XPU backend is not available.
      """
      return unittest.skipUnless(is_torch_xpu_available(), "test requires XPU device")(test_case)
  
  
  def require_non_xpu(test_case):
      """
      Decorator marking a test that should be skipped for XPU.
      """
      return unittest.skipUnless(torch_device != "xpu", "test requires a non-XPU")(test_case)
  
  
  def require_torch_multi_xpu(test_case):
      """
      Decorator marking a test that requires a multi-XPU setup (in PyTorch). These tests are skipped on a machine without
      multiple XPUs.
  
      To run *only* the multi_xpu tests, assuming all test names contain multi_xpu: $ pytest -sv ./tests -k "multi_xpu"
      """
      if not is_torch_xpu_available():
          return unittest.skip(reason="test requires PyTorch XPU")(test_case)
  
      return unittest.skipUnless(torch.xpu.device_count() > 1, "test requires multiple XPUs")(test_case)
  
  
  def require_torch_multi_hpu(test_case):
      """
      Decorator marking a test that requires a multi-HPU setup (in PyTorch). These tests are skipped on a machine without
      multiple HPUs.
  
      To run *only* the multi_hpu tests, assuming all test names contain multi_hpu: $ pytest -sv ./tests -k "multi_hpu"
      """
      if not is_torch_hpu_available():
          return unittest.skip(reason="test requires PyTorch HPU")(test_case)
  
      return unittest.skipUnless(torch.hpu.device_count() > 1, "test requires multiple HPUs")(test_case)
  
  
  if is_torch_available():
      # Set env var CUDA_VISIBLE_DEVICES="" to force cpu-mode
      import torch
  
      if "TRANSFORMERS_TEST_BACKEND" in os.environ:
          backend = os.environ["TRANSFORMERS_TEST_BACKEND"]
          try:
              _ = importlib.import_module(backend)
          except ModuleNotFoundError as e:
              raise ModuleNotFoundError(
                  f"Failed to import `TRANSFORMERS_TEST_BACKEND` '{backend}'! This should be the name of an installed module. The original error (look up to see its"
                  f" traceback):\n{e}"
              ) from e
  
      if "TRANSFORMERS_TEST_DEVICE" in os.environ:
          torch_device = os.environ["TRANSFORMERS_TEST_DEVICE"]
          if torch_device == "cuda" and not torch.cuda.is_available():
              raise ValueError(
                  f"TRANSFORMERS_TEST_DEVICE={torch_device}, but CUDA is unavailable. Please double-check your testing environment."
              )
          if torch_device == "xpu" and not is_torch_xpu_available():
              raise ValueError(
                  f"TRANSFORMERS_TEST_DEVICE={torch_device}, but XPU is unavailable. Please double-check your testing environment."
              )
          if torch_device == "npu" and not is_torch_npu_available():
              raise ValueError(
                  f"TRANSFORMERS_TEST_DEVICE={torch_device}, but NPU is unavailable. Please double-check your testing environment."
              )
          if torch_device == "mlu" and not is_torch_mlu_available():
              raise ValueError(
                  f"TRANSFORMERS_TEST_DEVICE={torch_device}, but MLU is unavailable. Please double-check your testing environment."
              )
          if torch_device == "hpu" and not is_torch_hpu_available():
              raise ValueError(
                  f"TRANSFORMERS_TEST_DEVICE={torch_device}, but HPU is unavailable. Please double-check your testing environment."
              )
  
          try:
              # try creating device to see if provided device is valid
              _ = torch.device(torch_device)
          except RuntimeError as e:
              raise RuntimeError(
                  f"Unknown testing device specified by environment variable `TRANSFORMERS_TEST_DEVICE`: {torch_device}"
              ) from e
      elif torch.cuda.is_available():
          torch_device = "cuda"
      elif is_torch_npu_available():
          torch_device = "npu"
      elif is_torch_mlu_available():
          torch_device = "mlu"
      elif is_torch_hpu_available():
          torch_device = "hpu"
      elif is_torch_xpu_available():
          torch_device = "xpu"
      else:
          torch_device = "cpu"
  else:
      torch_device = None
  
  
  def require_torchao(test_case):
      """Decorator marking a test that requires torchao"""
      return unittest.skipUnless(is_torchao_available(), "test requires torchao")(test_case)
  
  
  def require_torchao_version_greater_or_equal(torchao_version):
      def decorator(test_case):
          correct_torchao_version = is_torchao_available() and version.parse(
              version.parse(importlib.metadata.version("torchao")).base_version
          ) >= version.parse(torchao_version)
          return unittest.skipUnless(
              correct_torchao_version, f"Test requires torchao with the version greater than {torchao_version}."
          )(test_case)
  
      return decorator
  
  
  def require_torch_tensorrt_fx(test_case):
      """Decorator marking a test that requires Torch-TensorRT FX"""
      return unittest.skipUnless(is_torch_tensorrt_fx_available(), "test requires Torch-TensorRT FX")(test_case)
  
  
  def require_torch_gpu(test_case):
      """Decorator marking a test that requires CUDA and PyTorch."""
      return unittest.skipUnless(torch_device == "cuda", "test requires CUDA")(test_case)
  
  
  def require_torch_mps(test_case):
      """Decorator marking a test that requires CUDA and PyTorch."""
      return unittest.skipUnless(torch_device == "mps", "test requires MPS")(test_case)
  
  
  def require_large_cpu_ram(test_case, memory: float = 80):
      """Decorator marking a test that requires a CPU RAM with more than `memory` GiB of memory."""
      if not is_psutil_available():
          return test_case
  
      import psutil
  
      return unittest.skipUnless(
          psutil.virtual_memory().total / 1024**3 > memory,
          f"test requires a machine with more than {memory} GiB of CPU RAM memory",
      )(test_case)
  
  
  def require_torch_large_gpu(test_case, memory: float = 20):
      """Decorator marking a test that requires a CUDA GPU with more than `memory` GiB of memory."""
      if torch_device != "cuda":
          return unittest.skip(reason=f"test requires a CUDA GPU with more than {memory} GiB of memory")(test_case)
  
      return unittest.skipUnless(
          torch.cuda.get_device_properties(0).total_memory / 1024**3 > memory,
          f"test requires a GPU with more than {memory} GiB of memory",
      )(test_case)
  
  
  def require_torch_large_accelerator(test_case=None, *, memory: float = 20):
      """Decorator marking a test that requires an accelerator with more than `memory` GiB of memory."""
  
      def memory_decorator(tc):
          if torch_device not in ("cuda", "xpu"):
              return unittest.skip(f"test requires a GPU or XPU with more than {memory} GiB of memory")(tc)
  
          torch_accel = getattr(torch, torch_device)
          return unittest.skipUnless(
              torch_accel.get_device_properties(0).total_memory / 1024**3 > memory,
              f"test requires a GPU or XPU with more than {memory} GiB of memory",
          )(tc)
  
      return memory_decorator if test_case is None else memory_decorator(test_case)
  
  
  def require_torch_accelerator(test_case):
      """Decorator marking a test that requires an accessible accelerator and PyTorch."""
      return unittest.skipUnless(torch_device is not None and torch_device != "cpu", "test requires accelerator")(
          test_case
      )
  
  
  def require_torch_fp16(test_case):
      """Decorator marking a test that requires a device that supports fp16"""
      return unittest.skipUnless(
          is_torch_fp16_available_on_device(torch_device), "test requires device with fp16 support"
      )(test_case)
  
  
  def require_fp8(test_case):
      """Decorator marking a test that requires supports for fp8"""
      return unittest.skipUnless(is_accelerate_available() and is_fp8_available(), "test requires fp8 support")(
          test_case
      )
  
  
  def require_cuda_capability_at_least(major, minor):
      """Decorator to skip tests when CUDA capability is below the given version."""
      import torch
  
      if not torch.cuda.is_available():
          return unittest.skip("CUDA not available")
      capability = torch.cuda.get_device_capability()
      return unittest.skipIf(capability < (major, minor), f"Requires CUDA capability >= {major}.{minor}")
  
  
  def require_torch_bf16(test_case):
      """Decorator marking a test that requires a device that supports bf16"""
      return unittest.skipUnless(
          is_torch_bf16_available_on_device(torch_device), "test requires device with bf16 support"
      )(test_case)
  
  
  def require_deterministic_for_xpu(test_case):
      @wraps(test_case)
      def wrapper(*args, **kwargs):
          if is_torch_xpu_available():
              original_state = torch.are_deterministic_algorithms_enabled()
              try:
                  torch.use_deterministic_algorithms(True)
                  return test_case(*args, **kwargs)
              finally:
                  torch.use_deterministic_algorithms(original_state)
          else:
              return test_case(*args, **kwargs)
  
      return wrapper
  
  
  def require_torch_tf32(test_case):
      """Decorator marking a test that requires Ampere or a newer GPU arch, cuda>=11 and torch>=1.7."""
      return unittest.skipUnless(
          is_torch_tf32_available(), "test requires Ampere or a newer GPU arch, cuda>=11 and torch>=1.7"
      )(test_case)
  
  
  def require_detectron2(test_case):
      """Decorator marking a test that requires detectron2."""
      return unittest.skipUnless(is_detectron2_available(), "test requires `detectron2`")(test_case)
  
  
  def require_faiss(test_case):
      """Decorator marking a test that requires faiss."""
      return unittest.skipUnless(is_faiss_available(), "test requires `faiss`")(test_case)
  
  
  def require_ipython(test_case):
      """Decorator marking a test that requires IPython. These tests are skipped when IPython isn't installed."""
      return unittest.skipUnless(is_ipython_available(), "test requires `IPython`")(test_case)
  
  
  def require_optuna(test_case):
      """
      Decorator marking a test that requires optuna.
  
      These tests are skipped when optuna isn't installed.
  
      """
      return unittest.skipUnless(is_optuna_available(), "test requires optuna")(test_case)
  
  
  def require_ray(test_case):
      """
      Decorator marking a test that requires Ray/tune.
  
      These tests are skipped when Ray/tune isn't installed.
  
      """
      return unittest.skipUnless(is_ray_available(), "test requires Ray/tune")(test_case)
  
  
  def require_swanlab(test_case):
      """
      Decorator marking a test that requires swanlab.
  
      These tests are skipped when swanlab isn't installed.
  
      """
      return unittest.skipUnless(is_swanlab_available(), "test requires swanlab")(test_case)
  
  
  def require_trackio(test_case):
      """
      Decorator marking a test that requires trackio.
  
      These tests are skipped when trackio isn't installed.
  
      """
      return unittest.skipUnless(is_trackio_available(), "test requires trackio")(test_case)
  
  
  def require_wandb(test_case):
      """
      Decorator marking a test that requires wandb.
  
      These tests are skipped when wandb isn't installed.
  
      """
      return unittest.skipUnless(is_wandb_available(), "test requires wandb")(test_case)
  
  
  def require_clearml(test_case):
      """
      Decorator marking a test requires clearml.
  
      These tests are skipped when clearml isn't installed.
  
      """
      return unittest.skipUnless(is_clearml_available(), "test requires clearml")(test_case)
  
  
  def require_deepspeed(test_case):
      """
      Decorator marking a test that requires deepspeed
      """
      return unittest.skipUnless(is_deepspeed_available(), "test requires deepspeed")(test_case)
  
  
  def require_apex(test_case):
      """
      Decorator marking a test that requires apex
      """
      return unittest.skipUnless(is_apex_available(), "test requires apex")(test_case)
  
  
  def require_aqlm(test_case):
      """
      Decorator marking a test that requires aqlm
      """
      return unittest.skipUnless(is_aqlm_available(), "test requires aqlm")(test_case)
  
  
  def require_vptq(test_case):
      """
      Decorator marking a test that requires vptq
      """
      return unittest.skipUnless(is_vptq_available(), "test requires vptq")(test_case)
  
  
  def require_spqr(test_case):
      """
      Decorator marking a test that requires spqr
      """
      return unittest.skipUnless(is_spqr_available(), "test requires spqr")(test_case)
  
  
  def require_av(test_case):
      """
      Decorator marking a test that requires av
      """
      return unittest.skipUnless(is_av_available(), "test requires av")(test_case)
  
  
  def require_decord(test_case):
      """
      Decorator marking a test that requires decord
      """
      return unittest.skipUnless(is_decord_available(), "test requires decord")(test_case)
  
  
  def require_bitsandbytes(test_case):
      """
      Decorator marking a test that requires the bitsandbytes library. Will be skipped when the library or its hard dependency torch is not installed.
      """
      return unittest.skipUnless(is_bitsandbytes_available(), "test requires bitsandbytes")(test_case)
  
  
  def require_optimum(test_case):
      """
      Decorator for optimum dependency
      """
      return unittest.skipUnless(is_optimum_available(), "test requires optimum")(test_case)
  
  
  def require_tensorboard(test_case):
      """
      Decorator for `tensorboard` dependency
      """
      return unittest.skipUnless(is_tensorboard_available(), "test requires tensorboard")
  
  
  def require_gptqmodel(test_case):
      """
      Decorator for gptqmodel dependency
      """
      return unittest.skipUnless(is_gptqmodel_available(), "test requires gptqmodel")(test_case)
  
  
  def require_hqq(test_case):
      """
      Decorator for hqq dependency
      """
      return unittest.skipUnless(is_hqq_available(), "test requires hqq")(test_case)
  
  
  def require_auto_round(test_case):
      """
      Decorator for auto_round dependency
      """
      return unittest.skipUnless(is_auto_round_available(), "test requires autoround")(test_case)
  
  
  def require_optimum_quanto(test_case):
      """
      Decorator for quanto dependency
      """
      return unittest.skipUnless(is_optimum_quanto_available(), "test requires optimum-quanto")(test_case)
  
  
  def require_compressed_tensors(test_case):
      """
      Decorator for compressed_tensors dependency
      """
      return unittest.skipUnless(is_compressed_tensors_available(), "test requires compressed_tensors")(test_case)
  
  
  def require_fbgemm_gpu(test_case):
      """
      Decorator for fbgemm_gpu dependency
      """
      return unittest.skipUnless(is_fbgemm_gpu_available(), "test requires fbgemm-gpu")(test_case)
  
  
  def require_quark(test_case):
      """
      Decorator for quark dependency
      """
      return unittest.skipUnless(is_quark_available(), "test requires quark")(test_case)
  
  
  def require_flute_hadamard(test_case):
      """
      Decorator marking a test that requires higgs and hadamard
      """
      return unittest.skipUnless(
          is_flute_available() and is_hadamard_available(), "test requires flute and fast_hadamard_transform"
      )(test_case)
  
  
  def require_fouroversix(test_case):
      """
      Decorator marking a test that requires fouroversix
      """
      return unittest.skipUnless(is_fouroversix_available(), "test requires fouroversix")(test_case)
  
  
  def require_fp_quant(test_case):
      """
      Decorator marking a test that requires fp_quant and qutlass
      """
      return unittest.skipUnless(is_fp_quant_available(), "test requires fp_quant")(test_case)
  
  
  def require_qutlass(test_case):
      """
      Decorator marking a test that requires qutlass
      """
      return unittest.skipUnless(is_qutlass_available(), "test requires qutlass")(test_case)
  
  
  def require_phonemizer(test_case):
      """
      Decorator marking a test that requires phonemizer
      """
      return unittest.skipUnless(is_phonemizer_available(), "test requires phonemizer")(test_case)
  
  
  def require_pyctcdecode(test_case):
      """
      Decorator marking a test that requires pyctcdecode
      """
      return unittest.skipUnless(is_pyctcdecode_available(), "test requires pyctcdecode")(test_case)
  
  
  def require_numba(test_case):
      """
      Decorator marking a test that requires numba
      """
      return unittest.skipUnless(is_numba_available(), "test requires numba")(test_case)
  
  
  def require_librosa(test_case):
      """
      Decorator marking a test that requires librosa
      """
      return unittest.skipUnless(is_librosa_available(), "test requires librosa")(test_case)
  
  
  def require_multipart(test_case):
      """
      Decorator marking a test that requires python-multipart
      """
      return unittest.skipUnless(is_multipart_available(), "test requires python-multipart")(test_case)
  
  
  def require_liger_kernel(test_case):
      """
      Decorator marking a test that requires liger_kernel
      """
      return unittest.skipUnless(is_liger_kernel_available(), "test requires liger_kernel")(test_case)
  
  
  def require_essentia(test_case):
      """
      Decorator marking a test that requires essentia
      """
      return unittest.skipUnless(is_essentia_available(), "test requires essentia")(test_case)
  
  
  def require_pretty_midi(test_case):
      """
      Decorator marking a test that requires pretty_midi
      """
      return unittest.skipUnless(is_pretty_midi_available(), "test requires pretty_midi")(test_case)
  
  
  def cmd_exists(cmd):
      return shutil.which(cmd) is not None
  
  
  def require_usr_bin_time(test_case):
      """
      Decorator marking a test that requires `/usr/bin/time`
      """
      return unittest.skipUnless(cmd_exists("/usr/bin/time"), "test requires /usr/bin/time")(test_case)
  
  
  def require_sudachi(test_case):
      """
      Decorator marking a test that requires sudachi
      """
      return unittest.skipUnless(is_sudachi_available(), "test requires sudachi")(test_case)
  
  
  def require_sudachi_projection(test_case):
      """
      Decorator marking a test that requires sudachi_projection
      """
      return unittest.skipUnless(is_sudachi_projection_available(), "test requires sudachi which supports projection")(
          test_case
      )
  
  
  def require_jumanpp(test_case):
      """
      Decorator marking a test that requires jumanpp
      """
      return unittest.skipUnless(is_jumanpp_available(), "test requires jumanpp")(test_case)
  
  
  def require_cython(test_case):
      """
      Decorator marking a test that requires jumanpp
      """
      return unittest.skipUnless(is_cython_available(), "test requires cython")(test_case)
  
  
  def require_tiktoken(test_case):
      """
      Decorator marking a test that requires TikToken. These tests are skipped when TikToken isn't installed.
      """
      return unittest.skipUnless(is_tiktoken_available(), "test requires TikToken")(test_case)
  
  
  def require_speech(test_case):
      """
      Decorator marking a test that requires speech. These tests are skipped when speech isn't available.
      """
      return unittest.skipUnless(is_speech_available(), "test requires torchaudio")(test_case)
  
  
  def require_openai(test_case):
      """
      Decorator marking a test that requires openai
      """
      return unittest.skipUnless(is_openai_available(), "test requires openai")(test_case)
  
  
  def require_serve(test_case):
      """
      Decorator marking a test that requires the serving dependencies (fastapi, uvicorn, pydantic, openai).
      """
      return unittest.skipUnless(is_serve_available(), "test requires serving dependencies")(test_case)
  
  
  def require_mistral_common(test_case):
      """
      Decorator marking a test that requires mistral-common. These tests are skipped when mistral-common isn't available.
      """
      return unittest.skipUnless(is_mistral_common_available(), "test requires mistral-common")(test_case)
  
  
  def get_gpu_count():
      """
      Return the number of available gpus
      """
      if is_torch_available():
          import torch
  
          return torch.cuda.device_count()
      else:
          return 0
  
  
  def get_tests_dir(append_path=None):
      """
      Args:
          append_path: optional path to append to the tests dir path
  
      Return:
          The full path to the `tests` dir, so that the tests can be invoked from anywhere. Optionally `append_path` is
          joined after the `tests` dir the former is provided.
  
      """
      # this function caller's __file__
      caller__file__ = inspect.stack()[1][1]
      tests_dir = os.path.abspath(os.path.dirname(caller__file__))
  
      while not tests_dir.endswith("tests"):
          tests_dir = os.path.dirname(tests_dir)
  
      if append_path:
          return os.path.join(tests_dir, append_path)
      else:
          return tests_dir
  
  
  def get_steps_per_epoch(trainer: Trainer) -> int:
      training_args = trainer.args
      train_dataloader = trainer.get_train_dataloader()
  
      initial_training_values = trainer.set_initial_training_values(args=training_args, dataloader=train_dataloader)
      steps_per_epoch = initial_training_values[5]
  
      return steps_per_epoch
  
  
  def evaluate_side_effect_factory(
      side_effect_values: list[dict[str, float]],
  ) -> Generator[dict[str, float], None, None]:
      """
      Function that returns side effects for the _evaluate method.
      Used when we're unsure of exactly how many times _evaluate will be called.
      """
      yield from side_effect_values
  
      while True:
          yield side_effect_values[-1]
  
  
  #
  # Helper functions for dealing with testing text outputs
  # The original code came from:
  # https://github.com/fastai/fastai/blob/master/tests/utils/text.py
  
  
  # When any function contains print() calls that get overwritten, like progress bars,
  # a special care needs to be applied, since under pytest -s captured output (capsys
  # or contextlib.redirect_stdout) contains any temporary printed strings, followed by
  # \r's. This helper function ensures that the buffer will contain the same output
  # with and without -s in pytest, by turning:
  # foo bar\r tar mar\r final message
  # into:
  # final message
  # it can handle a single string or a multiline buffer
  def apply_print_resets(buf):
      return re.sub(r"^.*\r", "", buf, 0, re.MULTILINE)
  
  
  def assert_screenout(out, what):
      out_pr = apply_print_resets(out).lower()
      match_str = out_pr.find(what.lower())
      assert match_str != -1, f"expecting to find {what} in output: f{out_pr}"
  
  
  def set_config_for_less_flaky_test(config):
      target_attrs = [
          "rms_norm_eps",
          "layer_norm_eps",
          "norm_eps",
          "norm_epsilon",
          "layer_norm_epsilon",
          "batch_norm_eps",
      ]
      for target_attr in target_attrs:
          setattr(config, target_attr, 1.0)
  
      # norm layers (layer/group norm, etc.) could cause flaky tests when the tensors have very small variance.
      # (We don't need the original epsilon values to check eager/sdpa matches)
      attrs = ["text_config", "vision_config", "audio_config", "text_encoder", "audio_encoder", "decoder"]
      for attr in attrs:
          if hasattr(config, attr) and getattr(config, attr) is not None:
              for target_attr in target_attrs:
                  setattr(getattr(config, attr), target_attr, 1.0)
  
  
  def set_model_for_less_flaky_test(model):
      # Another way to make sure norm layers have desired epsilon. (Some models don't set it from its config.)
      target_names = (
          "LayerNorm",
          "GroupNorm",
          "BatchNorm",
          "RMSNorm",
          "BatchNorm2d",
          "BatchNorm1d",
          "BitGroupNormActivation",
          "WeightStandardizedConv2d",
      )
      target_attrs = ["eps", "epsilon", "variance_epsilon"]
      if is_torch_available() and isinstance(model, torch.nn.Module):
          for module in model.modules():
              if type(module).__name__.endswith(target_names):
                  for attr in target_attrs:
                      if hasattr(module, attr):
                          setattr(module, attr, 1.0)
  
  
  class CaptureStd:
      """
      Context manager to capture:
  
          - stdout: replay it, clean it up and make it available via `obj.out`
          - stderr: replay it and make it available via `obj.err`
  
      Args:
          out (`bool`, *optional*, defaults to `True`): Whether to capture stdout or not.
          err (`bool`, *optional*, defaults to `True`): Whether to capture stderr or not.
          replay (`bool`, *optional*, defaults to `True`): Whether to replay or not.
              By default each captured stream gets replayed back on context's exit, so that one can see what the test was
              doing. If this is a not wanted behavior and the captured data shouldn't be replayed, pass `replay=False` to
              disable this feature.
  
      Examples:
  
      ```python
      # to capture stdout only with auto-replay
      with CaptureStdout() as cs:
          print("Secret message")
      assert "message" in cs.out
  
      # to capture stderr only with auto-replay
      import sys
  
      with CaptureStderr() as cs:
          print("Warning: ", file=sys.stderr)
      assert "Warning" in cs.err
  
      # to capture both streams with auto-replay
      with CaptureStd() as cs:
          print("Secret message")
          print("Warning: ", file=sys.stderr)
      assert "message" in cs.out
      assert "Warning" in cs.err
  
      # to capture just one of the streams, and not the other, with auto-replay
      with CaptureStd(err=False) as cs:
          print("Secret message")
      assert "message" in cs.out
      # but best use the stream-specific subclasses
  
      # to capture without auto-replay
      with CaptureStd(replay=False) as cs:
          print("Secret message")
      assert "message" in cs.out
      ```"""
  
      def __init__(self, out=True, err=True, replay=True):
          self.replay = replay
  
          if out:
              self.out_buf = StringIO()
              self.out = "error: CaptureStd context is unfinished yet, called too early"
          else:
              self.out_buf = None
              self.out = "not capturing stdout"
  
          if err:
              self.err_buf = StringIO()
              self.err = "error: CaptureStd context is unfinished yet, called too early"
          else:
              self.err_buf = None
              self.err = "not capturing stderr"
  
      def __enter__(self):
          if self.out_buf:
              self.out_old = sys.stdout
              sys.stdout = self.out_buf
  
          if self.err_buf:
              self.err_old = sys.stderr
              sys.stderr = self.err_buf
  
          return self
  
      def __exit__(self, *exc):
          if self.out_buf:
              sys.stdout = self.out_old
              captured = self.out_buf.getvalue()
              if self.replay:
                  sys.stdout.write(captured)
              self.out = apply_print_resets(captured)
  
          if self.err_buf:
              sys.stderr = self.err_old
              captured = self.err_buf.getvalue()
              if self.replay:
                  sys.stderr.write(captured)
              self.err = captured
  
      def __repr__(self):
          msg = ""
          if self.out_buf:
              msg += f"stdout: {self.out}\n"
          if self.err_buf:
              msg += f"stderr: {self.err}\n"
          return msg
  
  
  # in tests it's the best to capture only the stream that's wanted, otherwise
  # it's easy to miss things, so unless you need to capture both streams, use the
  # subclasses below (less typing). Or alternatively, configure `CaptureStd` to
  # disable the stream you don't need to test.
  
  
  class CaptureStdout(CaptureStd):
      """Same as CaptureStd but captures only stdout"""
  
      def __init__(self, replay=True):
          super().__init__(err=False, replay=replay)
  
  
  class CaptureStderr(CaptureStd):
      """Same as CaptureStd but captures only stderr"""
  
      def __init__(self, replay=True):
          super().__init__(out=False, replay=replay)
  
  
  class CaptureLogger:
      """
      Context manager to capture `logging` streams
  
      Args:
          logger: 'logging` logger object
  
      Returns:
          The captured output is available via `self.out`
  
      Example:
  
      ```python
      >>> from transformers import logging
      >>> from transformers.testing_utils import CaptureLogger
  
      >>> msg = "Testing 1, 2, 3"
      >>> logging.set_verbosity_info()
      >>> logger = logging.get_logger("transformers.models.bart.tokenization_bart")
      >>> with CaptureLogger(logger) as cl:
      ...     logger.info(msg)
      >>> assert cl.out, msg + "\n"
      ```
      """
  
      def __init__(self, logger):
          self.logger = logger
          self.io = StringIO()
          self.sh = logging.StreamHandler(self.io)
          self.out = ""
  
      def __enter__(self):
          self.logger.addHandler(self.sh)
          return self
  
      def __exit__(self, *exc):
          self.logger.removeHandler(self.sh)
          self.out = self.io.getvalue()
  
      def __repr__(self):
          return f"captured: {self.out}\n"
  
  
  @contextlib.contextmanager
  def LoggingLevel(level):
      """
      This is a context manager to temporarily change transformers modules logging level to the desired value and have it
      restored to the original setting at the end of the scope.
  
      Example:
  
      ```python
      with LoggingLevel(logging.INFO):
          AutoModel.from_pretrained("openai-community/gpt2")  # calls logger.info() several times
      ```
      """
      orig_level = transformers_logging.get_verbosity()
      try:
          transformers_logging.set_verbosity(level)
          yield
      finally:
          transformers_logging.set_verbosity(orig_level)
  
  
  class TemporaryHubRepo:
      """Create a temporary Hub repository and return its `RepoUrl` object. This is similar to
      `tempfile.TemporaryDirectory` and can be used as a context manager. For example:
  
          with TemporaryHubRepo(token=self._token) as temp_repo:
              ...
  
      Upon exiting the context, the repository and everything contained in it are removed.
  
      Example:
  
      ```python
      with TemporaryHubRepo(token=self._token) as temp_repo:
          model.push_to_hub(tmp_repo.repo_id, token=self._token)
      ```
      """
  
      def __init__(self, namespace: str | None = None, token: str | None = None) -> None:
          self.token = token
          with tempfile.TemporaryDirectory() as tmp_dir:
              repo_id = Path(tmp_dir).name
              if namespace is not None:
                  repo_id = f"{namespace}/{repo_id}"
              self.repo_url = create_repo(repo_id, token=self.token)
  
      def __enter__(self):
          return self.repo_url
  
      def __exit__(self, exc, value, tb):
          delete_repo(repo_id=self.repo_url.repo_id, token=self.token, missing_ok=True)
  
  
  @contextlib.contextmanager
  # adapted from https://stackoverflow.com/a/64789046/9201239
  def ExtendSysPath(path: str | os.PathLike) -> Iterator[None]:
      """
      Temporary add given path to `sys.path`.
  
      Usage :
  
      ```python
      with ExtendSysPath("/path/to/dir"):
          mymodule = importlib.import_module("mymodule")
      ```
      """
  
      path = os.fspath(path)
      try:
          sys.path.insert(0, path)
          yield
      finally:
          sys.path.remove(path)
  
  
  class TestCasePlus(unittest.TestCase):
      """
      This class extends *unittest.TestCase* with additional features.
  
      Feature 1: A set of fully resolved important file and dir path accessors.
  
      In tests often we need to know where things are relative to the current test file, and it's not trivial since the
      test could be invoked from more than one directory or could reside in sub-directories with different depths. This
      class solves this problem by sorting out all the basic paths and provides easy accessors to them:
  
      - `pathlib` objects (all fully resolved):
  
         - `test_file_path` - the current test file path (=`__file__`)
         - `test_file_dir` - the directory containing the current test file
         - `tests_dir` - the directory of the `tests` test suite
         - `examples_dir` - the directory of the `examples` test suite
         - `repo_root_dir` - the directory of the repository
         - `src_dir` - the directory of `src` (i.e. where the `transformers` sub-dir resides)
  
      - stringified paths---same as above but these return paths as strings, rather than `pathlib` objects:
  
         - `test_file_path_str`
         - `test_file_dir_str`
         - `tests_dir_str`
         - `examples_dir_str`
         - `repo_root_dir_str`
         - `src_dir_str`
  
      Feature 2: Flexible auto-removable temporary dirs which are guaranteed to get removed at the end of test.
  
      1. Create a unique temporary dir:
  
      ```python
      def test_whatever(self):
          tmp_dir = self.get_auto_remove_tmp_dir()
      ```
  
      `tmp_dir` will contain the path to the created temporary dir. It will be automatically removed at the end of the
      test.
  
  
      2. Create a temporary dir of my choice, ensure it's empty before the test starts and don't
      empty it after the test.
  
      ```python
      def test_whatever(self):
          tmp_dir = self.get_auto_remove_tmp_dir("./xxx")
      ```
  
      This is useful for debug when you want to monitor a specific directory and want to make sure the previous tests
      didn't leave any data in there.
  
      3. You can override the first two options by directly overriding the `before` and `after` args, leading to the
          following behavior:
  
      `before=True`: the temporary dir will always be cleared at the beginning of the test.
  
      `before=False`: if the temporary dir already existed, any existing files will remain there.
  
      `after=True`: the temporary dir will always be deleted at the end of the test.
  
      `after=False`: the temporary dir will always be left intact at the end of the test.
  
      Note 1: In order to run the equivalent of `rm -r` safely, only subdirs of the project repository checkout are
      allowed if an explicit `tmp_dir` is used, so that by mistake no `/tmp` or similar important part of the filesystem
      will get nuked. i.e. please always pass paths that start with `./`
  
      Note 2: Each test can register multiple temporary dirs and they all will get auto-removed, unless requested
      otherwise.
  
      Feature 3: Get a copy of the `os.environ` object that sets up `PYTHONPATH` specific to the current test suite. This
      is useful for invoking external programs from the test suite - e.g. distributed training.
  
  
      ```python
      def test_whatever(self):
          env = self.get_env()
      ```"""
  
      def setUp(self):
          # get_auto_remove_tmp_dir feature:
          self.teardown_tmp_dirs = []
  
          # figure out the resolved paths for repo_root, tests, examples, etc.
          self._test_file_path = inspect.getfile(self.__class__)
          path = Path(self._test_file_path).resolve()
          self._test_file_dir = path.parents[0]
          for up in [1, 2, 3]:
              tmp_dir = path.parents[up]
              if (tmp_dir / "src").is_dir() and (tmp_dir / "tests").is_dir():
                  break
          if tmp_dir:
              self._repo_root_dir = tmp_dir
          else:
              raise ValueError(f"can't figure out the root of the repo from {self._test_file_path}")
          self._tests_dir = self._repo_root_dir / "tests"
          self._examples_dir = self._repo_root_dir / "examples"
          self._src_dir = self._repo_root_dir / "src"
  
      @property
      def test_file_path(self):
          return self._test_file_path
  
      @property
      def test_file_path_str(self):
          return str(self._test_file_path)
  
      @property
      def test_file_dir(self):
          return self._test_file_dir
  
      @property
      def test_file_dir_str(self):
          return str(self._test_file_dir)
  
      @property
      def tests_dir(self):
          return self._tests_dir
  
      @property
      def tests_dir_str(self):
          return str(self._tests_dir)
  
      @property
      def examples_dir(self):
          return self._examples_dir
  
      @property
      def examples_dir_str(self):
          return str(self._examples_dir)
  
      @property
      def repo_root_dir(self):
          return self._repo_root_dir
  
      @property
      def repo_root_dir_str(self):
          return str(self._repo_root_dir)
  
      @property
      def src_dir(self):
          return self._src_dir
  
      @property
      def src_dir_str(self):
          return str(self._src_dir)
  
      def get_env(self):
          """
          Return a copy of the `os.environ` object that sets up `PYTHONPATH` correctly, depending on the test suite it's
          invoked from. This is useful for invoking external programs from the test suite - e.g. distributed training.
  
          It always inserts `./src` first, then `./tests` or `./examples` depending on the test suite type and finally
          the preset `PYTHONPATH` if any (all full resolved paths).
  
          """
          env = os.environ.copy()
          paths = [self.repo_root_dir_str, self.src_dir_str]
          if "/examples" in self.test_file_dir_str:
              paths.append(self.examples_dir_str)
          paths.append(env.get("PYTHONPATH", ""))
  
          env["PYTHONPATH"] = ":".join(paths)
          return env
  
      def get_auto_remove_tmp_dir(self, tmp_dir=None, before=None, after=None, return_pathlib_obj=False):
          """
          Args:
              tmp_dir (`string`, *optional*):
                  if `None`:
  
                     - a unique temporary path will be created
                     - sets `before=True` if `before` is `None`
                     - sets `after=True` if `after` is `None`
                  else:
  
                     - `tmp_dir` will be created
                     - sets `before=True` if `before` is `None`
                     - sets `after=False` if `after` is `None`
              before (`bool`, *optional*):
                  If `True` and the `tmp_dir` already exists, make sure to empty it right away if `False` and the
                  `tmp_dir` already exists, any existing files will remain there.
              after (`bool`, *optional*):
                  If `True`, delete the `tmp_dir` at the end of the test if `False`, leave the `tmp_dir` and its contents
                  intact at the end of the test.
              return_pathlib_obj (`bool`, *optional*):
                  If `True` will return a pathlib.Path object
  
          Returns:
              tmp_dir(`string`): either the same value as passed via *tmp_dir* or the path to the auto-selected tmp dir
          """
          if tmp_dir is not None:
              # defining the most likely desired behavior for when a custom path is provided.
              # this most likely indicates the debug mode where we want an easily locatable dir that:
              # 1. gets cleared out before the test (if it already exists)
              # 2. is left intact after the test
              if before is None:
                  before = True
              if after is None:
                  after = False
  
              # using provided path
              path = Path(tmp_dir).resolve()
  
              # to avoid nuking parts of the filesystem, only relative paths are allowed
              if not tmp_dir.startswith("./"):
                  raise ValueError(
                      f"`tmp_dir` can only be a relative path, i.e. `./some/path`, but received `{tmp_dir}`"
                  )
  
              # ensure the dir is empty to start with
              if before is True and path.exists():
                  shutil.rmtree(tmp_dir, ignore_errors=True)
  
              path.mkdir(parents=True, exist_ok=True)
  
          else:
              # defining the most likely desired behavior for when a unique tmp path is auto generated
              # (not a debug mode), here we require a unique tmp dir that:
              # 1. is empty before the test (it will be empty in this situation anyway)
              # 2. gets fully removed after the test
              if before is None:
                  before = True
              if after is None:
                  after = True
  
              # using unique tmp dir (always empty, regardless of `before`)
              tmp_dir = tempfile.mkdtemp()
  
          if after is True:
              # register for deletion
              self.teardown_tmp_dirs.append(tmp_dir)
  
          return Path(tmp_dir).resolve() if return_pathlib_obj else tmp_dir
  
      def python_one_liner_max_rss(self, one_liner_str):
          """
          Runs the passed python one liner (just the code) and returns how much max cpu memory was used to run the
          program.
  
          Args:
              one_liner_str (`string`):
                  a python one liner code that gets passed to `python -c`
  
          Returns:
              max cpu memory bytes used to run the program. This value is likely to vary slightly from run to run.
  
          Requirements:
              this helper needs `/usr/bin/time` to be installed (`apt install time`)
  
          Example:
  
          ```
          one_liner_str = 'from transformers import AutoModel; AutoModel.from_pretrained("google-t5/t5-large")'
          max_rss = self.python_one_liner_max_rss(one_liner_str)
          ```
          """
  
          if not cmd_exists("/usr/bin/time"):
              raise ValueError("/usr/bin/time is required, install with `apt install time`")
  
          cmd = shlex.split(f"/usr/bin/time -f %M python -c '{one_liner_str}'")
          with CaptureStd() as cs:
              execute_subprocess_async(cmd, env=self.get_env())
          # returned data is in KB so convert to bytes
          max_rss = int(cs.err.split("\n")[-2].replace("stderr: ", "")) * 1024
          return max_rss
  
      def tearDown(self):
          # get_auto_remove_tmp_dir feature: remove registered temp dirs
          for path in self.teardown_tmp_dirs:
              shutil.rmtree(path, ignore_errors=True)
          self.teardown_tmp_dirs = []
          if is_accelerate_available():
              AcceleratorState._reset_state()
              PartialState._reset_state()
  
              # delete all the env variables having `ACCELERATE` in them
              for k in list(os.environ.keys()):
                  if "ACCELERATE" in k:
                      del os.environ[k]
  
  
  def mockenv(**kwargs):
      """
      this is a convenience wrapper, that allows this ::
  
      @mockenv(RUN_SLOW=True, USE_TF=False) def test_something():
          run_slow = os.getenv("RUN_SLOW", False) use_tf = os.getenv("USE_TF", False)
  
      """
      return mock.patch.dict(os.environ, kwargs)
  
  
  # from https://stackoverflow.com/a/34333710/9201239
  @contextlib.contextmanager
  def mockenv_context(*remove, **update):
      """
      Temporarily updates the `os.environ` dictionary in-place. Similar to mockenv
  
      The `os.environ` dictionary is updated in-place so that the modification is sure to work in all situations.
  
      Args:
        remove: Environment variables to remove.
        update: Dictionary of environment variables and values to add/update.
      """
      env = os.environ
      update = update or {}
      remove = remove or []
  
      # List of environment variables being updated or removed.
      stomped = (set(update.keys()) | set(remove)) & set(env.keys())
      # Environment variables and values to restore on exit.
      update_after = {k: env[k] for k in stomped}
      # Environment variables and values to remove on exit.
      remove_after = frozenset(k for k in update if k not in env)
  
      try:
          env.update(update)
          [env.pop(k, None) for k in remove]
          yield
      finally:
          env.update(update_after)
          [env.pop(k) for k in remove_after]
  
  
  # --- pytest conf functions --- #
  
  # to avoid multiple invocation from tests/conftest.py and examples/conftest.py - make sure it's called only once
  pytest_opt_registered = {}
  
  
  def pytest_addoption_shared(parser):
      """
      This function is to be called from `conftest.py` via `pytest_addoption` wrapper that has to be defined there.
  
      It allows loading both `conftest.py` files at once without causing a failure due to adding the same `pytest`
      option.
  
      """
      option = "--make-reports"
      if option not in pytest_opt_registered:
          parser.addoption(
              option,
              action="store",
              default=False,
              help="generate report files. The value of this option is used as a prefix to report names",
          )
          pytest_opt_registered[option] = 1
  
  
  def pytest_terminal_summary_main(tr, id):
      """
      Generate multiple reports at the end of test suite run - each report goes into a dedicated file in the current
      directory. The report files are prefixed with the test suite name.
  
      This function emulates --duration and -rA pytest arguments.
  
      This function is to be called from `conftest.py` via `pytest_terminal_summary` wrapper that has to be defined
      there.
  
      Args:
      - tr: `terminalreporter` passed from `conftest.py`
      - id: unique id like `tests` or `examples` that will be incorporated into the final reports filenames - this is
        needed as some jobs have multiple runs of pytest, so we can't have them overwrite each other.
  
      NB: this functions taps into a private _pytest API and while unlikely, it could break should pytest do internal
      changes - also it calls default internal methods of terminalreporter which can be hijacked by various `pytest-`
      plugins and interfere.
  
      """
      from _pytest.config import create_terminal_writer
  
      if not len(id):
          id = "tests"
  
      config = tr.config
      orig_writer = config.get_terminal_writer()
      orig_tbstyle = config.option.tbstyle
      orig_reportchars = tr.reportchars
  
      dir = f"reports/{id}"
      Path(dir).mkdir(parents=True, exist_ok=True)
      report_files = {
          k: f"{dir}/{k}.txt"
          for k in [
              "durations",
              "errors",
              "failures_long",
              "failures_short",
              "failures_line",
              "passes",
              "stats",
              "summary_short",
              "warnings",
          ]
      }
  
      # custom durations report
      # note: there is no need to call pytest --durations=XX to get this separate report
      # adapted from https://github.com/pytest-dev/pytest/blob/897f151e/src/_pytest/runner.py#L66
      dlist = []
      for replist in tr.stats.values():
          for rep in replist:
              if hasattr(rep, "duration"):
                  dlist.append(rep)
      if dlist:
          dlist.sort(key=lambda x: x.duration, reverse=True)
          with open(report_files["durations"], "w") as f:
              durations_min = 0.05  # sec
              f.write("slowest durations\n")
              for i, rep in enumerate(dlist):
                  if rep.duration < durations_min:
                      f.write(f"{len(dlist) - i} durations < {durations_min} secs were omitted")
                      break
                  f.write(f"{rep.duration:02.2f}s {rep.when:<8} {rep.nodeid}\n")
  
      def summary_failures_short(tr):
          # expecting that the reports were --tb=long (default) so we chop them off here to the last frame
          reports = tr.getreports("failed")
          if not reports:
              return
          tr.write_sep("=", "FAILURES SHORT STACK")
          for rep in reports:
              msg = tr._getfailureheadline(rep)
              tr.write_sep("_", msg, red=True, bold=True)
              # chop off the optional leading extra frames, leaving only the last one
              longrepr = re.sub(r".*_ _ _ (_ ){10,}_ _ ", "", rep.longreprtext, 0, re.MULTILINE | re.DOTALL)
              tr._tw.line(longrepr)
              # note: not printing out any rep.sections to keep the report short
  
      # use ready-made report funcs, we are just hijacking the filehandle to log to a dedicated file each
      # adapted from https://github.com/pytest-dev/pytest/blob/897f151e/src/_pytest/terminal.py#L814
      # note: some pytest plugins may interfere by hijacking the default `terminalreporter` (e.g.
      # pytest-instafail does that)
  
      # report failures with line/short/long styles
      config.option.tbstyle = "auto"  # full tb
      with open(report_files["failures_long"], "w") as f:
          tr._tw = create_terminal_writer(config, f)
          tr.summary_failures()
  
      # config.option.tbstyle = "short" # short tb
      with open(report_files["failures_short"], "w") as f:
          tr._tw = create_terminal_writer(config, f)
          summary_failures_short(tr)
  
      config.option.tbstyle = "line"  # one line per error
      with open(report_files["failures_line"], "w") as f:
          tr._tw = create_terminal_writer(config, f)
          tr.summary_failures()
  
      with open(report_files["errors"], "w") as f:
          tr._tw = create_terminal_writer(config, f)
          tr.summary_errors()
  
      with open(report_files["warnings"], "w") as f:
          tr._tw = create_terminal_writer(config, f)
          tr.summary_warnings()  # normal warnings
          tr.summary_warnings()  # final warnings
  
      tr.reportchars = "wPpsxXEf"  # emulate -rA (used in summary_passes() and short_test_summary())
  
      # Skip the `passes` report, as it starts to take more than 5 minutes, and sometimes it timeouts on CircleCI if it
      # takes > 10 minutes (as this part doesn't generate any output on the terminal).
      # (also, it seems there is no useful information in this report, and we rarely need to read it)
      # with open(report_files["passes"], "w") as f:
      #     tr._tw = create_terminal_writer(config, f)
      #     tr.summary_passes()
  
      with open(report_files["summary_short"], "w") as f:
          tr._tw = create_terminal_writer(config, f)
          tr.short_test_summary()
  
      with open(report_files["stats"], "w") as f:
          tr._tw = create_terminal_writer(config, f)
          tr.summary_stats()
  
      # restore:
      tr._tw = orig_writer
      tr.reportchars = orig_reportchars
      config.option.tbstyle = orig_tbstyle
  
  
  # --- distributed testing functions --- #
  
  # adapted from https://stackoverflow.com/a/59041913/9201239
  import asyncio  # noqa
  
  
  class _RunOutput:
      def __init__(self, returncode, stdout, stderr):
          self.returncode = returncode
          self.stdout = stdout
          self.stderr = stderr
  
  
  async def _read_stream(stream, callback):
      while True:
          line = await stream.readline()
          if line:
              callback(line)
          else:
              break
  
  
  async def _stream_subprocess(cmd, env=None, stdin=None, timeout=None, quiet=False, echo=False) -> _RunOutput:
      if echo:
          print("\nRunning: ", " ".join(cmd))
  
      p = await asyncio.create_subprocess_exec(
          cmd[0],
          *cmd[1:],
          stdin=stdin,
          stdout=asyncio.subprocess.PIPE,
          stderr=asyncio.subprocess.PIPE,
          env=env,
      )
  
      # note: there is a warning for a possible deadlock when using `wait` with huge amounts of data in the pipe
      # https://docs.python.org/3/library/asyncio-subprocess.html#asyncio.asyncio.subprocess.Process.wait
      #
      # If it starts hanging, will need to switch to the following code. The problem is that no data
      # will be seen until it's done and if it hangs for example there will be no debug info.
      # out, err = await p.communicate()
      # return _RunOutput(p.returncode, out, err)
  
      out = []
      err = []
  
      def tee(line, sink, pipe, label=""):
          line = line.decode("utf-8").rstrip()
          sink.append(line)
          if not quiet:
              print(label, line, file=pipe)
  
      # XXX: the timeout doesn't seem to make any difference here
      await asyncio.wait(
          [
              asyncio.create_task(_read_stream(p.stdout, lambda l: tee(l, out, sys.stdout, label="stdout:"))),
              asyncio.create_task(_read_stream(p.stderr, lambda l: tee(l, err, sys.stderr, label="stderr:"))),
          ],
          timeout=timeout,
      )
      return _RunOutput(await p.wait(), out, err)
  
  
  def execute_subprocess_async(cmd, env=None, stdin=None, timeout=180, quiet=False, echo=True) -> _RunOutput:
      loop = asyncio.get_event_loop()
      result = loop.run_until_complete(
          _stream_subprocess(cmd, env=env, stdin=stdin, timeout=timeout, quiet=quiet, echo=echo)
      )
  
      cmd_str = " ".join(cmd)
      if result.returncode > 0:
          stderr = "\n".join(result.stderr)
          raise RuntimeError(
              f"'{cmd_str}' failed with returncode {result.returncode}\n\n"
              f"The combined stderr from workers follows:\n{stderr}"
          )
  
      # check that the subprocess actually did run and produced some output, should the test rely on
      # the remote side to do the testing
      if not result.stdout and not result.stderr:
          raise RuntimeError(f"'{cmd_str}' produced no output.")
  
      return result
  
  
  def pytest_xdist_worker_id():
      """
      Returns an int value of worker's numerical id under `pytest-xdist`'s concurrent workers `pytest -n N` regime, or 0
      if `-n 1` or `pytest-xdist` isn't being used.
      """
      worker = os.environ.get("PYTEST_XDIST_WORKER", "gw0")
      worker = re.sub(r"^gw", "", worker, 0, re.MULTILINE)
      return int(worker)
  
  
  def get_torch_dist_unique_port():
      """
      Returns a free port number that can be fed to `torch.distributed.launch`'s `--master_port` argument.
  
      Binds to port 0 to let the OS assign an available port, avoiding collisions from hardcoded ports
      and TCP TIME_WAIT issues between sequential subprocess launches.
      """
      import socket
  
      with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
          s.bind(("", 0))
          return s.getsockname()[1]
  
  
  def nested_simplify(obj, decimals=3):
      """
      Simplifies an object by rounding float numbers, and downcasting tensors/numpy arrays to get simple equality test
      within tests.
      """
      import numpy as np
  
      if isinstance(obj, list):
          return [nested_simplify(item, decimals) for item in obj]
      if isinstance(obj, tuple):
          return tuple(nested_simplify(item, decimals) for item in obj)
      elif isinstance(obj, np.ndarray):
          return nested_simplify(obj.tolist())
      elif isinstance(obj, Mapping):
          return {nested_simplify(k, decimals): nested_simplify(v, decimals) for k, v in obj.items()}
      elif isinstance(obj, (str, int, np.int64)) or obj is None:
          return obj
      elif is_torch_available() and isinstance(obj, torch.Tensor):
          return nested_simplify(obj.tolist(), decimals)
      elif isinstance(obj, float):
          return round(obj, decimals)
      elif isinstance(obj, (np.int32, np.float32, np.float16)):
          return nested_simplify(obj.item(), decimals)
      else:
          raise Exception(f"Not supported: {type(obj)}")
  
  
  def check_json_file_has_correct_format(file_path):
      with open(file_path) as f:
          lines = f.readlines()
          if len(lines) == 1:
              # length can only be 1 if dict is empty
              assert lines[0] == "{}"
          else:
              # otherwise make sure json has correct format (at least 3 lines)
              assert len(lines) >= 3
              # each key one line, ident should be 2, min length is 3
              assert lines[0].strip() == "{"
              for line in lines[1:-1]:
                  left_indent = len(lines[1]) - len(lines[1].lstrip())
                  assert left_indent == 2
              assert lines[-1].strip() == "}"
  
  
  def to_2tuple(x):
      if isinstance(x, collections.abc.Iterable):
          return x
      return (x, x)
  
  
  # These utils relate to ensuring the right error message is received when running scripts
  class SubprocessCallException(Exception):
      pass
  
  
  def run_command(command: list[str], return_stdout=False):
      """
      Runs `command` with `subprocess.check_output` and will potentially return the `stdout`. Will also properly capture
      if an error occurred while running `command`
      """
      try:
          output = subprocess.check_output(command, stderr=subprocess.STDOUT)
          if return_stdout:
              if hasattr(output, "decode"):
                  output = output.decode("utf-8")
              return output
      except subprocess.CalledProcessError as e:
          raise SubprocessCallException(
              f"Command `{' '.join(command)}` failed with the following error:\n\n{e.output.decode()}"
          ) from e
  
  
  class RequestCounter:
      """
      Helper class that will count all requests made online.
  
      Might not be robust if urllib3 changes its logging format but should be good enough for us.
  
      Usage:
      ```py
      with RequestCounter() as counter:
          _ = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-bert")
      assert counter["GET"] == 0
      assert counter["HEAD"] == 1
      assert counter.total_calls == 1
      ```
      """
  
      def __enter__(self):
          self._counter = defaultdict(int)
          self._thread_id = threading.get_ident()
          self._extra_info = []
  
          def patched_with_thread_info(func):
              def wrap(*args, **kwargs):
                  self._extra_info.append(threading.get_ident())
                  return func(*args, **kwargs)
  
              return wrap
  
          import urllib3
  
          self.patcher = patch.object(
              urllib3.connectionpool.log, "debug", side_effect=patched_with_thread_info(urllib3.connectionpool.log.debug)
          )
          self.mock = self.patcher.start()
          return self
  
      def __exit__(self, *args, **kwargs) -> None:
          assert len(self.mock.call_args_list) == len(self._extra_info)
          for thread_id, call in zip(self._extra_info, self.mock.call_args_list):
              if thread_id != self._thread_id:
                  continue
              # code 307: the URL being requested by the user has moved to a temporary location
              if call.args[-2] == 307:
                  continue
              log = call.args[0] % call.args[1:]
              for method in ("HEAD", "GET", "POST", "PUT", "DELETE", "CONNECT", "OPTIONS", "TRACE", "PATCH"):
                  if method in log:
                      self._counter[method] += 1
                      break
          self.patcher.stop()
  
      def __getitem__(self, key: str) -> int:
          return self._counter[key]
  
      @property
      def total_calls(self) -> int:
          return sum(self._counter.values())
  
  
  def is_flaky(max_attempts: int = 5, wait_before_retry: float | None = None, description: str | None = None):
      """
      To decorate flaky tests. They will be retried on failures.
  
      Please note that our push tests use `pytest-rerunfailures`, which prompts the CI to rerun certain types of
      failed tests. More specifically, if the test exception contains any substring in `FLAKY_TEST_FAILURE_PATTERNS`
      (in `.circleci/create_circleci_config.py`), it will be rerun. If you find a recurrent pattern of failures,
      expand `FLAKY_TEST_FAILURE_PATTERNS` in our CI configuration instead of using `is_flaky`.
  
      Args:
          max_attempts (`int`, *optional*, defaults to 5):
              The maximum number of attempts to retry the flaky test.
          wait_before_retry (`float`, *optional*):
              If provided, will wait that number of seconds before retrying the test.
          description (`str`, *optional*):
              A string to describe the situation (what / where / why is flaky, link to GH issue/PR comments, errors,
              etc.)
      """
  
      def decorator(test_func_ref):
          @functools.wraps(test_func_ref)
          def wrapper(*args, **kwargs):
              retry_count = 1
  
              while retry_count < max_attempts:
                  try:
                      return test_func_ref(*args, **kwargs)
  
                  except Exception as err:
                      logger.error(f"Test failed with {err} at try {retry_count}/{max_attempts}.")
                      if wait_before_retry is not None:
                          time.sleep(wait_before_retry)
                      retry_count += 1
  
              return test_func_ref(*args, **kwargs)
  
          return unittest.skipUnless(_run_flaky_tests, "test is flaky")(wrapper)
  
      return decorator
  
  
  def hub_retry(max_attempts: int = 5, wait_before_retry: float | None = 2):
      """
      To decorate tests that download from the Hub. They can fail due to a
      variety of network issues such as timeouts, connection resets, etc.
  
      Uses exponential backoff starting from `wait_before_retry`.
  
      Args:
          max_attempts (`int`, *optional*, defaults to 5):
              The maximum number of attempts to retry the flaky test.
          wait_before_retry (`float`, *optional*, defaults to 2):
              If provided, the initial delay in seconds before the first retry.
              Subsequent retries use exponential backoff with jitter.
      """
      from .utils.generic import retry
  
      return retry(
          max_retries=max_attempts,
          initial_delay=wait_before_retry or 0,
          jitter=wait_before_retry is not None,
          exceptions=(httpx.HTTPError,),
      )
  
  
  def run_first(test_case):
      """
      Decorator marking a test with order(1). When pytest-order plugin is installed, tests marked with this decorator
      are guaranteed to run first.
  
      This is especially useful in some test settings like on a Gaudi instance where a Gaudi device can only be used by a
      single process at a time. So we make sure all tests that run in a subprocess are launched first, to avoid device
      allocation conflicts.
      """
      # Without this check, we get unwanted warnings when it's not installed
      if is_pytest_order_available():
          import pytest
  
          return pytest.mark.order(1)(test_case)
      else:
          return test_case
  
  
  def run_test_in_subprocess(test_case, target_func, inputs=None, timeout=None):
      """
      To run a test in a subprocess. In particular, this can avoid (GPU) memory issue.
  
      Args:
          test_case (`unittest.TestCase`):
              The test that will run `target_func`.
          target_func (`Callable`):
              The function implementing the actual testing logic.
          inputs (`dict`, *optional*, defaults to `None`):
              The inputs that will be passed to `target_func` through an (input) queue.
          timeout (`int`, *optional*, defaults to `None`):
              The timeout (in seconds) that will be passed to the input and output queues. If not specified, the env.
              variable `PYTEST_TIMEOUT` will be checked. If still `None`, its value will be set to `600`.
      """
      if timeout is None:
          timeout = int(os.environ.get("PYTEST_TIMEOUT", "600"))
  
      start_methohd = "spawn"
      ctx = multiprocessing.get_context(start_methohd)
  
      input_queue = ctx.Queue(1)
      output_queue = ctx.JoinableQueue(1)
  
      # We can't send `unittest.TestCase` to the child, otherwise we get issues regarding pickle.
      input_queue.put(inputs, timeout=timeout)
  
      process = ctx.Process(target=target_func, args=(input_queue, output_queue, timeout))
      process.start()
      # Kill the child process if we can't get outputs from it in time: otherwise, the hanging subprocess prevents
      # the test to exit properly.
      try:
          results = output_queue.get(timeout=timeout)
          output_queue.task_done()
      except Exception as e:
          process.terminate()
          test_case.fail(e)
      process.join(timeout=timeout)
  
      if results["error"] is not None:
          test_case.fail(f"{results['error']}")
  
  
  def run_test_using_subprocess(func):
      """
      To decorate a test to run in a subprocess using the `subprocess` module. This could avoid potential GPU memory
      issues (GPU OOM or a test that causes many subsequential failing with `CUDA error: device-side assert triggered`).
      """
      import pytest
  
      @functools.wraps(func)
      def wrapper(*args, **kwargs):
          if os.getenv("_INSIDE_SUB_PROCESS", None) == "1":
              func(*args, **kwargs)
          else:
              test = " ".join(os.environ.get("PYTEST_CURRENT_TEST").split(" ")[:-1])
              try:
                  env = copy.deepcopy(os.environ)
                  env["_INSIDE_SUB_PROCESS"] = "1"
                  # This prevents the entries in `short test summary info` given by the subprocess being truncated. so the
                  # full information can be passed to the parent pytest process.
                  # See: https://docs.pytest.org/en/stable/explanation/ci.html
                  env["CI"] = "true"
  
                  # If not subclass of `unitTest.TestCase` and `pytestconfig` is used: try to grab and use the arguments
                  if "pytestconfig" in kwargs:
                      command = list(kwargs["pytestconfig"].invocation_params.args)
                      for idx, x in enumerate(command):
                          if x in kwargs["pytestconfig"].args:
                              test = test.split("::")[1:]
                              command[idx] = "::".join([f"{func.__globals__['__file__']}"] + test)
                      command = [f"{sys.executable}", "-m", "pytest"] + command
                      command = [x for x in command if x != "--no-summary"]
                  # Otherwise, simply run the test with no option at all
                  else:
                      command = [f"{sys.executable}", "-m", "pytest", f"{test}"]
  
                  subprocess.run(command, env=env, check=True, capture_output=True)
              except subprocess.CalledProcessError as e:
                  exception_message = e.stdout.decode()
                  lines = exception_message.split("\n")
                  # Add a first line with more informative information instead of just `= test session starts =`.
                  # This makes the `short test summary info` section more useful.
                  if "= test session starts =" in lines[0]:
                      text = ""
                      for line in lines[1:]:
                          if line.startswith("FAILED "):
                              text = line[len("FAILED ") :]
                              text = "".join(text.split(" - ")[1:])
                          elif line.startswith("=") and line.endswith("=") and " failed in " in line:
                              break
                          elif len(text) > 0:
                              text += f"\n{line}"
                      text = "(subprocess) " + text
                      lines = [text] + lines
                  exception_message = "\n".join(lines)
                  raise pytest.fail(exception_message, pytrace=False)
  
      return wrapper
  
  
  """
  The following contains utils to run the documentation tests without having to overwrite any files.
  
  The `preprocess_string` function adds `# doctest: +IGNORE_RESULT` markers on the fly anywhere a `load_dataset` call is
  made as a print would otherwise fail the corresponding line.
  
  To skip cuda tests, make sure to call `SKIP_CUDA_DOCTEST=1 pytest --doctest-modules <path_to_files_to_test>
  """
  
  
  def preprocess_string(string, skip_cuda_tests):
      """Prepare a docstring or a `.md` file to be run by doctest.
  
      The argument `string` would be the whole file content if it is a `.md` file. For a python file, it would be one of
      its docstring. In each case, it may contain multiple python code examples. If `skip_cuda_tests` is `True` and a
      cuda stuff is detective (with a heuristic), this method will return an empty string so no doctest will be run for
      `string`.
      """
      codeblock_pattern = r"(```(?:python|py)\s*\n\s*>>> )(.*?```)"
      codeblocks = re.split(codeblock_pattern, string, flags=re.DOTALL)
      is_cuda_found = False
      for i, codeblock in enumerate(codeblocks):
          if "load_dataset(" in codeblock and "# doctest: +IGNORE_RESULT" not in codeblock:
              codeblocks[i] = re.sub(r"(>>> .*load_dataset\(.*)", r"\1 # doctest: +IGNORE_RESULT", codeblock)
          if (
              (">>>" in codeblock or "..." in codeblock)
              and re.search(r"cuda|to\(0\)|device=0", codeblock)
              and skip_cuda_tests
          ):
              is_cuda_found = True
              break
  
      modified_string = ""
      if not is_cuda_found:
          modified_string = "".join(codeblocks)
  
      return modified_string
  
  
  class HfDocTestParser(doctest.DocTestParser):
      """
      Overwrites the DocTestParser from doctest to properly parse the codeblocks that are formatted with black. This
      means that there are no extra lines at the end of our snippets. The `# doctest: +IGNORE_RESULT` marker is also
      added anywhere a `load_dataset` call is made as a print would otherwise fail the corresponding line.
  
      Tests involving cuda are skipped base on a naive pattern that should be updated if it is not enough.
      """
  
      # This regular expression is used to find doctest examples in a
      # string.  It defines three groups: `source` is the source code
      # (including leading indentation and prompts); `indent` is the
      # indentation of the first (PS1) line of the source code; and
      # `want` is the expected output (including leading indentation).
      # fmt: off
      _EXAMPLE_RE = re.compile(r'''
          # Source consists of a PS1 line followed by zero or more PS2 lines.
          (?P<source>
              (?:^(?P<indent> [ ]*) >>>    .*)    # PS1 line
              (?:\n           [ ]*  \.\.\. .*)*)  # PS2 lines
          \n?
          # Want consists of any non-blank lines that do not start with PS1.
          (?P<want> (?:(?![ ]*$)    # Not a blank line
               (?![ ]*>>>)          # Not a line starting with PS1
               # !!!!!!!!!!! HF Specific !!!!!!!!!!!
               (?:(?!```).)*        # Match any character except '`' until a '```' is found (this is specific to HF because black removes the last line)
               # !!!!!!!!!!! HF Specific !!!!!!!!!!!
               (?:\n|$)  # Match a new line or end of string
            )*)
          ''', re.MULTILINE | re.VERBOSE
      )
      # fmt: on
  
      # !!!!!!!!!!! HF Specific !!!!!!!!!!!
      skip_cuda_tests: bool = os.environ.get("SKIP_CUDA_DOCTEST", "0") == "1"
      # !!!!!!!!!!! HF Specific !!!!!!!!!!!
  
      def parse(self, string, name="<string>"):
          """
          Overwrites the `parse` method to incorporate a skip for CUDA tests, and remove logs and dataset prints before
          calling `super().parse`
          """
          string = preprocess_string(string, self.skip_cuda_tests)
          return super().parse(string, name)
  
  
  class HfDoctestModule(Module):
      """
      Overwrites the `DoctestModule` of the pytest package to make sure the HFDocTestParser is used when discovering
      tests.
      """
  
      def collect(self) -> Iterable[DoctestItem]:
          class MockAwareDocTestFinder(doctest.DocTestFinder):
              """A hackish doctest finder that overrides stdlib internals to fix a stdlib bug.
  
              https://github.com/pytest-dev/pytest/issues/3456 https://bugs.python.org/issue25532
              """
  
              def _find_lineno(self, obj, source_lines):
                  """Doctest code does not take into account `@property`, this
                  is a hackish way to fix it. https://bugs.python.org/issue17446
  
                  Wrapped Doctests will need to be unwrapped so the correct line number is returned. This will be
                  reported upstream. #8796
                  """
                  if isinstance(obj, property):
                      obj = getattr(obj, "fget", obj)
  
                  if hasattr(obj, "__wrapped__"):
                      # Get the main obj in case of it being wrapped
                      obj = inspect.unwrap(obj)
  
                  # Type ignored because this is a private function.
                  return super()._find_lineno(  # type:ignore[misc]
                      obj,
                      source_lines,
                  )
  
              def _find(self, tests, obj, name, module, source_lines, globs, seen) -> None:
                  if _is_mocked(obj):
                      return
                  with _patch_unwrap_mock_aware():
                      # Type ignored because this is a private function.
                      super()._find(  # type:ignore[misc]
                          tests, obj, name, module, source_lines, globs, seen
                      )
  
          if self.path.name == "conftest.py":
              module = self.config.pluginmanager._importconftest(
                  self.path,
                  self.config.getoption("importmode"),
                  rootpath=self.config.rootpath,
              )
          else:
              try:
                  module = import_path(
                      self.path,
                      root=self.config.rootpath,
                      mode=self.config.getoption("importmode"),
                  )
              except ImportError:
                  if self.config.getvalue("doctest_ignore_import_errors"):
                      skip("unable to import module %r" % self.path)
                  else:
                      raise
  
          # !!!!!!!!!!! HF Specific !!!!!!!!!!!
          finder = MockAwareDocTestFinder(parser=HfDocTestParser())
          # !!!!!!!!!!! HF Specific !!!!!!!!!!!
          optionflags = get_optionflags(self)
          runner = _get_runner(
              verbose=False,
              optionflags=optionflags,
              checker=_get_checker(),
              continue_on_failure=_get_continue_on_failure(self.config),
          )
          for test in finder.find(module, module.__name__):
              if test.examples:  # skip empty doctests and cuda
                  yield DoctestItem.from_parent(self, name=test.name, runner=runner, dtest=test)
  
  
  def _device_agnostic_dispatch(device: str, dispatch_table: dict[str, Callable], *args, **kwargs):
      if device not in dispatch_table:
          if not callable(dispatch_table["default"]):
              return dispatch_table["default"]
  
          return dispatch_table["default"](*args, **kwargs)
  
      fn = dispatch_table[device]
  
      # Some device agnostic functions return values or None, will return then directly.
      if not callable(fn):
          return fn
  
      return fn(*args, **kwargs)
  
  
  if is_torch_available():
      # Mappings from device names to callable functions to support device agnostic
      # testing.
      BACKEND_MANUAL_SEED = {
          "cuda": torch.cuda.manual_seed,
          "cpu": torch.manual_seed,
          "default": torch.manual_seed,
      }
      BACKEND_EMPTY_CACHE = {
          "cuda": torch.cuda.empty_cache,
          "cpu": None,
          "default": None,
      }
      BACKEND_DEVICE_COUNT = {
          "cuda": torch.cuda.device_count,
          "cpu": lambda: 0,
          "default": lambda: 1,
      }
      BACKEND_RESET_MAX_MEMORY_ALLOCATED = {
          "cuda": torch.cuda.reset_max_memory_allocated,
          "cpu": None,
          "default": None,
      }
      BACKEND_MAX_MEMORY_ALLOCATED = {
          "cuda": torch.cuda.max_memory_allocated,
          "cpu": 0,
          "default": 0,
      }
      BACKEND_RESET_PEAK_MEMORY_STATS = {
          "cuda": torch.cuda.reset_peak_memory_stats,
          "cpu": None,
          "default": None,
      }
      BACKEND_MEMORY_ALLOCATED = {
          "cuda": torch.cuda.memory_allocated,
          "cpu": 0,
          "default": 0,
      }
      BACKEND_SYNCHRONIZE = {
          "cuda": torch.cuda.synchronize,
          "cpu": None,
          "default": None,
      }
      BACKEND_TORCH_ACCELERATOR_MODULE = {
          "cuda": torch.cuda,
          "cpu": None,
          "default": None,
      }
  else:
      BACKEND_MANUAL_SEED = {"default": None}
      BACKEND_EMPTY_CACHE = {"default": None}
      BACKEND_DEVICE_COUNT = {"default": lambda: 0}
      BACKEND_RESET_MAX_MEMORY_ALLOCATED = {"default": None}
      BACKEND_RESET_PEAK_MEMORY_STATS = {"default": None}
      BACKEND_MAX_MEMORY_ALLOCATED = {"default": 0}
      BACKEND_MEMORY_ALLOCATED = {"default": 0}
      BACKEND_SYNCHRONIZE = {"default": None}
      BACKEND_TORCH_ACCELERATOR_MODULE = {"default": None}
  
  
  if is_torch_hpu_available():
      BACKEND_MANUAL_SEED["hpu"] = torch.hpu.manual_seed
      BACKEND_DEVICE_COUNT["hpu"] = torch.hpu.device_count
      BACKEND_TORCH_ACCELERATOR_MODULE["hpu"] = torch.hpu
  
  if is_torch_mlu_available():
      BACKEND_EMPTY_CACHE["mlu"] = torch.mlu.empty_cache
      BACKEND_MANUAL_SEED["mlu"] = torch.mlu.manual_seed
      BACKEND_DEVICE_COUNT["mlu"] = torch.mlu.device_count
      BACKEND_TORCH_ACCELERATOR_MODULE["mlu"] = torch.mlu
  
  if is_torch_npu_available():
      BACKEND_EMPTY_CACHE["npu"] = torch.npu.empty_cache
      BACKEND_MANUAL_SEED["npu"] = torch.npu.manual_seed
      BACKEND_DEVICE_COUNT["npu"] = torch.npu.device_count
      BACKEND_TORCH_ACCELERATOR_MODULE["npu"] = torch.npu
  
  if is_torch_xpu_available():
      BACKEND_EMPTY_CACHE["xpu"] = torch.xpu.empty_cache
      BACKEND_MANUAL_SEED["xpu"] = torch.xpu.manual_seed
      BACKEND_DEVICE_COUNT["xpu"] = torch.xpu.device_count
      BACKEND_RESET_MAX_MEMORY_ALLOCATED["xpu"] = torch.xpu.reset_peak_memory_stats
      BACKEND_RESET_PEAK_MEMORY_STATS["xpu"] = torch.xpu.reset_peak_memory_stats
      BACKEND_MAX_MEMORY_ALLOCATED["xpu"] = torch.xpu.max_memory_allocated
      BACKEND_MEMORY_ALLOCATED["xpu"] = torch.xpu.memory_allocated
      BACKEND_SYNCHRONIZE["xpu"] = torch.xpu.synchronize
      BACKEND_TORCH_ACCELERATOR_MODULE["xpu"] = torch.xpu
  
  
  if is_torch_xla_available():
      BACKEND_EMPTY_CACHE["xla"] = torch.cuda.empty_cache
      BACKEND_MANUAL_SEED["xla"] = torch.cuda.manual_seed
      BACKEND_DEVICE_COUNT["xla"] = torch.cuda.device_count
  
  
  def backend_manual_seed(device: str, seed: int):
      return _device_agnostic_dispatch(device, BACKEND_MANUAL_SEED, seed)
  
  
  def backend_empty_cache(device: str):
      return _device_agnostic_dispatch(device, BACKEND_EMPTY_CACHE)
  
  
  def backend_device_count(device: str):
      return _device_agnostic_dispatch(device, BACKEND_DEVICE_COUNT)
  
  
  def backend_reset_max_memory_allocated(device: str):
      return _device_agnostic_dispatch(device, BACKEND_RESET_MAX_MEMORY_ALLOCATED)
  
  
  def backend_reset_peak_memory_stats(device: str):
      return _device_agnostic_dispatch(device, BACKEND_RESET_PEAK_MEMORY_STATS)
  
  
  def backend_max_memory_allocated(device: str):
      return _device_agnostic_dispatch(device, BACKEND_MAX_MEMORY_ALLOCATED)
  
  
  def backend_memory_allocated(device: str):
      return _device_agnostic_dispatch(device, BACKEND_MEMORY_ALLOCATED)
  
  
  def backend_synchronize(device: str):
      return _device_agnostic_dispatch(device, BACKEND_SYNCHRONIZE)
  
  
  def backend_torch_accelerator_module(device: str):
      return _device_agnostic_dispatch(device, BACKEND_TORCH_ACCELERATOR_MODULE)
  
  
  if is_torch_available():
      # If `TRANSFORMERS_TEST_DEVICE_SPEC` is enabled we need to import extra entries
      # into device to function mappings.
      if "TRANSFORMERS_TEST_DEVICE_SPEC" in os.environ:
          device_spec_path = os.environ["TRANSFORMERS_TEST_DEVICE_SPEC"]
          if not Path(device_spec_path).is_file():
              raise ValueError(
                  f"Specified path to device spec file is not a file or not found. Received '{device_spec_path}"
              )
  
          # Try to strip extension for later import – also verifies we are importing a
          # python file.
          device_spec_dir, _ = os.path.split(os.path.realpath(device_spec_path))
          sys.path.append(device_spec_dir)
          try:
              import_name = device_spec_path[: device_spec_path.index(".py")]
          except ValueError as e:
              raise ValueError(f"Provided device spec file was not a Python file! Received '{device_spec_path}") from e
  
          device_spec_module = importlib.import_module(import_name)
  
          # Imported file must contain `DEVICE_NAME`. If it doesn't, terminate early.
          try:
              device_name = device_spec_module.DEVICE_NAME
          except AttributeError as e:
              raise AttributeError("Device spec file did not contain `DEVICE_NAME`") from e
  
          if "TRANSFORMERS_TEST_DEVICE" in os.environ and torch_device != device_name:
              msg = f"Mismatch between environment variable `TRANSFORMERS_TEST_DEVICE` '{torch_device}' and device found in spec '{device_name}'\n"
              msg += "Either unset `TRANSFORMERS_TEST_DEVICE` or ensure it matches device spec name."
              raise ValueError(msg)
  
          torch_device = device_name
  
          def update_mapping_from_spec(device_fn_dict: dict[str, Callable], attribute_name: str):
              try:
                  # Try to import the function directly
                  spec_fn = getattr(device_spec_module, attribute_name)
                  device_fn_dict[torch_device] = spec_fn
              except AttributeError as e:
                  # If the function doesn't exist, and there is no default, throw an error
                  if "default" not in device_fn_dict:
                      raise AttributeError(
                          f"`{attribute_name}` not found in '{device_spec_path}' and no default fallback function found."
                      ) from e
  
          # Add one entry here for each `BACKEND_*` dictionary.
          update_mapping_from_spec(BACKEND_MANUAL_SEED, "MANUAL_SEED_FN")
          update_mapping_from_spec(BACKEND_EMPTY_CACHE, "EMPTY_CACHE_FN")
          update_mapping_from_spec(BACKEND_DEVICE_COUNT, "DEVICE_COUNT_FN")
  
  
  def compare_pipeline_output_to_hub_spec(output, hub_spec):
      missing_keys = []
      unexpected_keys = []
      all_field_names = {field.name for field in fields(hub_spec)}
      matching_keys = sorted([key for key in output if key in all_field_names])
  
      # Fields with a MISSING default are required and must be in the output
      for field in fields(hub_spec):
          if field.default is MISSING and field.name not in output:
              missing_keys.append(field.name)
  
      # All output keys must match either a required or optional field in the Hub spec
      for output_key in output:
          if output_key not in all_field_names:
              unexpected_keys.append(output_key)
  
      if missing_keys or unexpected_keys:
          error = ["Pipeline output does not match Hub spec!"]
          if matching_keys:
              error.append(f"Matching keys: {matching_keys}")
          if missing_keys:
              error.append(f"Missing required keys in pipeline output: {missing_keys}")
          if unexpected_keys:
              error.append(f"Keys in pipeline output that are not in Hub spec: {unexpected_keys}")
          raise KeyError("\n".join(error))
  
  
  @require_torch
  def cleanup(device: str, gc_collect=False):
      if gc_collect:
          gc.collect()
      backend_empty_cache(device)
      torch.compiler.reset()
  
  
  # Type definition of key used in `Expectations` class.
  DeviceProperties = tuple[str | None, int | None, int | None]
  # Helper type. Makes creating instances of `Expectations` smoother.
  PackedDeviceProperties = tuple[str | None, None | int | tuple[int, int]]
  
  
  @cache
  def get_device_properties() -> DeviceProperties:
      """
      Get environment device properties.
      """
      if IS_CUDA_SYSTEM or IS_ROCM_SYSTEM:
          import torch
  
          major, minor = torch.cuda.get_device_capability()
          if IS_ROCM_SYSTEM:
              return ("rocm", major, minor)
          else:
              return ("cuda", major, minor)
      elif IS_XPU_SYSTEM:
          import torch
  
          # To get more info of the architecture meaning and bit allocation, refer to https://github.com/intel/llvm/blob/sycl/sycl/include/sycl/ext/oneapi/experimental/device_architecture.def
          arch = torch.xpu.get_device_capability()["architecture"]
          gen_mask = 0x000000FF00000000
          gen = (arch & gen_mask) >> 32
          return ("xpu", gen, None)
      elif IS_NPU_SYSTEM:
          return ("npu", None, None)
      else:
          return (torch_device, None, None)
  
  
  def unpack_device_properties(
      properties: PackedDeviceProperties | None = None,
  ) -> DeviceProperties:
      """
      Unpack a `PackedDeviceProperties` tuple into consistently formatted `DeviceProperties` tuple. If properties is None, it is fetched.
      """
      if properties is None:
          return get_device_properties()
      device_type, major_minor = properties
      if major_minor is None:
          major, minor = None, None
      elif isinstance(major_minor, int):
          major, minor = major_minor, None
      else:
          major, minor = major_minor
      return device_type, major, minor
  
  
  class Expectations(UserDict[PackedDeviceProperties, Any]):
      def get_expectation(self) -> Any:
          """
          Find best matching expectation based on environment device properties. We look at device_type, major and minor
          versions of the drivers. Expectations are stored as a dictionary with keys of the form
          (device_type, (major, minor)). If the major and minor versions are not provided, we use None.
          """
          return self.find_expectation(get_device_properties())
  
      def unpacked(self) -> list[tuple[DeviceProperties, Any]]:
          return [(unpack_device_properties(k), v) for k, v in self.data.items()]
  
      @staticmethod
      def is_default(expectation_key: PackedDeviceProperties) -> bool:
          """
          This function returns True if the expectation_key is the Default expectation (None, None).
          When an Expectation dict contains a Default value, it is generally because the test existed before Expectations.
          When we modify a test to use Expectations for a specific hardware, we don't want to affect the tests on other
          hardwares. Thus we set the previous value as the Default expectation with key (None, None) and add a value for
          the specific hardware with key (hardware_type, (major, minor)).
          """
          return all(p is None for p in expectation_key)
  
      @staticmethod
      def score(properties: DeviceProperties, other: DeviceProperties) -> float:
          """
          Returns score indicating how similar two instances of the `Properties` tuple are.
          Rules are as follows:
              * Matching `type` adds one point, semi-matching `type` adds 0.1 point (e.g. cuda and rocm).
              * If types match, matching `major` adds another point, and then matching `minor` adds another.
              * The Default expectation (None, None) is worth 0.5 point, which is better than semi-matching. More on this
              in the `is_default` function.
          """
          device_type, major, minor = properties
          other_device_type, other_major, other_minor = other
  
          score = 0
          # Matching device type, maybe major and minor
          if device_type is not None and device_type == other_device_type:
              score += 1
              if major is not None and major == other_major:
                  score += 1
                  if minor is not None and minor == other_minor:
                      score += 1
          # Semi-matching device type, which carries less importance than the default expectation
          elif device_type in ["cuda", "rocm"] and other_device_type in ["cuda", "rocm"]:
              score = 0.1
  
          # Default expectation
          if Expectations.is_default(other):
              score = 0.5
  
          return score
  
      def find_expectation(self, properties: DeviceProperties = (None, None, None)) -> Any:
          """
          Find best matching expectation based on provided device properties. We score each expectation, and to
          distinguish between expectations with the same score, we use the major and minor version numbers, prioritizing
          most recent versions.
          """
          (result_key, result) = max(
              self.unpacked(),
              key=lambda x: (
                  Expectations.score(properties, x[0]),  # x[0] is a device properties tuple (device_type, major, minor)
                  x[0][1] if x[0][1] is not None else -1,  # This key is the major version, -1 if major is None
                  x[0][2] if x[0][2] is not None else -1,  # This key is the minor version, -1 if minor is None
              ),
          )
  
          if Expectations.score(properties, result_key) == 0:
              raise ValueError(f"No matching expectation found for {properties}")
  
          return result
  
      def __repr__(self):
          return f"{self.data}"
  
  
  def patch_torch_compile_force_graph():
      """
      Patch `torch.compile` to always use `fullgraph=True`.
  
      This is useful when some `torch.compile` tests are running with `fullgraph=False` and we want to be able to run
      them with `fullgraph=True` in some occasion (without introducing new tests) to make sure there is no graph break.
  
      After PR #40137, `CompileConfig.fullgraph` is `False` by default, this patch is necessary.
      """
  
      force_fullgraph = os.environ.get("TORCH_COMPILE_FORCE_FULLGRAPH", "")
      force_fullgraph = force_fullgraph.lower() in ("yes", "true", "on", "t", "y", "1")
  
      if force_fullgraph:
          import torch
  
          orig_method = torch.compile
  
          def patched(*args, **kwargs):
              # In `torch_compile`, all arguments except `model` is keyword only argument.
              kwargs["fullgraph"] = True
              return orig_method(*args, **kwargs)
  
          torch.compile = patched
  
  
  def _get_test_info():
      """
      Collect some information about the current test.
  
      For example, test full name, line number, stack, traceback, etc.
      """
  
      full_test_name = os.environ.get("PYTEST_CURRENT_TEST", "").split(" ")[0]
      test_file, test_class, test_name = full_test_name.split("::")
  
      # from the most recent frame to the top frame
      stack_from_inspect = inspect.stack()
      # but visit from the top frame to the most recent frame
  
      actual_test_file, _actual_test_class = test_file, test_class
      test_frame, test_obj, test_method = None, None, None
      for frame in reversed(stack_from_inspect):
          # if test_file in str(frame).replace(r"\\", "/"):
          # check frame's function + if it has `self` as locals; double check if self has the (function) name
          # TODO: Question: How about expanded?
          if (
              test_name.startswith(frame.function)
              and "self" in frame.frame.f_locals
              and hasattr(frame.frame.f_locals["self"], test_name)
          ):
              # if test_name == frame.frame.f_locals["self"]._testMethodName:
              test_frame = frame
              # The test instance
              test_obj = frame.frame.f_locals["self"]
              # TODO: Do we get the (relative?) path or it's just a file name?
              # TODO: Does `test_obj` always have `tearDown` object?
              actual_test_file = frame.filename
              # TODO: check `test_method` will work used at the several places!
              test_method = getattr(test_obj, test_name)
              break
  
      if test_frame is not None:
          line_number = test_frame.lineno
  
      # The frame of `patched` being called (the one and the only one calling `_get_test_info`)
      # This is used to get the original method being patched in order to get the context.
      frame_of_patched_obj = None
  
      captured_frames = []
      to_capture = False
      # From the most outer (i.e. python's `runpy.py`) frame to most inner frame (i.e. the frame of this method)
      # Between `the test method being called` and `before entering `patched``.
      for frame in reversed(stack_from_inspect):
          if (
              test_name.startswith(frame.function)
              and "self" in frame.frame.f_locals
              and hasattr(frame.frame.f_locals["self"], test_name)
          ):
              to_capture = True
          # TODO: check simply with the name is not robust.
          elif frame.frame.f_code.co_name == "patched":
              frame_of_patched_obj = frame
              to_capture = False
              break
          if to_capture:
              captured_frames.append(frame)
  
      tb_next = None
      for frame_info in reversed(captured_frames):
          tb = types.TracebackType(tb_next, frame_info.frame, frame_info.frame.f_lasti, frame_info.frame.f_lineno)
          tb_next = tb
      test_traceback = tb
  
      origin_method_being_patched = frame_of_patched_obj.frame.f_locals["orig_method"]
  
      # An iterable of type `traceback.StackSummary` with each element of type `FrameSummary`
      stack = traceback.extract_stack()
      # The frame which calls `the original method being patched`
      caller_frame = None
      # From the most inner (i.e. recent) frame to the most outer frame
      for frame in reversed(stack):
          if origin_method_being_patched.__name__ in frame.line:
              caller_frame = frame
  
      caller_path = os.path.relpath(caller_frame.filename)
      caller_lineno = caller_frame.lineno
  
      test_lineno = line_number
  
      # Get the code context in the test function/method.
      from _pytest._code.source import Source
  
      with open(actual_test_file) as fp:
          s = fp.read()
          source = Source(s)
          test_code_context = "\n".join(source.getstatement(test_lineno - 1).lines)
  
      # Get the code context in the caller (to the patched function/method).
      with open(caller_path) as fp:
          s = fp.read()
          source = Source(s)
          caller_code_context = "\n".join(source.getstatement(caller_lineno - 1).lines)
  
      test_info = f"test:\n\n{full_test_name}\n\n{'-' * 80}\n\ntest context: {actual_test_file}:{test_lineno}\n\n{test_code_context}"
      test_info = f"{test_info}\n\n{'-' * 80}\n\ncaller context: {caller_path}:{caller_lineno}\n\n{caller_code_context}"
  
      return (
          full_test_name,
          test_file,
          test_lineno,
          test_obj,
          test_method,
          test_frame,
          test_traceback,
          test_code_context,
          caller_path,
          caller_lineno,
          caller_code_context,
          test_info,
      )
  
  
  def _get_call_arguments(code_context):
      """
      Analyze the positional and keyword arguments in a call expression.
  
      This will extract the expressions of the positional and kwyword arguments, and associate them to the positions and
      the keyword argument names.
      """
  
      def get_argument_name(node):
          """Extract the name/expression from an AST node"""
          if isinstance(node, ast.Name):
              return node.id
          elif isinstance(node, ast.Attribute):
              return ast.unparse(node)
          elif isinstance(node, ast.Constant):
              return repr(node.value)
          else:
              return ast.unparse(node)
  
      indent = len(code_context) - len(code_context.lstrip())
      code_context = code_context.replace(" " * indent, "")
  
      try:
          # Parse the line
          tree = ast.parse(code_context, mode="eval")
  
          assert isinstance(tree.body, ast.Call)
          call_node = tree.body
  
          if call_node:
              result = {
                  "positional_args": [],
                  "keyword_args": {},
                  "starargs": None,  # *args
                  "kwargs": None,  # **kwargs
              }
  
              # Extract positional arguments
              for arg in call_node.args:
                  arg_name = get_argument_name(arg)
                  result["positional_args"].append(arg_name)
  
              # Extract keyword arguments
              for keyword in call_node.keywords:
                  if keyword.arg is None:
                      # This is **kwargs
                      result["kwargs"] = get_argument_name(keyword.value)
                  else:
                      # Regular keyword argument
                      arg_name = get_argument_name(keyword.value)
                      result["keyword_args"][keyword.arg] = arg_name
  
              return result
  
      except (SyntaxError, AttributeError) as e:
          print(f"Error parsing: {e}")
  
      return None
  
  
  def _prepare_debugging_info(test_info, info):
      """Combine the information about the test and the call information to a patched function/method within it."""
  
      info = f"{test_info}\n\n{info}"
      p = os.path.join(os.environ.get("_PATCHED_TESTING_METHODS_OUTPUT_DIR", ""), "captured_info.txt")
      # TODO (ydshieh): This is not safe when we use pytest-xdist with more than 1 worker.
      with open(p, "a") as fp:
          fp.write(f"{info}\n\n{'=' * 120}\n\n")
  
      return info
  
  
  def _patched_tearDown(self, *args, **kwargs):
      """Used to report a test that has failures captured and handled by patched functions/methods (without re-raise).
  
      The patched functions/methods refer to the `patched` defined in `_patch_with_call_info`, which is applied to
      `torch.testing.assert_close` and `unittest.case.TestCase.assertEqual`.
  
      The objective is to avoid a failure being silence after being processed.
  
      If there is any failure that is not handled by the patched functions/methods, we add custom error message for them
      along with the usual pytest failure report.
      """
  
      # Check for regular failures before clearing:
      # when `_patched_tearDown` is called, the current test fails due to an assertion error given by a method being
      # patched by `_patch_with_call_info`. The patched method catches such an error and continue running the remaining
      # statements within the test. If the test fails with another error not handled by the patched methods, we don't let
      # pytest to fail and report it but the original failure (the first one that was processed) instead.
      # We still record those failures not handled by the patched methods, and add custom messages along with the usual
      # pytest failure report.
      regular_failures_info = []
  
      errors = None
      if hasattr(self._outcome, "errors"):
          errors = self._outcome.errors
      elif hasattr(self._outcome, "result") and hasattr(self._outcome.result, "errors"):
          errors = self._outcome.result.errors
  
      if hasattr(self, "_outcome") and errors:
          for error_entry in errors:
              test_instance, (exc_type, exc_obj, exc_tb) = error_entry
              # breakpoint()
              regular_failures_info.append(
                  {
                      "message": f"{str(exc_obj)}\n\n",
                      "type": exc_type.__name__,
                      "file": "test_modeling_vit.py",
                      "line": 237,  # get_deepest_frame_line(exc_tb)  # Your helper function
                  }
              )
  
          # Clear the regular failure (i.e. that is not from any of our patched assertion methods) from pytest's records.
          if hasattr(self._outcome, "errors"):
              self._outcome.errors.clear()
          elif hasattr(self._outcome, "result") and hasattr(self._outcome.result, "errors"):
              self._outcome.result.errors.clear()
  
      # reset back to the original tearDown method, so `_patched_tearDown` won't be run by the subsequent tests if they
      # have only test failures that are not handle by the patched methods (or no test failure at all).
      orig_tearDown = _patched_tearDown.orig_tearDown
      type(self).tearDown = orig_tearDown
  
      # Call the original tearDown
      orig_tearDown(self, *args, **kwargs)
  
      # Get the failure
      test_method = getattr(self, self._testMethodName)
      captured_failures = test_method.__func__.captured_failures[id(test_method)]
  
      # TODO: How could we show several exceptions in a sinigle test on the terminal? (Maybe not a good idea)
      captured_exceptions = captured_failures[0]["exception"]
      captured_traceback = captured_failures[0]["traceback"]
      # Show the captured information on the terminal.
      capturued_info = [x["info"] for x in captured_failures]
      capturued_info_str = f"\n\n{'=' * 80}\n\n".join(capturued_info)
  
      # Enhance the exception message if there were suppressed failures
      if regular_failures_info:
          enhanced_message = f"""{str(captured_exceptions)}
  
  {"=" * 80}
  Handled Failures: ({len(capturued_info)} handled):
  {"-" * 80}\n
  {capturued_info_str}
  
  {"=" * 80}
  Unhandled Failures: ({len(regular_failures_info)} unhandled):
  {"-" * 80}\n
  {", ".join(f"{info['type']}: {info['message']}{info['file']}:{info['line']}" for info in regular_failures_info)}
  
  {"-" * 80}
  Note: This failure occurred after other failures analyzed by the patched assertion methods.
  To see the full details, temporarily disable assertion patching.
  {"=" * 80}"""
  
          # Create new exception with enhanced message
          enhanced_exception = type(captured_exceptions)(enhanced_message)
          enhanced_exception.__cause__ = captured_exceptions.__cause__
          enhanced_exception.__context__ = captured_exceptions.__context__
  
          # Raise with your existing traceback reconstruction
          captured_exceptions = enhanced_exception
  
      # clean up the recorded status
      del test_method.__func__.captured_failures
  
      raise captured_exceptions.with_traceback(captured_traceback)
  
  
  def _patch_with_call_info(module_or_class, attr_name, _parse_call_info_func, target_args):
      """
      Patch a callerable `attr_name` of a module or class `module_or_class`.
  
      This will allow us to collect the call information, e.g. the argument names and values, also the literal expressions
      passed as the arguments.
      """
      orig_method = getattr(module_or_class, attr_name)
      if not callable(orig_method):
          return
  
      def patched(*args, **kwargs):
          # If the target callable is not called within a test, simply call it without modification.
          if not os.environ.get("PYTEST_CURRENT_TEST", ""):
              return orig_method(*args, **kwargs)
  
          try:
              orig_method(*args, **kwargs)
          except AssertionError as e:
              captured_exception = e
              # captured_traceback = e.__traceback__
              (
                  full_test_name,
                  test_file,
                  test_lineno,
                  test_obj,
                  test_method,
                  test_frame,
                  test_traceback,
                  test_code_context,
                  caller_path,
                  caller_lineno,
                  caller_code_context,
                  test_info,
              ) = _get_test_info()
              test_info = f"{test_info}\n\n{'-' * 80}\n\npatched method: {orig_method.__module__}.{orig_method.__name__}"
              call_argument_expressions = _get_call_arguments(caller_code_context)
  
              # This is specific
              info = _parse_call_info_func(orig_method, args, kwargs, call_argument_expressions, target_args)
              info = _prepare_debugging_info(test_info, info)
  
              # If the test is running in a CI environment (e.g. not a manual run), let's raise and fail the test, so it
              # behaves as usual.
              # On Github Actions or CircleCI, this is set automatically.
              # When running manually, it's the user to determine if to set it.
              # This is to avoid the patched function being called `with self.assertRaises(AssertionError):` and fails
              # because of the missing expected `AssertionError`.
              # TODO (ydshieh): If there is way to raise only when we are inside such context managers?
              # TODO (ydshieh): How not to record the failure if it happens inside `self.assertRaises(AssertionError)`?
              if os.getenv("CI") == "true":
                  raise captured_exception.with_traceback(test_traceback)
  
              # Save this, so we can raise at the end of the current test
              captured_failure = {
                  "result": "failed",
                  "exception": captured_exception,
                  "traceback": test_traceback,
                  "info": info,
              }
  
              # Record the failure status and its information, so we can raise it later.
              # We are modifying the (unbound) function at class level: not its logic but only adding a new extra
              # attribute.
              if getattr(test_method.__func__, "captured_failures", None) is None:
                  test_method.__func__.captured_failures = {}
              if id(test_method) not in test_method.__func__.captured_failures:
                  test_method.__func__.captured_failures[id(test_method)] = []
              test_method.__func__.captured_failures[id(test_method)].append(captured_failure)
  
              # This modifies the `tearDown` which will be called after every tests, but we reset it back inside
              # `_patched_tearDown`.
              if not hasattr(type(test_obj).tearDown, "orig_tearDown"):
                  orig_tearDown = type(test_obj).tearDown
                  _patched_tearDown.orig_tearDown = orig_tearDown
                  type(test_obj).tearDown = _patched_tearDown
  
      setattr(module_or_class, attr_name, patched)
  
  
  def _parse_call_info(func, args, kwargs, call_argument_expressions, target_args):
      """
      Prepare a string containing the call info to `func`, e.g. argument names/values/expressions.
      """
      signature = inspect.signature(func)
      signature_names = [param.name for param_name, param in signature.parameters.items()]
  
      # called as `self.method_name()` or `xxx.method_name()`.
      if len(args) == len(call_argument_expressions["positional_args"]) + 1:
          # We simply add "self" as the expression despite it might not be the actual argument name.
          # (This part is very unlikely what a user would be interest to know)
          call_argument_expressions["positional_args"] = ["self"] + call_argument_expressions["positional_args"]
  
      param_position_mapping = {param_name: idx for idx, param_name in enumerate(signature_names)}
  
      arg_info = {}
      for arg_name in target_args:
          if arg_name in kwargs:
              arg_value = kwargs[arg_name]
              arg_expr = call_argument_expressions["keyword_args"][arg_name]
          else:
              arg_pos = param_position_mapping[arg_name]
              arg_value = args[arg_pos]
              arg_expr = call_argument_expressions["positional_args"][arg_pos]
  
          arg_value_str = _format_py_obj(arg_value)
          arg_info[arg_name] = {"arg_expr": arg_expr, "arg_value_str": arg_value_str}
  
      info = ""
      for arg_name in arg_info:
          arg_expr, arg_value_str = arg_info[arg_name]["arg_expr"], arg_info[arg_name]["arg_value_str"]
          info += f"{'-' * 80}\n\nargument name: `{arg_name}`\nargument expression: `{arg_expr}`\n\nargument value:\n\n{arg_value_str}\n\n"
  
      # remove the trailing \n\n
      info = info[:-2]
  
      return info
  
  
  def patch_testing_methods_to_collect_info():
      """
      Patch some methods (`torch.testing.assert_close`, `unittest.case.TestCase.assertEqual`, etc).
  
      This will allow us to collect the call information, e.g. the argument names and values, also the literal expressions
      passed as the arguments.
      """
      p = os.path.join(os.environ.get("_PATCHED_TESTING_METHODS_OUTPUT_DIR", ""), "captured_info.txt")
      Path(p).unlink(missing_ok=True)
  
      if is_torch_available():
          import torch
  
          _patch_with_call_info(torch.testing, "assert_close", _parse_call_info, target_args=("actual", "expected"))
  
      _patch_with_call_info(unittest.case.TestCase, "assertEqual", _parse_call_info, target_args=("first", "second"))
      _patch_with_call_info(unittest.case.TestCase, "assertListEqual", _parse_call_info, target_args=("list1", "list2"))
      _patch_with_call_info(
          unittest.case.TestCase, "assertTupleEqual", _parse_call_info, target_args=("tuple1", "tuple2")
      )
      _patch_with_call_info(unittest.case.TestCase, "assertSetEqual", _parse_call_info, target_args=("set1", "set1"))
      _patch_with_call_info(unittest.case.TestCase, "assertDictEqual", _parse_call_info, target_args=("d1", "d2"))
      _patch_with_call_info(unittest.case.TestCase, "assertIn", _parse_call_info, target_args=("member", "container"))
      _patch_with_call_info(unittest.case.TestCase, "assertNotIn", _parse_call_info, target_args=("member", "container"))
      _patch_with_call_info(unittest.case.TestCase, "assertLess", _parse_call_info, target_args=("a", "b"))
      _patch_with_call_info(unittest.case.TestCase, "assertLessEqual", _parse_call_info, target_args=("a", "b"))
      _patch_with_call_info(unittest.case.TestCase, "assertGreater", _parse_call_info, target_args=("a", "b"))
      _patch_with_call_info(unittest.case.TestCase, "assertGreaterEqual", _parse_call_info, target_args=("a", "b"))
  
  
  def torchrun(script: str, nproc_per_node: int, is_torchrun: bool = True, env: dict | None = None):
      """Run the `script` using `torchrun` command for multi-processing in a subprocess. Captures errors as necessary."""
      with tempfile.NamedTemporaryFile(mode="w+", suffix=".py") as tmp:
          tmp.write(script)
          tmp.flush()
          tmp.seek(0)
          if is_torchrun:
              cmd = (
                  f"torchrun --nproc_per_node {nproc_per_node} --master_port {get_torch_dist_unique_port()} {tmp.name}"
              ).split()
          else:
              cmd = ["python3", tmp.name]
  
          # Note that the subprocess will be waited for here, and raise an error if not successful
          try:
              _ = subprocess.run(cmd, capture_output=True, env=env, text=True, check=True)
          except subprocess.CalledProcessError as e:
              raise Exception(f"The following error was captured: {e.stderr}")
  
  
  def _format_tensor(t, indent_level=0, sci_mode=None):
      """Format torch's tensor in a pretty way to be shown 👀 in the test report."""
  
      # `torch.testing.assert_close` could accept python int/float numbers.
      if not isinstance(t, torch.Tensor):
          t = torch.tensor(t)
  
      # Simply make the processing below simpler (not to handle both cases)
      is_scalar = False
      if t.ndim == 0:
          t = torch.tensor([t])
          is_scalar = True
  
      # For scalar or one-dimensional tensor, keep it as one-line. If there is only one element along any dimension except
      # the last one, we also keep it as one-line.
      if t.ndim <= 1 or set(t.shape[0:-1]) == {1}:
          # Use `detach` to remove `grad_fn=<...>`, and use `to("cpu")` to remove `device='...'`
          t = t.detach().to("cpu")
  
          # We work directly with the string representation instead the tensor itself
          t_str = str(t)
  
          # remove `tensor( ... )` so keep only the content
          t_str = t_str.replace("tensor(", "").replace(")", "")
  
          # Sometimes there are extra spaces between `[` and the first digit of the first value (for alignment).
          # For example `[[ 0.06, -0.51], [-0.76, -0.49]]`. It may have multiple consecutive spaces.
          # Let's remove such extra spaces.
          while "[ " in t_str:
              t_str = t_str.replace("[ ", "[")
  
          # Put everything in a single line. We replace `\n` by a space ` ` so we still keep `,\n` as `, `.
          t_str = t_str.replace("\n", " ")
  
          # Remove repeated spaces (introduced by the previous step)
          while "  " in t_str:
              t_str = t_str.replace("  ", " ")
  
          # remove leading `[` and `]` for scalar tensor
          if is_scalar:
              t_str = t_str[1:-1]
  
          t_str = " " * 4 * indent_level + t_str
  
          return t_str
  
      # Otherwise, we separate the representations of each element along an outer dimension by new lines (after a `,`).
      # The representation of each element is obtained by calling this function recursively with current `indent_level`.
      else:
          t_str = str(t)
  
          # (For the recursive calls should receive this value)
          if sci_mode is None:
              sci_mode = "e+" in t_str or "e-" in t_str
  
          # Use the original content to determine the scientific mode to use. This is required as the representation of
          # t[index] (computed below) maybe have different format regarding scientific notation.
          torch.set_printoptions(sci_mode=sci_mode)
  
          t_str = " " * 4 * indent_level + "[\n"
          # Keep the ending `,` for all outer dimensions whose representations are not put in one-line, even if there is
          # only one element along that dimension.
          t_str += ",\n".join(_format_tensor(x, indent_level=indent_level + 1, sci_mode=sci_mode) for x in t)
          t_str += ",\n" + " " * 4 * indent_level + "]"
  
          torch.set_printoptions(sci_mode=None)
  
      return t_str
  
  
  def _quote_string(s):
      """Given a string `s`, return a python literal expression that give `s` when it is used in a python source code.
  
      For example, if `s` is the string `abc`, the return value is `"abc"`.
  
      We choice double quotes over single quote despite `str(s)` would give `'abc'` instead of `"abc"`.
      """
      has_single_quote = "'" in s
      has_double_quote = '"' in s
  
      if has_single_quote and has_double_quote:
          # replace any double quote by the raw string r'\"'.
          s = s.replace('"', r"\"")
          return f'"{s}"'
      elif has_single_quote:
          return f'"{s}"'
      elif has_double_quote:
          return f"'{s}'"
      else:
          return f'"{s}"'
  
  
  def _format_py_obj(obj, indent=0, mode="", cache=None, prefix=""):
      """Format python objects of basic built-in type in a pretty way so we could copy-past them to code editor easily.
  
      Currently, this support int, float, str, list, tuple, and dict.
  
      It also works with `torch.Tensor` via calling `format_tesnor`.
      """
  
      if cache is None:
          cache = {}
      else:
          if (id(obj), indent, mode, prefix) in cache:
              return cache[(id(obj), indent, mode, prefix)]
  
      # special format method for `torch.Tensor`
      if str(obj.__class__) == "<class 'torch.Tensor'>":
          return _format_tensor(obj)
  
      elif obj.__class__.__name__ == "str":
          quoted_string = _quote_string(obj)
          # we don't want the newline being interpreted
          quoted_string = quoted_string.replace("\n", r"\n")
          output = quoted_string
  
      elif obj.__class__.__name__ in ["int", "float"]:
          # for float like `1/3`, we will get `0.3333333333333333`
          output = str(obj)
  
      elif obj.__class__.__name__ in ["list", "tuple", "dict"]:
          parenthesis = {
              "list": "[]",
              "tuple": "()",
              "dict": "{}",
          }
          p1, p2 = parenthesis[obj.__class__.__name__]
  
          elements_without_indent = []
          if isinstance(obj, dict):
              for idx, (k, v) in enumerate(obj.items()):
                  last_element = idx == len(obj) - 1
                  ok = _format_py_obj(k, indent=indent + 1, mode="one-line", cache=cache)
                  ov = _format_py_obj(
                      v,
                      indent=indent + 1,
                      mode=mode,
                      cache=cache,
                      prefix=ok.lstrip() + ": " + "," if not last_element else "",
                  )
                  # Each element could be multiple-line, but the indent of its first line is removed
                  elements_without_indent.append(f"{ok.lstrip()}: {ov.lstrip()}")
  
          else:
              for idx, x in enumerate(obj):
                  last_element = idx == len(obj) - 1
                  o = _format_py_obj(
                      x, indent=indent + 1, mode=mode, cache=cache, prefix="," if not last_element else ""
                  )
                  # Each element could be multiple-line, but the indent of its first line is removed
                  elements_without_indent.append(o.lstrip())
  
          groups = []
          buf = []
          for idx, x in enumerate(elements_without_indent):
              buf.append(x)
  
              x_expanded = "\n" in buf[-1]
              not_last_element = idx != len(elements_without_indent) - 1
              # if `x` should be separated from subsequent elements
              should_finalize_x = x_expanded or len(f"{' ' * (4 * (indent + 1))}") + len(
                  ", ".join(buf[-1:])
              ) > 120 - int(not_last_element)
  
              # if `buf[:-1]` (i.e. without `x`) should be combined together (into one line)
              should_finalize_buf = x_expanded
  
              # the recursive call returns single line, so we can use it to determine if we can fit the width limit
              if not should_finalize_buf:
                  buf_not_fit_into_one_line = len(f"{' ' * (4 * (indent + 1))}") + len(", ".join(buf)) > 120 - int(
                      not_last_element
                  )
                  should_finalize_buf = buf_not_fit_into_one_line
  
              # any element of iterable type need to be on its own line
              if (type(obj[idx]) if type(obj) is not dict else type(list(obj.values())[idx])) in [list, tuple, dict]:
                  should_finalize_x = True
                  should_finalize_buf = True
  
              # any type change --> need to be added after a new line
              prev_type = None
              current_type = type(obj[idx]) if type(obj) is not dict else type(list(obj.values())[idx])
              if len(buf) > 1:
                  prev_type = type(obj[idx - 1]) if type(obj) is not dict else type(list(obj.values())[idx - 1])
                  type_changed = current_type != prev_type
                  if type_changed:
                      should_finalize_buf = True
  
              # all elements in the buf are string --> don't finalize the buf by width limit
              if prev_type is None or (prev_type is str and current_type is str):
                  should_finalize_buf = False
  
              # collect as many elements of string type as possible (without width limit).
              # These will be examined as a whole (if not fit into the width, each element would be in its own line)
              if current_type is str:
                  should_finalize_x = False
                  # `len(buf) == 1` or `obj[idx-1]` is a string
                  if prev_type in [None, str]:
                      should_finalize_buf = False
  
              if should_finalize_buf:
                  orig_buf_len = len(buf)
  
                  if orig_buf_len > 1:
                      not_fit_into_one_line = None
  
                      # all elements in `obj` that give `buf[:-1]` are string.
                      if prev_type is str:
                          # `-1` at the end: because buf[-2] is not the last element
                          not_fit_into_one_line = len(f"{' ' * (4 * (indent + 1))}") + len(", ".join(buf[:-1])) > 120 - 1
  
                      if not_fit_into_one_line:
                          for x in buf[:-1]:
                              groups.append([x])
                      else:
                          groups.append(buf[:-1])
  
                      buf = buf[-1:]
  
                  if should_finalize_x:
                      groups.append(buf)
                      buf = []
  
          # The last buf
          if len(buf) > 0:
              not_fit_into_one_line = None
              if current_type is str:
                  # no `-1` at the end: because buf[-1] is the last element
                  not_fit_into_one_line = len(f"{' ' * (4 * (indent + 1))}") + len(", ".join(buf)) > 120
  
              if not_fit_into_one_line:
                  for x in buf:
                      groups.append([x])
              else:
                  groups.append(buf)
  
          output = f"{' ' * 4 * indent}{p1}\n"
          element_strings = [f"{' ' * (4 * (indent + 1))}" + ", ".join(buf) for buf in groups]
          output += ",\n".join(element_strings)
          output += f"\n{' ' * 4 * indent}{p2}"
  
          # if all elements are in one-line
          no_new_line_in_elements = all("\n" not in x for x in element_strings)
          # if yes, we can form a one-line representation of `obj`
          could_use_one_line = no_new_line_in_elements
  
          # if mode == "one-line", this function always returns one-line representation, so `no_new_line_in_elements`
          # will be `True`.
          if could_use_one_line:
              one_line_form = ", ".join([x.lstrip() for x in element_strings])
              one_line_form = f"{p1}{one_line_form}{p2}"
  
              if mode == "one-line":
                  return output
  
              # check with the width limit
              could_use_one_line = len(f"{' ' * 4 * indent}") + len(prefix) + len(one_line_form) <= 120
  
              # extra conditions for returning one-line representation
              def use_one_line_repr(obj):
                  # iterable types
                  if type(obj) in (list, tuple, dict):
                      # get all types
                      element_types = []
                      if type(obj) is dict:
                          element_types.extend(type(x) for x in obj.values())
                      elif type(obj) in [list, tuple]:
                          element_types.extend(type(x) for x in obj)
  
                      # At least one element is of iterable type
                      if any(x in (list, tuple, dict) for x in element_types):
                          # If `obj` has more than one element and at least one of them is iterable --> no one line repr.
                          if len(obj) > 1:
                              return False
  
                          # only one element that is iterable, but not the same type as `obj` --> no one line repr.
                          if type(obj) is not type(obj[0]):
                              return False
  
                          # one-line repr. if possible, without width limit
                          return no_new_line_in_elements
  
                      # all elements are of simple types, but more than one type --> no one line repr.
                      if len(set(element_types)) > 1:
                          return False
  
                      # all elements are of the same simple type
                      if element_types[0] in [int, float]:
                          # one-line repr. without width limit
                          return no_new_line_in_elements
                      elif element_types[0] is str:
                          if len(obj) == 1:
                              # one single string element --> one-line repr. without width limit
                              return no_new_line_in_elements
                          else:
                              # multiple string elements --> one-line repr. if fit into width limit
                              return could_use_one_line
  
                  # simple types (int, flat, string)
                  return True
  
              # width condition combined with specific mode conditions
              if use_one_line_repr(obj):
                  output = f"{' ' * 4 * indent}{one_line_form}"
  
      cache[(id(obj), indent, mode, prefix)] = output
  
      return output
  
  
  def write_file(file, content):
      with open(file, "w") as f:
          f.write(content)
  
  
  def read_json_file(file):
      with open(file, "r") as fh:
          return json.load(fh)
  
  
  # =============================================================================
  # Training CI Utilities - Logging and Memory Monitoring
  # =============================================================================
  
  
  # ANSI color codes for terminal output
  class Colors:
      """ANSI color codes for terminal output formatting."""
  
      RESET = "\033[0m"
      BOLD = "\033[1m"
      DIM = "\033[2m"
  
      # Foreground colors
      RED = "\033[31m"
      GREEN = "\033[32m"
      YELLOW = "\033[33m"
      BLUE = "\033[34m"
      MAGENTA = "\033[35m"
      CYAN = "\033[36m"
      WHITE = "\033[37m"
  
      # Bright variants
      BRIGHT_RED = "\033[91m"
      BRIGHT_GREEN = "\033[92m"
      BRIGHT_YELLOW = "\033[93m"
      BRIGHT_BLUE = "\033[94m"
      BRIGHT_CYAN = "\033[96m"
  
  
  class ColoredFormatter(logging.Formatter):
      """Custom formatter that adds colors based on log level."""
  
      LEVEL_COLORS = {
          logging.DEBUG: Colors.DIM + Colors.CYAN,
          logging.INFO: Colors.WHITE,
          logging.WARNING: Colors.BRIGHT_YELLOW,
          logging.ERROR: Colors.BRIGHT_RED,
          logging.CRITICAL: Colors.BOLD + Colors.BRIGHT_RED,
      }
  
      # Loggers that should be dimmed (less important/verbose)
      DIMMED_LOGGERS = {"httpx", "httpcore", "urllib3", "requests"}
  
      def __init__(self, fmt: str | None = None, datefmt: str | None = None):
          super().__init__(fmt, datefmt)
  
      def format(self, record: logging.LogRecord) -> str:
          # Check if this logger should be dimmed
          is_dimmed = record.name in self.DIMMED_LOGGERS
  
          if is_dimmed:
              # Dim the entire log line for httpx and similar
              timestamp = self.formatTime(record, self.datefmt)
              message = record.getMessage()
              return f"{Colors.DIM}{timestamp} - {record.name} - {record.levelname:8} - {message}{Colors.RESET}"
  
          # Get color for this level
          color = self.LEVEL_COLORS.get(record.levelno, Colors.RESET)
  
          # Color the level name
          levelname = record.levelname
          colored_levelname = f"{color}{levelname:8}{Colors.RESET}"
  
          # Color the timestamp
          colored_time = f"{Colors.DIM}{self.formatTime(record, self.datefmt)}{Colors.RESET}"
  
          # Color the logger name
          colored_name = f"{Colors.BLUE}{record.name}{Colors.RESET}"
  
          # Get message
          message = record.getMessage()
  
          return f"{colored_time} - {colored_name} - {colored_levelname} - {message}"
  
  
  _warn_once_logged: set[str] = set()
  
  
  def init_test_logger() -> logging.Logger:
      """Initialize a test-specific logger with colored stderr handler and INFO level for tests.
  
      Uses a named logger instead of root logger to avoid conflicts with pytest-xdist parallel execution.
      Uses stderr instead of stdout to avoid deadlocks with pytest-xdist output capture.
      """
      logger = logging.getLogger("transformers.training_test")
      logger.setLevel(logging.INFO)
  
      # Only add handler if not already present (avoid duplicate handlers on repeated calls)
      if not logger.handlers:
          # Use stderr instead of stdout - pytest-xdist captures stdout which can cause deadlocks
          ch = logging.StreamHandler(sys.stderr)
          ch.setLevel(logging.INFO)
  
          # Use colored formatter if terminal supports it, plain otherwise
          if sys.stderr.isatty():
              formatter = ColoredFormatter(datefmt="%Y-%m-%d %H:%M:%S")
          else:
              formatter = logging.Formatter(
                  "%(asctime)s - %(name)s - %(levelname)s - %(message)s", datefmt="%Y-%m-%d %H:%M:%S"
              )
  
          ch.setFormatter(formatter)
          logger.addHandler(ch)
  
      logger.propagate = False  # Don't propagate to root logger to avoid duplicate output
      return logger
  
  
  def warn_once(logger_instance: logging.Logger, msg: str) -> None:
      """Log a warning message only once per unique message.
  
      Uses a global set to track messages that have already been logged
      to prevent duplicate warning messages from cluttering the output.
  
      Args:
          logger_instance: The logger instance to use for warning.
          msg: The warning message to log.
      """
      if msg not in _warn_once_logged:
          logger_instance.warning(msg)
          _warn_once_logged.add(msg)
  
  
  # Named tuple for passing memory stats for logging
  MemoryStats = collections.namedtuple(
      "MemoryStats",
      [
          "rss_gib",  # Resident Set Size in GiB
          "rss_pct",  # RSS as percentage of total memory
          "vms_gib",  # Virtual Memory Size in GiB
          "peak_rss_gib",  # Peak RSS in GiB
          "peak_rss_pct",  # Peak RSS as percentage of total memory
          "available_gib",  # Available system memory in GiB
          "total_gib",  # Total system memory in GiB
      ],
  )
  
  
  class CPUMemoryMonitor:
      """Monitor CPU memory usage for the current process."""
  
      def __init__(self):
          self.device_name = "CPU"
          self._peak_rss = 0
          self._process = None
          self.total_memory = 0
          self.total_memory_gib = 0
  
          if is_psutil_available():
              import psutil
  
              self._process = psutil.Process(os.getpid())
              mem_info = psutil.virtual_memory()
              self.total_memory = mem_info.total
              self.total_memory_gib = self._to_gib(self.total_memory)
  
      def _to_gib(self, memory_in_bytes: int) -> float:
          """Convert bytes to GiB."""
          return memory_in_bytes / (1024 * 1024 * 1024)
  
      def _to_pct(self, memory_in_bytes: int) -> float:
          """Convert bytes to percentage of total memory."""
          if self.total_memory == 0:
              return 0.0
          return 100.0 * memory_in_bytes / self.total_memory
  
      def _update_peak(self) -> None:
          """Update peak memory tracking."""
          if self._process is not None:
              current_rss = self._process.memory_info().rss
              self._peak_rss = max(self._peak_rss, current_rss)
  
      def get_stats(self) -> MemoryStats:
          """Get current memory statistics."""
          if not is_psutil_available():
              return MemoryStats(0, 0, 0, 0, 0, 0, 0)
  
          import psutil
  
          self._update_peak()
  
          mem_info = self._process.memory_info()
          sys_mem = psutil.virtual_memory()
  
          return MemoryStats(
              rss_gib=self._to_gib(mem_info.rss),
              rss_pct=self._to_pct(mem_info.rss),
              vms_gib=self._to_gib(mem_info.vms),
              peak_rss_gib=self._to_gib(self._peak_rss),
              peak_rss_pct=self._to_pct(self._peak_rss),
              available_gib=self._to_gib(sys_mem.available),
              total_gib=self._to_gib(sys_mem.total),
          )
  
      def reset_peak_stats(self) -> None:
          """Reset peak memory tracking."""
          if self._process is not None:
              self._peak_rss = self._process.memory_info().rss
  
  
  def build_cpu_memory_monitor(logger_instance: logging.Logger | None = None) -> CPUMemoryMonitor:
      """Build and initialize a CPU memory monitor.
  
      Args:
          logger_instance: Optional logger to log initialization info. If None, no logging is done.
  
      Returns:
          CPUMemoryMonitor instance.
      """
      monitor = CPUMemoryMonitor()
      if logger_instance is not None:
          if is_psutil_available():
              logger_instance.info(f"CPU memory monitor initialized: {monitor.total_memory_gib:.2f} GiB total")
          else:
              logger_instance.warning("psutil not available, memory monitoring disabled")
      return monitor
  
  
  def convert_all_safetensors_to_bins(folder: str):
      """Convert all safetensors files into torch bin files, to mimic saving with torch (since we still support loading
      bin files, but not saving them anymore)"""
      for file in os.listdir(folder):
          path = os.path.join(folder, file)
          if file.endswith(".safetensors"):
              new_path = path.replace(".safetensors", ".bin").replace("model", "pytorch_model")
              state_dict = load_file(path)
              os.remove(path)
              torch.save(state_dict, new_path)
          # Adapt the index as well
          elif file == SAFE_WEIGHTS_INDEX_NAME:
              new_path = os.path.join(folder, WEIGHTS_INDEX_NAME)
              with open(path) as f:
                  index = json.loads(f.read())
              os.remove(path)
              if "weight_map" in index.keys():
                  weight_map = index["weight_map"]
                  new_weight_map = {}
                  for k, v in weight_map.items():
                      new_weight_map[k] = v.replace(".safetensors", ".bin").replace("model", "pytorch_model")
              index["weight_map"] = new_weight_map
              with open(new_path, "w") as f:
                  f.write(json.dumps(index, indent=4))
  
  
  @contextmanager
  def force_serialization_as_bin_files():
      """Since we don't support saving with torch `.bin` files anymore, but still support loading them, we use this context
      to easily create the bin files and try to load them back"""
      try:
          # Monkey patch the method to save as bin files
          original_save = PreTrainedModel.save_pretrained
  
          def new_save(self, save_directory, *args, **kwargs):
              original_save(self, save_directory, *args, **kwargs)
              convert_all_safetensors_to_bins(save_directory)
  
          PreTrainedModel.save_pretrained = new_save
  
          yield
      finally:
          PreTrainedModel.save_pretrained = original_save
  