

'''simple docstring'''







from copy import deepcopy

import torch
import torch.nn.functional as F
from torch.optim import AdamW
from torch.optim.lr_scheduler import LambdaLR
from torch.utils.data import DataLoader

from accelerate.accelerator import Accelerator
from accelerate.state import GradientState
from accelerate.test_utils import RegressionDataset, RegressionModel
from accelerate.utils import DistributedType, is_torch_version, set_seed
def    UpperCAmelCase__     (   UpperCAmelCase__,							UpperCAmelCase__,							UpperCAmelCase__,							UpperCAmelCase__					)      ->      List[Any]:
       for param, grad_param in zip(model_a.parameters(),							model_b.parameters()					):
              if not param.requires_grad:
                     continue
              if not did_step:
                     # Grads should not be in sync
                     assert (
                         torch.allclose(param.grad,							grad_param.grad					) is False
                     ), F'''Gradients in sync when they should not be at iteration {iteration}:\nmodel_a grad ({param.grad}) == model_b grad ({grad_param.grad})'''
              else:
                     # Grads should be in sync
                     assert (
                         torch.allclose(param.grad,							grad_param.grad					) is True
                     ), F'''Gradients not in sync when they should be at iteration {iteration}:\nmodel_a grad ({param.grad}) != model_b grad ({grad_param.grad})'''
def    UpperCAmelCase__     (   UpperCAmelCase__,							UpperCAmelCase__,							UpperCAmelCase__,							UpperCAmelCase__,							UpperCAmelCase__=True					)      ->      Optional[int]:
       model.train()
       A_						       =   model(UpperCAmelCase__					)
       A_						       =   F.mse_loss(UpperCAmelCase__,							target.to(output.device					)					)
       if not do_backward:
              loss /= accelerator.gradient_accumulation_steps
              loss.backward()
       else:
              accelerator.backward(UpperCAmelCase__					)
def    UpperCAmelCase__     (   UpperCAmelCase__,							UpperCAmelCase__=False					)      ->      int:
       set_seed(42					)
       A_						       =   RegressionModel()
       A_						       =   deepcopy(UpperCAmelCase__					)
       A_						       =   RegressionDataset(length=80					)
       A_						       =   DataLoader(UpperCAmelCase__,							batch_size=16					)
       model.to(accelerator.device					)
       if sched:
              A_						       =   AdamW(params=model.parameters(),							lr=1e-3					)
              A_						       =   AdamW(params=ddp_model.parameters(),							lr=1e-3					)
              A_						       =   LambdaLR(UpperCAmelCase__,							lr_lambda=lambda UpperCAmelCase__					: epoch**0.65					)
              A_						       =   LambdaLR(UpperCAmelCase__,							lr_lambda=lambda UpperCAmelCase__					: epoch**0.65					)
       # Make a copy of `model`
       if sched:
              A_       ,      A_       ,      A_       ,      A_						       =   accelerator.prepare(UpperCAmelCase__,							UpperCAmelCase__,							UpperCAmelCase__,							UpperCAmelCase__					)
       else:
              A_       ,      A_						       =   accelerator.prepare(UpperCAmelCase__,							UpperCAmelCase__					)
       if sched:
              return (model, opt, sched, dataloader, ddp_model, ddp_opt, ddp_sched)
       return model, ddp_model, dataloader
def    UpperCAmelCase__     (   UpperCAmelCase__					)      ->      Tuple:
       # Test when on a single CPU or GPU that the context manager does nothing
       A_       ,      A_       ,      A_						       =   get_training_setup(UpperCAmelCase__					)
       # Use a single batch
       A_       ,      A_						       =   next(iter(UpperCAmelCase__					)					).values()
       for iteration in range(3					):
              # Gather the distributed inputs and targs for the base model
              A_       ,      A_						       =   accelerator.gather((ddp_input, ddp_target)					)
              A_       ,      A_						       =   input.to(accelerator.device					), target.to(accelerator.device					)
              # Perform our initial ground truth step in non "DDP"
              step_model(UpperCAmelCase__,							UpperCAmelCase__,							UpperCAmelCase__,							UpperCAmelCase__					)
              # Do "gradient accumulation" (noop)
              if iteration % 2 == 0:
                     # Accumulate grads locally
                     with accelerator.no_sync(UpperCAmelCase__					):
                            step_model(UpperCAmelCase__,							UpperCAmelCase__,							UpperCAmelCase__,							UpperCAmelCase__					)
              else:
                     # Sync grads
                     step_model(UpperCAmelCase__,							UpperCAmelCase__,							UpperCAmelCase__,							UpperCAmelCase__					)

              # Since `no_sync` is a noop, `ddp_model` and `model` grads should always be in sync
              check_model_parameters(UpperCAmelCase__,							UpperCAmelCase__,							UpperCAmelCase__,							UpperCAmelCase__					)
              for param, ddp_param in zip(model.parameters(),							ddp_model.parameters()					):
                     if not param.requires_grad:
                            continue
                     assert torch.allclose(
                         param.grad,							ddp_param.grad					), F'''Gradients not in sync when they should be:\nModel grad ({param.grad}) != DDP grad ({ddp_param.grad})'''

              # Shuffle ddp_input on each iteration
              torch.manual_seed(13_37 + iteration					)
              A_						       =   ddp_input[torch.randperm(len(UpperCAmelCase__					)					)]
def    UpperCAmelCase__     (   UpperCAmelCase__					)      ->      Union[str, Any]:
       # Test on distributed setup that context manager behaves properly
       A_       ,      A_       ,      A_						       =   get_training_setup(UpperCAmelCase__					)
       # Use a single batch
       A_       ,      A_						       =   next(iter(UpperCAmelCase__					)					).values()
       for iteration in range(3					):
              # Gather the distributed inputs and targs for the base model
              A_       ,      A_						       =   accelerator.gather((ddp_input, ddp_target)					)
              A_       ,      A_						       =   input.to(accelerator.device					), target.to(accelerator.device					)
              # Perform our initial ground truth step in non "DDP"
              step_model(UpperCAmelCase__,							UpperCAmelCase__,							UpperCAmelCase__,							UpperCAmelCase__					)
              # Do "gradient accumulation" (noop)
              if iteration % 2 == 0:
                     # Accumulate grads locally
                     with accelerator.no_sync(UpperCAmelCase__					):
                            step_model(UpperCAmelCase__,							UpperCAmelCase__,							UpperCAmelCase__,							UpperCAmelCase__					)
              else:
                     # Sync grads
                     step_model(UpperCAmelCase__,							UpperCAmelCase__,							UpperCAmelCase__,							UpperCAmelCase__					)

              # DDP model and model should only be in sync when not (iteration % 2 == 0)
              for param, ddp_param in zip(model.parameters(),							ddp_model.parameters()					):
                     if not param.requires_grad:
                            continue
                     if iteration % 2 == 0:
                            # Grads should not be in sync
                            assert (
                                torch.allclose(param.grad,							ddp_param.grad					) is False
                            ), F'''Gradients in sync when they should not be:\nModel grad ({param.grad}) == DDP grad ({ddp_param.grad})'''
                     else:
                            # Grads should be in sync
                            assert (
                                torch.allclose(param.grad,							ddp_param.grad					) is True
                            ), F'''Gradients not in sync when they should be:\nModel grad ({param.grad}) != DDP grad ({ddp_param.grad})'''

        # Shuffle ddp_input on each iteration
              torch.manual_seed(13_37 + iteration					)
              A_						       =   ddp_input[torch.randperm(len(UpperCAmelCase__					)					)]
def    UpperCAmelCase__     (   UpperCAmelCase__=False,							UpperCAmelCase__=False					)      ->      Tuple:
       A_						       =   Accelerator(
           split_batches=UpperCAmelCase__,							dispatch_batches=UpperCAmelCase__,							gradient_accumulation_steps=2					)
       # Test that context manager behaves properly
       A_       ,      A_       ,      A_						       =   get_training_setup(UpperCAmelCase__					)
       for iteration, batch in enumerate(UpperCAmelCase__					):
              A_       ,      A_						       =   batch.values()
              # Gather the distributed inputs and targs for the base model
              A_       ,      A_						       =   accelerator.gather((ddp_input, ddp_target)					)
              A_       ,      A_						       =   input.to(accelerator.device					), target.to(accelerator.device					)
              # Perform our initial ground truth step in non "DDP"
              step_model(UpperCAmelCase__,							UpperCAmelCase__,							UpperCAmelCase__,							UpperCAmelCase__,							UpperCAmelCase__					)
              # Do "gradient accumulation" (noop)
              with accelerator.accumulate(UpperCAmelCase__					):
                     step_model(UpperCAmelCase__,							UpperCAmelCase__,							UpperCAmelCase__,							UpperCAmelCase__					)

              # DDP model and model should only be in sync when not (iteration % 2 == 0)
              for param, ddp_param in zip(model.parameters(),							ddp_model.parameters()					):
                     if not param.requires_grad:
                            continue
                     if ((iteration + 1) % 2 == 0) or (iteration == len(UpperCAmelCase__					) - 1):
                            # Grads should be in sync
                            assert (
                                torch.allclose(param.grad,							ddp_param.grad					) is True
                            ), F'''Gradients not in sync when they should be at iteration {iteration}:\nModel grad ({param.grad}) != DDP grad ({ddp_param.grad})'''
                     else:
                            # Grads should not be in sync
                            assert (
                                torch.allclose(param.grad,							ddp_param.grad					) is False
                            ), F'''Gradients in sync when they should not be at iteration {iteration}:\nModel grad ({param.grad}) == DDP grad ({ddp_param.grad})'''

        # Shuffle ddp_input on each iteration
              torch.manual_seed(13_37 + iteration					)
              A_						       =   ddp_input[torch.randperm(len(UpperCAmelCase__					)					)]
       GradientState._reset_state()
def    UpperCAmelCase__     (   UpperCAmelCase__=False,							UpperCAmelCase__=False					)      ->      str:
       A_						       =   Accelerator(
           split_batches=UpperCAmelCase__,							dispatch_batches=UpperCAmelCase__,							gradient_accumulation_steps=2					)
       # Test that context manager behaves properly
       A_       ,      A_       ,      A_       ,      A_       ,      A_       ,      A_       ,      A_						       =   get_training_setup(UpperCAmelCase__,							UpperCAmelCase__					)
       for iteration, batch in enumerate(UpperCAmelCase__					):
              A_       ,      A_						       =   batch.values()
              # Gather the distributed inputs and targs for the base model
              A_       ,      A_						       =   accelerator.gather((ddp_input, ddp_target)					)
              A_       ,      A_						       =   input.to(accelerator.device					), target.to(accelerator.device					)
              # Perform our initial ground truth step in non "DDP"
              model.train()
              ddp_model.train()
              step_model(UpperCAmelCase__,							UpperCAmelCase__,							UpperCAmelCase__,							UpperCAmelCase__,							UpperCAmelCase__					)
              opt.step()

              if ((iteration + 1) % 2 == 0) or ((iteration + 1) == len(UpperCAmelCase__					)):
                     if split_batches:
                            sched.step()
                     else:
                            for _ in range(accelerator.num_processes					):
                                   sched.step()
              opt.zero_grad()
              # Perform gradient accumulation under wrapper
              with accelerator.accumulate(UpperCAmelCase__					):
                     step_model(UpperCAmelCase__,							UpperCAmelCase__,							UpperCAmelCase__,							UpperCAmelCase__					)
                     ddp_opt.step()
                     ddp_sched.step()
                     ddp_opt.zero_grad()

              # Learning rates should be the same
              assert (
                  opt.param_groups[0]["lr"] == ddp_opt.param_groups[0]["lr"]
              ), F'''Learning rates found in each optimizer did not align\nopt: {opt.param_groups[0]["lr"]}\nDDP opt: {ddp_opt.param_groups[0]["lr"]}\n'''
              A_						       =   (((iteration + 1) % 2) == 0) or ((iteration + 1) == len(UpperCAmelCase__					))
              if accelerator.num_processes > 1:
                     check_model_parameters(UpperCAmelCase__,							UpperCAmelCase__,							UpperCAmelCase__,							UpperCAmelCase__					)
              # Shuffle ddp_input on each iteration
              torch.manual_seed(13_37 + iteration					)
       GradientState._reset_state()
def    UpperCAmelCase__     (   )      ->      Dict:
       A_						       =   Accelerator()

       A_						       =   RegressionDataset(length=80					)
       A_						       =   DataLoader(UpperCAmelCase__,							batch_size=16					)
       A_						       =   RegressionDataset(length=96					)
       A_						       =   DataLoader(UpperCAmelCase__,							batch_size=16					)
       A_       ,      A_						       =   accelerator.prepare(UpperCAmelCase__,							UpperCAmelCase__					)
       assert accelerator.gradient_state.active_dataloader is None
       for iteration, _ in enumerate(UpperCAmelCase__					):
              assert id(accelerator.gradient_state.active_dataloader					) == id(UpperCAmelCase__					)
              if iteration < len(UpperCAmelCase__					) - 1:
                     assert not accelerator.gradient_state.end_of_dataloader
                     if iteration == 1:
                            for batch_num, _ in enumerate(UpperCAmelCase__					):
                                   assert id(accelerator.gradient_state.active_dataloader					) == id(UpperCAmelCase__					)
                                   if batch_num < len(UpperCAmelCase__					) - 1:
                                          assert not accelerator.gradient_state.end_of_dataloader
                                   else:
                                          assert accelerator.gradient_state.end_of_dataloader
              else:
                     assert accelerator.gradient_state.end_of_dataloader
       assert accelerator.gradient_state.active_dataloader is None
def    UpperCAmelCase__     (   )      ->      Dict:
       A_						       =   Accelerator()
       A_						       =   accelerator.state
       if state.local_process_index == 0:
              print("""**Test `accumulate` gradient accumulation with dataloader break**"""					)
       test_dataloader_break()
       if state.distributed_type == DistributedType.NO:
              if state.local_process_index == 0:
                     print("""**Test NOOP `no_sync` context manager**"""					)
              test_noop_sync(UpperCAmelCase__					)
       if state.distributed_type in (DistributedType.MULTI_GPU, DistributedType.MULTI_CPU):
              if state.local_process_index == 0:
                     print("""**Test Distributed `no_sync` context manager**"""					)
              test_distributed_sync(UpperCAmelCase__					)
       if state.distributed_type == DistributedType.MULTI_GPU:
              for split_batch in [True, False]:
                     for dispatch_batches in [True, False]:
                            if state.local_process_index == 0:
                                   print(
                                       """**Test `accumulate` gradient accumulation, """,							F'''`split_batches={split_batch}` and `dispatch_batches={dispatch_batches}`**''',							)
                            test_gradient_accumulation(UpperCAmelCase__,							UpperCAmelCase__					)

    # Currently will break on torch 2.0 +, need to investigate why
       if is_torch_version("""<""",							"""2.0"""					) or state.distributed_type == DistributedType.NO:
              if state.local_process_index == 0:
                     print(
                         """**Test `accumulate` gradient accumulation with optimizer and scheduler, """,							"""`split_batches=False`, `dispatch_batches=False`**""",							)
              test_gradient_accumulation_with_opt_and_scheduler()
              if state.distributed_type == DistributedType.MULTI_GPU:
                     for split_batch in [True, False]:
                            for dispatch_batches in [True, False]:
                                   if not split_batch and not dispatch_batches:
                                          continue
                                   if state.local_process_index == 0:
                                          print(
                                              """**Test `accumulate` gradient accumulation with optimizer and scheduler, """,							F'''`split_batches={split_batch}` and `dispatch_batches={dispatch_batches}`**''',							)
                                   test_gradient_accumulation_with_opt_and_scheduler(UpperCAmelCase__,							UpperCAmelCase__					)
def    UpperCAmelCase__     (   UpperCAmelCase__					)      ->      Optional[int]:
       # For xla_spawn (TPUs)
       main()


if __name__ == "__main__":
      main()



'''simple docstring'''







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        A__   (   _snake_case      ):



    def __init__(       self		,					UpperCamelCase__		,					UpperCamelCase__=13		,					UpperCamelCase__=7		,					UpperCamelCase__=True		,					UpperCamelCase__=True		,					UpperCamelCase__=False		,					UpperCamelCase__=True		,					UpperCamelCase__=99		,					UpperCamelCase__=32		,					UpperCamelCase__=5		,					UpperCamelCase__=4		,					UpperCamelCase__=37		,					UpperCamelCase__="gelu"		,					UpperCamelCase__=0.1		,					UpperCamelCase__=0.1		,					UpperCamelCase__=512		,					UpperCamelCase__=16		,					UpperCamelCase__=2		,					UpperCamelCase__=0.02		,					UpperCamelCase__=3		,					UpperCamelCase__=4		,					UpperCamelCase__=None		,					)    ->   Dict:





           '''simple docstring'''
           A_						       =   parent
           A_						       =   batch_size
           A_						       =   seq_length
           A_						       =   is_training
           A_						       =   use_input_mask
           A_						       =   use_token_type_ids
           A_						       =   use_labels
           A_						       =   vocab_size
           A_						       =   hidden_size
           A_						       =   num_hidden_layers
           A_						       =   num_attention_heads
           A_						       =   intermediate_size
           A_						       =   hidden_act
           A_						       =   hidden_dropout_prob
           A_						       =   attention_probs_dropout_prob
           A_						       =   max_position_embeddings
           A_						       =   type_vocab_size
           A_						       =   type_sequence_label_size
           A_						       =   initializer_range
           A_						       =   num_labels
           A_						       =   num_choices
           A_						       =   scope



    def        snake_case_      (       self							)    ->   List[Any]:





           '''simple docstring'''
           A_						       =   ids_tensor([self.batch_size, self.seq_length]		,					self.vocab_size							)

           A_						       =   None
           if self.use_input_mask:
                  A_						       =   random_attention_mask([self.batch_size, self.seq_length]							)

           A_						       =   None
           A_						       =   None
           A_						       =   None
           if self.use_labels:
                  A_						       =   ids_tensor([self.batch_size]		,					self.type_sequence_label_size							)
                  A_						       =   ids_tensor([self.batch_size, self.seq_length]		,					self.num_labels							)
                  A_						       =   ids_tensor([self.batch_size]		,					self.num_choices							)

           A_						       =   self.get_config()

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



    def        snake_case_      (       self							)    ->   List[str]:





           '''simple docstring'''
           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        snake_case_      (       self		,					UpperCamelCase__		,					UpperCamelCase__		,					UpperCamelCase__		,					UpperCamelCase__		,					UpperCamelCase__		,					UpperCamelCase__							)    ->   Union[str, Any]:





           '''simple docstring'''
           A_						       =   DistilBertModel(config=UpperCamelCase__							)
           model.to(UpperCamelCase__							)
           model.eval()
           A_						       =   model(UpperCamelCase__		,					UpperCamelCase__							)
           A_						       =   model(UpperCamelCase__							)
           self.parent.assertEqual(result.last_hidden_state.shape		,					(self.batch_size, self.seq_length, self.hidden_size)							)



    def        snake_case_      (       self		,					UpperCamelCase__		,					UpperCamelCase__		,					UpperCamelCase__		,					UpperCamelCase__		,					UpperCamelCase__		,					UpperCamelCase__							)    ->   Union[str, Any]:





           '''simple docstring'''
           A_						       =   DistilBertForMaskedLM(config=UpperCamelCase__							)
           model.to(UpperCamelCase__							)
           model.eval()
           A_						       =   model(UpperCamelCase__		,					attention_mask=UpperCamelCase__		,					labels=UpperCamelCase__							)
           self.parent.assertEqual(result.logits.shape		,					(self.batch_size, self.seq_length, self.vocab_size)							)



    def        snake_case_      (       self		,					UpperCamelCase__		,					UpperCamelCase__		,					UpperCamelCase__		,					UpperCamelCase__		,					UpperCamelCase__		,					UpperCamelCase__							)    ->   Dict:





           '''simple docstring'''
           A_						       =   DistilBertForQuestionAnswering(config=UpperCamelCase__							)
           model.to(UpperCamelCase__							)
           model.eval()
           A_						       =   model(
               UpperCamelCase__		,					attention_mask=UpperCamelCase__		,					start_positions=UpperCamelCase__		,					end_positions=UpperCamelCase__							)
           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        snake_case_      (       self		,					UpperCamelCase__		,					UpperCamelCase__		,					UpperCamelCase__		,					UpperCamelCase__		,					UpperCamelCase__		,					UpperCamelCase__							)    ->   Union[str, Any]:





           '''simple docstring'''
           A_						       =   self.num_labels
           A_						       =   DistilBertForSequenceClassification(UpperCamelCase__							)
           model.to(UpperCamelCase__							)
           model.eval()
           A_						       =   model(UpperCamelCase__		,					attention_mask=UpperCamelCase__		,					labels=UpperCamelCase__							)
           self.parent.assertEqual(result.logits.shape		,					(self.batch_size, self.num_labels)							)



    def        snake_case_      (       self		,					UpperCamelCase__		,					UpperCamelCase__		,					UpperCamelCase__		,					UpperCamelCase__		,					UpperCamelCase__		,					UpperCamelCase__							)    ->   Any:





           '''simple docstring'''
           A_						       =   self.num_labels
           A_						       =   DistilBertForTokenClassification(config=UpperCamelCase__							)
           model.to(UpperCamelCase__							)
           model.eval()

           A_						       =   model(UpperCamelCase__		,					attention_mask=UpperCamelCase__		,					labels=UpperCamelCase__							)
           self.parent.assertEqual(result.logits.shape		,					(self.batch_size, self.seq_length, self.num_labels)							)



    def        snake_case_      (       self		,					UpperCamelCase__		,					UpperCamelCase__		,					UpperCamelCase__		,					UpperCamelCase__		,					UpperCamelCase__		,					UpperCamelCase__							)    ->   List[Any]:





           '''simple docstring'''
           A_						       =   self.num_choices
           A_						       =   DistilBertForMultipleChoice(config=UpperCamelCase__							)
           model.to(UpperCamelCase__							)
           model.eval()
           A_						       =   input_ids.unsqueeze(1							).expand(-1		,					self.num_choices		,					-1							).contiguous()
           A_						       =   input_mask.unsqueeze(1							).expand(-1		,					self.num_choices		,					-1							).contiguous()
           A_						       =   model(
               UpperCamelCase__		,					attention_mask=UpperCamelCase__		,					labels=UpperCamelCase__		,					)
           self.parent.assertEqual(result.logits.shape		,					(self.batch_size, self.num_choices)							)







    def        snake_case_      (       self							)    ->   Tuple:





           '''simple docstring'''
           A_						       =   self.prepare_config_and_inputs()
           ((A_)       ,      (A_)       ,      (A_)       ,      (A_)       ,      (A_)       ,      (A_))						       =   config_and_inputs
           A_						       =   {"""input_ids""": input_ids, """attention_mask""": input_mask}
           return config, inputs_dict


@require_torch
class        A__   (   _snake_case    ,    _snake_case    ,    unittest.TestCase      ):
    lowercase						       =		(
        (
            DistilBertModel,
            DistilBertForMaskedLM,
            DistilBertForMultipleChoice,
            DistilBertForQuestionAnswering,
            DistilBertForSequenceClassification,
            DistilBertForTokenClassification,
        )
        if is_torch_available()
        else None
    )
    lowercase						       =		(
        {
            "feature-extraction": DistilBertModel,
            "fill-mask": DistilBertForMaskedLM,
            "question-answering": DistilBertForQuestionAnswering,
            "text-classification": DistilBertForSequenceClassification,
            "token-classification": DistilBertForTokenClassification,
            "zero-shot": DistilBertForSequenceClassification,
        }
        if is_torch_available()
        else {}
    )
    lowercase						       =		True
    lowercase						       =		True
    lowercase						       =		True
    lowercase						       =		True



    def        snake_case_      (       self							)    ->   Tuple:





           '''simple docstring'''
           A_						       =   DistilBertModelTester(self							)
           A_						       =   ConfigTester(self		,					config_class=UpperCamelCase__		,					dim=37							)



    def        snake_case_      (       self							)    ->   List[Any]:





           '''simple docstring'''
           self.config_tester.run_common_tests()



    def        snake_case_      (       self							)    ->   Tuple:





           '''simple docstring'''
           A_						       =   self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_distilbert_model(*UpperCamelCase__							)



    def        snake_case_      (       self							)    ->   List[Any]:





           '''simple docstring'''
           A_						       =   self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_distilbert_for_masked_lm(*UpperCamelCase__							)



    def        snake_case_      (       self							)    ->   Union[str, Any]:





           '''simple docstring'''
           A_						       =   self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_distilbert_for_question_answering(*UpperCamelCase__							)



    def        snake_case_      (       self							)    ->   Any:





           '''simple docstring'''
           A_						       =   self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_distilbert_for_sequence_classification(*UpperCamelCase__							)



    def        snake_case_      (       self							)    ->   Dict:





           '''simple docstring'''
           A_						       =   self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_distilbert_for_token_classification(*UpperCamelCase__							)



    def        snake_case_      (       self							)    ->   Union[str, Any]:





           '''simple docstring'''
           A_						       =   self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_distilbert_for_multiple_choice(*UpperCamelCase__							)



    @slow
    def        snake_case_      (       self							)    ->   Tuple:





           '''simple docstring'''
           for model_name in DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                  A_						       =   DistilBertModel.from_pretrained(UpperCamelCase__							)
                  self.assertIsNotNone(UpperCamelCase__							)







    @slow
    @require_torch_gpu
    def        snake_case_      (       self							)    ->   str:





           '''simple docstring'''
           A_       ,      A_						       =   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

                  A_						       =   True
                  A_						       =   model_class(config=UpperCamelCase__							)

                  A_						       =   self._prepare_for_class(UpperCamelCase__		,					UpperCamelCase__							)
                  A_						       =   torch.jit.trace(
                      UpperCamelCase__		,					(inputs_dict["""input_ids"""].to("""cpu"""							), inputs_dict["""attention_mask"""].to("""cpu"""							))							)

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


@require_torch
class        A__   (   unittest.TestCase      ):



    @slow
    def        snake_case_      (       self							)    ->   int:





           '''simple docstring'''
           A_						       =   DistilBertModel.from_pretrained("""distilbert-base-uncased"""							)
           A_						       =   torch.tensor([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 2]]							)
           A_						       =   torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]							)
           with torch.no_grad():
                  A_						       =   model(UpperCamelCase__		,					attention_mask=UpperCamelCase__							)[0]
           A_						       =   torch.Size((1, 11, 768)							)
           self.assertEqual(output.shape		,					UpperCamelCase__							)
           A_						       =   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]		,					UpperCamelCase__		,					atol=1e-4							)							)
