
import os
import tempfile
import unittest

from transformers import DistilBertConfig, is_torch_available
from transformers.testing_utils import require_torch, require_torch_gpu, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
		import torch

		from transformers import (
		    DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
		    DistilBertForMaskedLM,
		    DistilBertForMultipleChoice,
		    DistilBertForQuestionAnswering,
		    DistilBertForSequenceClassification,
		    DistilBertForTokenClassification,
		    DistilBertModel,
		)




class 				lowerCamelCase__							(       __lowercase):




	'''simple docstring'''




	def __init__(      self	:str  ,     a	:List[str]  ,     a	:Union[str, Any]=1_3  ,     a	:Optional[int]=7  ,     a	:List[str]=True  ,     a	:str=True  ,     a	:Optional[Any]=False  ,     a	:Tuple=True  ,     a	:Optional[Any]=9_9  ,     a	:str=3_2  ,     a	:Dict=5  ,     a	:int=4  ,     a	:int=3_7  ,     a	:Union[str, Any]="gelu"  ,     a	:Union[str, Any]=0.1  ,     a	:Union[str, Any]=0.1  ,     a	:Any=5_1_2  ,     a	:Dict=1_6  ,     a	:Tuple=2  ,     a	:int=0.02  ,     a	:Dict=3  ,     a	:Tuple=4  ,     a	:List[Any]=None  ,     )	->				Tuple:
								__UpperCamelCase			:     Tuple         =      parent
								__UpperCamelCase			:     List[Any]         =      batch_size
								__UpperCamelCase			:     Optional[int]         =      seq_length
								__UpperCamelCase			:     List[Any]         =      is_training
								__UpperCamelCase			:     List[str]         =      use_input_mask
								__UpperCamelCase			:     List[str]         =      use_token_type_ids
								__UpperCamelCase			:     Union[str, Any]         =      use_labels
								__UpperCamelCase			:     Union[str, Any]         =      vocab_size
								__UpperCamelCase			:     int         =      hidden_size
								__UpperCamelCase			:     Any         =      num_hidden_layers
								__UpperCamelCase			:     Optional[int]         =      num_attention_heads
								__UpperCamelCase			:     Tuple         =      intermediate_size
								__UpperCamelCase			:     Optional[int]         =      hidden_act
								__UpperCamelCase			:     List[str]         =      hidden_dropout_prob
								__UpperCamelCase			:     Union[str, Any]         =      attention_probs_dropout_prob
								__UpperCamelCase			:     Tuple         =      max_position_embeddings
								__UpperCamelCase			:     Tuple         =      type_vocab_size
								__UpperCamelCase			:     Any         =      type_sequence_label_size
								__UpperCamelCase			:     Union[str, Any]         =      initializer_range
								__UpperCamelCase			:     List[Any]         =      num_labels
								__UpperCamelCase			:     Optional[Any]         =      num_choices
								__UpperCamelCase			:     Optional[Any]         =      scope







	def      _lowerCamelCase       (      self	:Tuple		)	->				List[str]:
								__UpperCamelCase			:     Any         =      ids_tensor([self.batch_size, self.seq_length]  ,     self.vocab_size		)

								__UpperCamelCase			:     Optional[int]         =      None
								if self.use_input_mask:
															__UpperCamelCase			:     Dict         =      random_attention_mask([self.batch_size, self.seq_length]		)

								__UpperCamelCase			:     Any         =      None
								__UpperCamelCase			:     Dict         =      None
								__UpperCamelCase			:     Optional[int]         =      None
								if self.use_labels:
															__UpperCamelCase			:     int         =      ids_tensor([self.batch_size]  ,     self.type_sequence_label_size		)
															__UpperCamelCase			:     Optional[Any]         =      ids_tensor([self.batch_size, self.seq_length]  ,     self.num_labels		)
															__UpperCamelCase			:     Optional[int]         =      ids_tensor([self.batch_size]  ,     self.num_choices		)

								__UpperCamelCase			:     Optional[Any]         =      self.get_config()

								return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels







	def      _lowerCamelCase       (      self	:Optional[int]		)	->				List[str]:
								return DistilBertConfig(
								    vocab_size=self.vocab_size  ,     dim=self.hidden_size  ,     n_layers=self.num_hidden_layers  ,     n_heads=self.num_attention_heads  ,     hidden_dim=self.intermediate_size  ,     hidden_act=self.hidden_act  ,     dropout=self.hidden_dropout_prob  ,     attention_dropout=self.attention_probs_dropout_prob  ,     max_position_embeddings=self.max_position_embeddings  ,     initializer_range=self.initializer_range  ,     )







	def      _lowerCamelCase       (      self	:Optional[Any]  ,     a	:str  ,     a	:Optional[Any]  ,     a	:Any  ,     a	:List[Any]  ,     a	:List[Any]  ,     a	:List[Any]		)	->				Any:
								__UpperCamelCase			:     Optional[int]         =      DistilBertModel(config=a		)
								model.to(a		)
								model.eval()
								__UpperCamelCase			:     int         =      model(a  ,     a		)
								__UpperCamelCase			:     List[Any]         =      model(a		)
								self.parent.assertEqual(result.last_hidden_state.shape  ,     (self.batch_size, self.seq_length, self.hidden_size)		)







	def      _lowerCamelCase       (      self	:List[str]  ,     a	:int  ,     a	:str  ,     a	:Any  ,     a	:Optional[int]  ,     a	:List[Any]  ,     a	:List[str]		)	->				Optional[Any]:
								__UpperCamelCase			:     Optional[Any]         =      DistilBertForMaskedLM(config=a		)
								model.to(a		)
								model.eval()
								__UpperCamelCase			:     str         =      model(a  ,     attention_mask=a  ,     labels=a		)
								self.parent.assertEqual(result.logits.shape  ,     (self.batch_size, self.seq_length, self.vocab_size)		)







	def      _lowerCamelCase       (      self	:Tuple  ,     a	:List[Any]  ,     a	:str  ,     a	:int  ,     a	:List[Any]  ,     a	:List[str]  ,     a	:str		)	->				Optional[Any]:
								__UpperCamelCase			:     Tuple         =      DistilBertForQuestionAnswering(config=a		)
								model.to(a		)
								model.eval()
								__UpperCamelCase			:     List[Any]         =      model(
								    a  ,     attention_mask=a  ,     start_positions=a  ,     end_positions=a		)
								self.parent.assertEqual(result.start_logits.shape  ,     (self.batch_size, self.seq_length)		)
								self.parent.assertEqual(result.end_logits.shape  ,     (self.batch_size, self.seq_length)		)







	def      _lowerCamelCase       (      self	:Any  ,     a	:Union[str, Any]  ,     a	:Optional[int]  ,     a	:int  ,     a	:Dict  ,     a	:Optional[int]  ,     a	:str		)	->				Optional[Any]:
								__UpperCamelCase			:     int         =      self.num_labels
								__UpperCamelCase			:     Union[str, Any]         =      DistilBertForSequenceClassification(a		)
								model.to(a		)
								model.eval()
								__UpperCamelCase			:     Tuple         =      model(a  ,     attention_mask=a  ,     labels=a		)
								self.parent.assertEqual(result.logits.shape  ,     (self.batch_size, self.num_labels)		)







	def      _lowerCamelCase       (      self	:Tuple  ,     a	:Tuple  ,     a	:Dict  ,     a	:Any  ,     a	:Dict  ,     a	:Union[str, Any]  ,     a	:Union[str, Any]		)	->				str:
								__UpperCamelCase			:     List[str]         =      self.num_labels
								__UpperCamelCase			:     Any         =      DistilBertForTokenClassification(config=a		)
								model.to(a		)
								model.eval()

								__UpperCamelCase			:     Any         =      model(a  ,     attention_mask=a  ,     labels=a		)
								self.parent.assertEqual(result.logits.shape  ,     (self.batch_size, self.seq_length, self.num_labels)		)







	def      _lowerCamelCase       (      self	:Any  ,     a	:int  ,     a	:Optional[int]  ,     a	:str  ,     a	:Optional[int]  ,     a	:List[Any]  ,     a	:Any		)	->				int:
								__UpperCamelCase			:     Union[str, Any]         =      self.num_choices
								__UpperCamelCase			:     Optional[Any]         =      DistilBertForMultipleChoice(config=a		)
								model.to(a		)
								model.eval()
								__UpperCamelCase			:     Union[str, Any]         =      input_ids.unsqueeze(1		).expand(-1  ,     self.num_choices  ,     -1		).contiguous()
								__UpperCamelCase			:     Tuple         =      input_mask.unsqueeze(1		).expand(-1  ,     self.num_choices  ,     -1		).contiguous()
								__UpperCamelCase			:     Optional[Any]         =      model(
								    a  ,     attention_mask=a  ,     labels=a  ,     )
								self.parent.assertEqual(result.logits.shape  ,     (self.batch_size, self.num_choices)		)







	def      _lowerCamelCase       (      self	:Union[str, Any]		)	->				Union[str, Any]:
								__UpperCamelCase			:     List[Any]         =      self.prepare_config_and_inputs()
								((__UpperCamelCase)     ,		(__UpperCamelCase)     ,		(__UpperCamelCase)     ,		(__UpperCamelCase)     ,		(__UpperCamelCase)     ,		(__UpperCamelCase))			:     Union[str, Any]         =      config_and_inputs
								__UpperCamelCase			:     Dict         =      {"input_ids": input_ids, "attention_mask": input_mask}
								return config, inputs_dict




@require_torch
class 				lowerCamelCase__							(       __lowercase  ,		__lowercase  ,		unittest.TestCase):




	'''simple docstring'''




	_A			=  (
	    (
	        DistilBertModel,
	        DistilBertForMaskedLM,
	        DistilBertForMultipleChoice,
	        DistilBertForQuestionAnswering,
	        DistilBertForSequenceClassification,
	        DistilBertForTokenClassification,
	    )
	    if is_torch_available()
	    else None
	)
	_A			=  (
	    {
	        'feature-extraction': DistilBertModel,
	        'fill-mask': DistilBertForMaskedLM,
	        'question-answering': DistilBertForQuestionAnswering,
	        'text-classification': DistilBertForSequenceClassification,
	        'token-classification': DistilBertForTokenClassification,
	        'zero-shot': DistilBertForSequenceClassification,
	    }
	    if is_torch_available()
	    else {}
	)
	_A			=  True
	_A			=  True
	_A			=  True
	_A			=  True







	def      _lowerCamelCase       (      self	:Optional[int]		)	->				Tuple:
								__UpperCamelCase			:     str         =      DistilBertModelTester(self		)
								__UpperCamelCase			:     Union[str, Any]         =      ConfigTester(self  ,     config_class=a  ,     dim=3_7		)







	def      _lowerCamelCase       (      self	:List[str]		)	->				str:
								self.config_tester.run_common_tests()







	def      _lowerCamelCase       (      self	:Dict		)	->				int:
								__UpperCamelCase			:     int         =      self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_distilbert_model(*a		)







	def      _lowerCamelCase       (      self	:List[str]		)	->				int:
								__UpperCamelCase			:     str         =      self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_distilbert_for_masked_lm(*a		)







	def      _lowerCamelCase       (      self	:Union[str, Any]		)	->				str:
								__UpperCamelCase			:     Optional[Any]         =      self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_distilbert_for_question_answering(*a		)







	def      _lowerCamelCase       (      self	:int		)	->				Any:
								__UpperCamelCase			:     str         =      self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_distilbert_for_sequence_classification(*a		)







	def      _lowerCamelCase       (      self	:Dict		)	->				Optional[int]:
								__UpperCamelCase			:     str         =      self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_distilbert_for_token_classification(*a		)







	def      _lowerCamelCase       (      self	:Any		)	->				Any:
								__UpperCamelCase			:     Dict         =      self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_distilbert_for_multiple_choice(*a		)







	@slow
	def      _lowerCamelCase       (      self	:Any		)	->				List[str]:
								for model_name in DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
															__UpperCamelCase			:     Optional[Any]         =      DistilBertModel.from_pretrained(a		)
															self.assertIsNotNone(a		)







	@slow
	@require_torch_gpu
	def      _lowerCamelCase       (      self	:Union[str, Any]		)	->				Optional[int]:
								__UpperCamelCase     ,		__UpperCamelCase			:     List[str]         =      self.model_tester.prepare_config_and_inputs_for_common()
								for model_class in self.all_model_classes:
															# BertForMultipleChoice behaves incorrectly in JIT environments.
															if model_class == DistilBertForMultipleChoice:
																						return

															__UpperCamelCase			:     Tuple         =      True
															__UpperCamelCase			:     str         =      model_class(config=a		)

															__UpperCamelCase			:     Dict         =      self._prepare_for_class(a  ,     a		)
															__UpperCamelCase			:     Dict         =      torch.jit.trace(
															    a  ,     (inputs_dict["input_ids"].to("cpu"		), inputs_dict["attention_mask"].to("cpu"		))		)

															with tempfile.TemporaryDirectory() as tmp:
																						torch.jit.save(a  ,     os.path.join(a  ,     "traced_model.pt"		)		)
																						__UpperCamelCase			:     Optional[int]         =      torch.jit.load(os.path.join(a  ,     "traced_model.pt"		)  ,     map_location=a		)
																						loaded(inputs_dict["input_ids"].to(a		)  ,     inputs_dict["attention_mask"].to(a		)		)




@require_torch
class 				lowerCamelCase__							(       unittest.TestCase):




	'''simple docstring'''




	@slow
	def      _lowerCamelCase       (      self	:List[Any]		)	->				Optional[int]:
								__UpperCamelCase			:     Tuple         =      DistilBertModel.from_pretrained("distilbert-base-uncased"		)
								__UpperCamelCase			:     Union[str, Any]         =      torch.tensor([[0, 3_4_5, 2_3_2, 3_2_8, 7_4_0, 1_4_0, 1_6_9_5, 6_9, 6_0_7_8, 1_5_8_8, 2]]		)
								__UpperCamelCase			:     Tuple         =      torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]		)
								with torch.no_grad():
															__UpperCamelCase			:     Dict         =      model(a  ,     attention_mask=a		)[0]
								__UpperCamelCase			:     Dict         =      torch.Size((1, 1_1, 7_6_8)		)
								self.assertEqual(output.shape  ,     a		)
								__UpperCamelCase			:     Dict         =      torch.tensor(
								    [[[-0.1639, 0.3299, 0.1648], [-0.1746, 0.3289, 0.1710], [-0.1884, 0.3357, 0.1810]]]		)

								self.assertTrue(torch.allclose(output[:, 1:4, 1:4]  ,     a  ,     atol=1E-4		)		)

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