"""simple docstring"""


from __future__ import annotations

import copy
import tempfile
import unittest

from transformers import CONFIG_MAPPING, AutoConfig, BertConfig, GPTaConfig, TaConfig, TapasConfig, is_tf_available
from transformers.testing_utils import (
    DUMMY_UNKNOWN_IDENTIFIER,
    SMALL_MODEL_IDENTIFIER,
    RequestCounter,
    require_tensorflow_probability,
    require_tf,
    slow,
)

from ..bert.test_modeling_bert import BertModelTester


if is_tf_available():
    from transformers import (
        TFAutoModel,
        TFAutoModelForCausalLM,
        TFAutoModelForMaskedLM,
        TFAutoModelForPreTraining,
        TFAutoModelForQuestionAnswering,
        TFAutoModelForSeqaSeqLM,
        TFAutoModelForSequenceClassification,
        TFAutoModelForTableQuestionAnswering,
        TFAutoModelForTokenClassification,
        TFAutoModelWithLMHead,
        TFBertForMaskedLM,
        TFBertForPreTraining,
        TFBertForQuestionAnswering,
        TFBertForSequenceClassification,
        TFBertModel,
        TFFunnelBaseModel,
        TFFunnelModel,
        TFGPTaLMHeadModel,
        TFRobertaForMaskedLM,
        TFTaForConditionalGeneration,
        TFTapasForQuestionAnswering,
    )
    from transformers.models.auto.modeling_tf_auto import (
        TF_MODEL_FOR_CAUSAL_LM_MAPPING,
        TF_MODEL_FOR_MASKED_LM_MAPPING,
        TF_MODEL_FOR_PRETRAINING_MAPPING,
        TF_MODEL_FOR_QUESTION_ANSWERING_MAPPING,
        TF_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING,
        TF_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING,
        TF_MODEL_MAPPING,
    )
    from transformers.models.bert.modeling_tf_bert import TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST
    from transformers.models.gpta.modeling_tf_gpta import TF_GPT2_PRETRAINED_MODEL_ARCHIVE_LIST
    from transformers.models.ta.modeling_tf_ta import TF_T5_PRETRAINED_MODEL_ARCHIVE_LIST
    from transformers.models.tapas.modeling_tf_tapas import TF_TAPAS_PRETRAINED_MODEL_ARCHIVE_LIST


class 				a (						lowerCAmelCase_      ):
    _snake_case   :     Optional[int]   							=							'new-model'


if is_tf_available():


    class 				a (						lowerCAmelCase_      ):
        _snake_case   :     Any   							=							NewModelConfig


@require_tf
class 				a (						unittest.TestCase      ):


    @slow
    def    lowerCAmelCase_     (    self		:     List[Any]    ):
      _UpperCAmelCase         =					"""bert-base-cased"""
      _UpperCAmelCase         =					AutoConfig.from_pretrained(__lowerCAmelCase    )
      self.assertIsNotNone(__lowerCAmelCase    )
      self.assertIsInstance(__lowerCAmelCase      , __lowerCAmelCase    )

      _UpperCAmelCase         =					TFAutoModel.from_pretrained(__lowerCAmelCase    )
      self.assertIsNotNone(__lowerCAmelCase    )
      self.assertIsInstance(__lowerCAmelCase      , __lowerCAmelCase    )


    @slow
    def    lowerCAmelCase_     (    self		:     str    ):
      _UpperCAmelCase         =					"""bert-base-cased"""
      _UpperCAmelCase         =					AutoConfig.from_pretrained(__lowerCAmelCase    )
      self.assertIsNotNone(__lowerCAmelCase    )
      self.assertIsInstance(__lowerCAmelCase      , __lowerCAmelCase    )

      _UpperCAmelCase         =					TFAutoModelForPreTraining.from_pretrained(__lowerCAmelCase    )
      self.assertIsNotNone(__lowerCAmelCase    )
      self.assertIsInstance(__lowerCAmelCase      , __lowerCAmelCase    )


    @slow
    def    lowerCAmelCase_     (    self		:     Any    ):
      for model_name in TF_GPT2_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
        _UpperCAmelCase         =					AutoConfig.from_pretrained(__lowerCAmelCase    )
        self.assertIsNotNone(__lowerCAmelCase    )
        self.assertIsInstance(__lowerCAmelCase      , __lowerCAmelCase    )

        _UpperCAmelCase         =					TFAutoModelForCausalLM.from_pretrained(__lowerCAmelCase    )
        _UpperCAmelCase       , _UpperCAmelCase         =					TFAutoModelForCausalLM.from_pretrained(__lowerCAmelCase      , output_loading_info=__lowerCAmelCase    )
        self.assertIsNotNone(__lowerCAmelCase    )
        self.assertIsInstance(__lowerCAmelCase      , __lowerCAmelCase    )


    @slow
    def    lowerCAmelCase_     (    self		:     Dict    ):
      for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
        _UpperCAmelCase         =					AutoConfig.from_pretrained(__lowerCAmelCase    )
        self.assertIsNotNone(__lowerCAmelCase    )
        self.assertIsInstance(__lowerCAmelCase      , __lowerCAmelCase    )

        _UpperCAmelCase         =					TFAutoModelWithLMHead.from_pretrained(__lowerCAmelCase    )
        self.assertIsNotNone(__lowerCAmelCase    )
        self.assertIsInstance(__lowerCAmelCase      , __lowerCAmelCase    )


    @slow
    def    lowerCAmelCase_     (    self		:     Union[str, Any]    ):
      for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
        _UpperCAmelCase         =					AutoConfig.from_pretrained(__lowerCAmelCase    )
        self.assertIsNotNone(__lowerCAmelCase    )
        self.assertIsInstance(__lowerCAmelCase      , __lowerCAmelCase    )

        _UpperCAmelCase         =					TFAutoModelForMaskedLM.from_pretrained(__lowerCAmelCase    )
        _UpperCAmelCase       , _UpperCAmelCase         =					TFAutoModelForMaskedLM.from_pretrained(__lowerCAmelCase      , output_loading_info=__lowerCAmelCase    )
        self.assertIsNotNone(__lowerCAmelCase    )
        self.assertIsInstance(__lowerCAmelCase      , __lowerCAmelCase    )


    @slow
    def    lowerCAmelCase_     (    self		:     Union[str, Any]    ):
      for model_name in TF_T5_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
        _UpperCAmelCase         =					AutoConfig.from_pretrained(__lowerCAmelCase    )
        self.assertIsNotNone(__lowerCAmelCase    )
        self.assertIsInstance(__lowerCAmelCase      , __lowerCAmelCase    )

        _UpperCAmelCase         =					TFAutoModelForSeqaSeqLM.from_pretrained(__lowerCAmelCase    )
        _UpperCAmelCase       , _UpperCAmelCase         =					TFAutoModelForSeqaSeqLM.from_pretrained(__lowerCAmelCase      , output_loading_info=__lowerCAmelCase    )
        self.assertIsNotNone(__lowerCAmelCase    )
        self.assertIsInstance(__lowerCAmelCase      , __lowerCAmelCase    )


    @slow
    def    lowerCAmelCase_     (    self		:     List[str]    ):
      # for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
      for model_name in ["bert-base-uncased"]:
        _UpperCAmelCase         =					AutoConfig.from_pretrained(__lowerCAmelCase    )
        self.assertIsNotNone(__lowerCAmelCase    )
        self.assertIsInstance(__lowerCAmelCase      , __lowerCAmelCase    )

        _UpperCAmelCase         =					TFAutoModelForSequenceClassification.from_pretrained(__lowerCAmelCase    )
        self.assertIsNotNone(__lowerCAmelCase    )
        self.assertIsInstance(__lowerCAmelCase      , __lowerCAmelCase    )


    @slow
    def    lowerCAmelCase_     (    self		:     Tuple    ):
      # for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
      for model_name in ["bert-base-uncased"]:
        _UpperCAmelCase         =					AutoConfig.from_pretrained(__lowerCAmelCase    )
        self.assertIsNotNone(__lowerCAmelCase    )
        self.assertIsInstance(__lowerCAmelCase      , __lowerCAmelCase    )

        _UpperCAmelCase         =					TFAutoModelForQuestionAnswering.from_pretrained(__lowerCAmelCase    )
        self.assertIsNotNone(__lowerCAmelCase    )
        self.assertIsInstance(__lowerCAmelCase      , __lowerCAmelCase    )


    @slow
    @require_tensorflow_probability
    def    lowerCAmelCase_     (    self		:     List[Any]    ):
      for model_name in TF_TAPAS_PRETRAINED_MODEL_ARCHIVE_LIST[5:6]:
        _UpperCAmelCase         =					AutoConfig.from_pretrained(__lowerCAmelCase    )
        self.assertIsNotNone(__lowerCAmelCase    )
        self.assertIsInstance(__lowerCAmelCase      , __lowerCAmelCase    )

        _UpperCAmelCase         =					TFAutoModelForTableQuestionAnswering.from_pretrained(__lowerCAmelCase    )
        _UpperCAmelCase       , _UpperCAmelCase         =					TFAutoModelForTableQuestionAnswering.from_pretrained(
            __lowerCAmelCase      , output_loading_info=__lowerCAmelCase    )
        self.assertIsNotNone(__lowerCAmelCase    )
        self.assertIsInstance(__lowerCAmelCase      , __lowerCAmelCase    )


    def    lowerCAmelCase_     (    self		:     Dict    ):
      _UpperCAmelCase         =					TFAutoModelWithLMHead.from_pretrained(__lowerCAmelCase    )
      self.assertIsInstance(__lowerCAmelCase      , __lowerCAmelCase    )
      self.assertEqual(model.num_parameters()      , 1_4410    )
      self.assertEqual(model.num_parameters(only_trainable=__lowerCAmelCase    )      , 1_4410    )


    def    lowerCAmelCase_     (    self		:     List[str]    ):
      _UpperCAmelCase         =					TFAutoModelWithLMHead.from_pretrained(__lowerCAmelCase    )
      self.assertIsInstance(__lowerCAmelCase      , __lowerCAmelCase    )
      self.assertEqual(model.num_parameters()      , 1_4410    )
      self.assertEqual(model.num_parameters(only_trainable=__lowerCAmelCase    )      , 1_4410    )


    def    lowerCAmelCase_     (    self		:     List[str]    ):
      # For the auto model mapping, FunnelConfig has two models: FunnelModel and FunnelBaseModel
      _UpperCAmelCase         =					TFAutoModel.from_pretrained("""sgugger/funnel-random-tiny"""    )
      self.assertIsInstance(__lowerCAmelCase      , __lowerCAmelCase    )

      _UpperCAmelCase         =					copy.deepcopy(model.config    )
      _UpperCAmelCase         =					["""FunnelBaseModel"""]
      _UpperCAmelCase         =					TFAutoModel.from_config(__lowerCAmelCase    )
      self.assertIsInstance(__lowerCAmelCase      , __lowerCAmelCase    )

      with tempfile.TemporaryDirectory() as tmp_dir:
        model.save_pretrained(__lowerCAmelCase    )
        _UpperCAmelCase         =					TFAutoModel.from_pretrained(__lowerCAmelCase    )
        self.assertIsInstance(__lowerCAmelCase      , __lowerCAmelCase    )


    def    lowerCAmelCase_     (    self		:     str    ):
      try:
        AutoConfig.register("""new-model"""      , __lowerCAmelCase    )

        _UpperCAmelCase         =					[
            TFAutoModel,
            TFAutoModelForCausalLM,
            TFAutoModelForMaskedLM,
            TFAutoModelForPreTraining,
            TFAutoModelForQuestionAnswering,
            TFAutoModelForSequenceClassification,
            TFAutoModelForTokenClassification,
        ]

        for auto_class in auto_classes:
          with self.subTest(auto_class.__name__    ):
            # Wrong config class will raise an error
            with self.assertRaises(__lowerCAmelCase    ):
              auto_class.register(__lowerCAmelCase      , __lowerCAmelCase    )
            auto_class.register(__lowerCAmelCase      , __lowerCAmelCase    )
            # Trying to register something existing in the Transformers library will raise an error
            with self.assertRaises(__lowerCAmelCase    ):
              auto_class.register(__lowerCAmelCase      , __lowerCAmelCase    )

            # Now that the config is registered, it can be used as any other config with the auto-API
            _UpperCAmelCase         =					BertModelTester(self    ).get_config()
            _UpperCAmelCase         =					NewModelConfig(**tiny_config.to_dict()    )
            _UpperCAmelCase         =					auto_class.from_config(__lowerCAmelCase    )
            self.assertIsInstance(__lowerCAmelCase      , __lowerCAmelCase    )

            with tempfile.TemporaryDirectory() as tmp_dir:
              model.save_pretrained(__lowerCAmelCase    )
              _UpperCAmelCase         =					auto_class.from_pretrained(__lowerCAmelCase    )
              self.assertIsInstance(__lowerCAmelCase      , __lowerCAmelCase    )

      finally:
        if "new-model" in CONFIG_MAPPING._extra_content:
          del CONFIG_MAPPING._extra_content["new-model"]
        for mapping in (
            TF_MODEL_MAPPING,
            TF_MODEL_FOR_PRETRAINING_MAPPING,
            TF_MODEL_FOR_QUESTION_ANSWERING_MAPPING,
            TF_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING,
            TF_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING,
            TF_MODEL_FOR_CAUSAL_LM_MAPPING,
            TF_MODEL_FOR_MASKED_LM_MAPPING,
        ):
          if NewModelConfig in mapping._extra_content:
            del mapping._extra_content[NewModelConfig]


    def    lowerCAmelCase_     (    self		:     List[Any]    ):
      with self.assertRaisesRegex(
          __lowerCAmelCase      , """bert-base is not a local folder and is not a valid model identifier"""    ):
        _UpperCAmelCase         =					TFAutoModel.from_pretrained("""bert-base"""    )


    def    lowerCAmelCase_     (    self		:     Tuple    ):
      with self.assertRaisesRegex(
          __lowerCAmelCase      , R"""aaaaaa is not a valid git identifier \(branch name, tag name or commit id\)"""    ):
        _UpperCAmelCase         =					TFAutoModel.from_pretrained(__lowerCAmelCase      , revision="""aaaaaa"""    )


    def    lowerCAmelCase_     (    self		:     Union[str, Any]    ):
      with self.assertRaisesRegex(
          __lowerCAmelCase      , """hf-internal-testing/config-no-model does not appear to have a file named pytorch_model.bin"""      , ):
        _UpperCAmelCase         =					TFAutoModel.from_pretrained("""hf-internal-testing/config-no-model"""    )


    def    lowerCAmelCase_     (    self		:     Any    ):
      with self.assertRaisesRegex(__lowerCAmelCase      , """Use `from_pt=True` to load this model"""    ):
        _UpperCAmelCase         =					TFAutoModel.from_pretrained("""hf-internal-testing/tiny-bert-pt-only"""    )


    def    lowerCAmelCase_     (    self		:     Any    ):
      # Make sure we have cached the model.
      _UpperCAmelCase         =					TFAutoModel.from_pretrained("""hf-internal-testing/tiny-random-bert"""    )
      with RequestCounter() as counter:
        _UpperCAmelCase         =					TFAutoModel.from_pretrained("""hf-internal-testing/tiny-random-bert"""    )
        self.assertEqual(counter.get_request_count      , 0    )
        self.assertEqual(counter.head_request_count      , 1    )
        self.assertEqual(counter.other_request_count      , 0    )

      # With a sharded checkpoint
      _UpperCAmelCase         =					TFAutoModel.from_pretrained("""ArthurZ/tiny-random-bert-sharded"""    )
      with RequestCounter() as counter:
        _UpperCAmelCase         =					TFAutoModel.from_pretrained("""ArthurZ/tiny-random-bert-sharded"""    )
        self.assertEqual(counter.get_request_count      , 0    )
        self.assertEqual(counter.head_request_count      , 1    )
        self.assertEqual(counter.other_request_count      , 0    )

"""simple docstring"""


from .constants import (
    MODEL_NAME,
    OPTIMIZER_NAME,
    RNG_STATE_NAME,
    SAFE_WEIGHTS_INDEX_NAME,
    SAFE_WEIGHTS_NAME,
    SCALER_NAME,
    SCHEDULER_NAME,
    TORCH_LAUNCH_PARAMS,
    WEIGHTS_INDEX_NAME,
    WEIGHTS_NAME,
)
from .dataclasses import (
    BnbQuantizationConfig,
    ComputeEnvironment,
    CustomDtype,
    DeepSpeedPlugin,
    DistributedDataParallelKwargs,
    DistributedType,
    DynamoBackend,
    FPaRecipeKwargs,
    FullyShardedDataParallelPlugin,
    GradientAccumulationPlugin,
    GradScalerKwargs,
    InitProcessGroupKwargs,
    KwargsHandler,
    LoggerType,
    MegatronLMPlugin,
    PrecisionType,
    ProjectConfiguration,
    RNGType,
    SageMakerDistributedType,
    TensorInformation,
    TorchDynamoPlugin,
)
from .environment import get_int_from_env, parse_choice_from_env, parse_flag_from_env
from .imports import (
    get_ccl_version,
    is_abit_bnb_available,
    is_abit_bnb_available,
    is_aim_available,
    is_bfaa_available,
    is_bnb_available,
    is_botoa_available,
    is_ccl_available,
    is_comet_ml_available,
    is_datasets_available,
    is_deepspeed_available,
    is_fpa_available,
    is_ipex_available,
    is_megatron_lm_available,
    is_mlflow_available,
    is_mps_available,
    is_npu_available,
    is_rich_available,
    is_safetensors_available,
    is_sagemaker_available,
    is_tensorboard_available,
    is_tpu_available,
    is_transformers_available,
    is_wandb_available,
    is_xpu_available,
)
from .modeling import (
    check_device_map,
    check_tied_parameters_in_config,
    check_tied_parameters_on_same_device,
    compute_module_sizes,
    convert_file_size_to_int,
    dtype_byte_size,
    find_tied_parameters,
    get_balanced_memory,
    get_max_layer_size,
    get_max_memory,
    get_mixed_precision_context_manager,
    id_tensor_storage,
    infer_auto_device_map,
    load_checkpoint_in_model,
    load_offloaded_weights,
    load_state_dict,
    named_module_tensors,
    retie_parameters,
    set_module_tensor_to_device,
    shard_checkpoint,
)
from .offload import (
    OffloadedWeightsLoader,
    PrefixedDataset,
    extract_submodules_state_dict,
    load_offloaded_weight,
    offload_state_dict,
    offload_weight,
    save_offload_index,
)
from .operations import (
    broadcast,
    broadcast_object_list,
    concatenate,
    convert_outputs_to_fpaa,
    convert_to_fpaa,
    find_batch_size,
    find_device,
    gather,
    gather_object,
    get_data_structure,
    honor_type,
    initialize_tensors,
    is_namedtuple,
    is_tensor_information,
    is_torch_tensor,
    listify,
    pad_across_processes,
    recursively_apply,
    reduce,
    send_to_device,
    slice_tensors,
)
from .versions import compare_versions, is_torch_version


if is_deepspeed_available():
    from .deepspeed import (
        DeepSpeedEngineWrapper,
        DeepSpeedOptimizerWrapper,
        DeepSpeedSchedulerWrapper,
        DummyOptim,
        DummyScheduler,
        HfDeepSpeedConfig,
    )

from .bnb import has_abit_bnb_layers, load_and_quantize_model
from .fsdp_utils import load_fsdp_model, load_fsdp_optimizer, save_fsdp_model, save_fsdp_optimizer
from .launch import (
    PrepareForLaunch,
    _filter_args,
    prepare_deepspeed_cmd_env,
    prepare_multi_gpu_env,
    prepare_sagemager_args_inputs,
    prepare_simple_launcher_cmd_env,
    prepare_tpu,
)
from .megatron_lm import (
    AbstractTrainStep,
    BertTrainStep,
    GPTTrainStep,
    MegatronEngine,
    MegatronLMDummyDataLoader,
    MegatronLMDummyScheduler,
    MegatronLMOptimizerWrapper,
    MegatronLMSchedulerWrapper,
    TaTrainStep,
    avg_losses_across_data_parallel_group,
    gather_across_data_parallel_groups,
)
from .megatron_lm import initialize as megatron_lm_initialize
from .megatron_lm import prepare_data_loader as megatron_lm_prepare_data_loader
from .megatron_lm import prepare_model as megatron_lm_prepare_model
from .megatron_lm import prepare_optimizer as megatron_lm_prepare_optimizer
from .megatron_lm import prepare_scheduler as megatron_lm_prepare_scheduler
from .memory import find_executable_batch_size, release_memory
from .other import (
    extract_model_from_parallel,
    get_pretty_name,
    is_port_in_use,
    merge_dicts,
    patch_environment,
    save,
    wait_for_everyone,
    write_basic_config,
)
from .random import set_seed, synchronize_rng_state, synchronize_rng_states
from .torch_xla import install_xla
from .tqdm import tqdm
from .transformer_engine import convert_model, has_transformer_engine_layers
