






"""simple docstring"""


import unittest

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

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


if is_torch_available():
   import torch

   from transformers import (
       RemBertConfig,
       RemBertForCausalLM,
       RemBertForMaskedLM,
       RemBertForMultipleChoice,
       RemBertForQuestionAnswering,
       RemBertForSequenceClassification,
       RemBertForTokenClassification,
       RemBertModel,
   )
   from transformers.models.rembert.modeling_rembert import REMBERT_PRETRAINED_MODEL_ARCHIVE_LIST


class  snake_case		:


   '''simple docstring'''

   def __init__(     self			:					Optional[Any],	_lowerCamelCase			:					Optional[Any],	_lowerCamelCase			:					List[str]=13,	_lowerCamelCase			:					Any=7,	_lowerCamelCase			:					Union[str, Any]=True,	_lowerCamelCase			:					Union[str, Any]=True,	_lowerCamelCase			:					List[str]=True,	_lowerCamelCase			:					List[Any]=True,	_lowerCamelCase			:					Optional[Any]=99,	_lowerCamelCase			:					Dict=32,	_lowerCamelCase			:					Optional[Any]=18,	_lowerCamelCase			:					Tuple=43,	_lowerCamelCase			:					List[str]=5,	_lowerCamelCase			:					Any=4,	_lowerCamelCase			:					Tuple=37,	_lowerCamelCase			:					Optional[int]="gelu",	_lowerCamelCase			:					Union[str, Any]=0.1,	_lowerCamelCase			:					List[Any]=0.1,	_lowerCamelCase			:					Any=5_12,	_lowerCamelCase			:					List[Any]=16,	_lowerCamelCase			:					Tuple=2,	_lowerCamelCase			:					List[Any]=0.02,	_lowerCamelCase			:					Tuple=3,	_lowerCamelCase			:					int=4,	_lowerCamelCase			:					Tuple=None,	):

     '''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						      =			input_embedding_size
     __A						      =			output_embedding_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     _SCREAMING_SNAKE_CASE    (     self			:					str  ):

     '''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
     if self.use_token_type_ids:
       __A						      =			ids_tensor([self.batch_size, self.seq_length],	self.type_vocab_size  )

     __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						      =			RemBertConfig(
         vocab_size=self.vocab_size,	hidden_size=self.hidden_size,	input_embedding_size=self.input_embedding_size,	output_embedding_size=self.output_embedding_size,	num_hidden_layers=self.num_hidden_layers,	num_attention_heads=self.num_attention_heads,	intermediate_size=self.intermediate_size,	hidden_act=self.hidden_act,	hidden_dropout_prob=self.hidden_dropout_prob,	attention_probs_dropout_prob=self.attention_probs_dropout_prob,	max_position_embeddings=self.max_position_embeddings,	type_vocab_size=self.type_vocab_size,	is_decoder=_lowerCamelCase,	initializer_range=self.initializer_range,	)

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


   def     _SCREAMING_SNAKE_CASE    (     self			:					List[str]  ):

     '''simple docstring'''
     (
         (
         __A
     )      ,   (
         __A
     )      ,   (
         __A
     )      ,   (
         __A
     )      ,   (
         __A
     )      ,   (
         __A
     )      ,   (
         __A
     )      ,   
     )						      =			self.prepare_config_and_inputs()

     __A						      =			True
     __A						      =			floats_tensor([self.batch_size, self.seq_length, self.hidden_size]  )
     __A						      =			ids_tensor([self.batch_size, self.seq_length],	vocab_size=2  )

     return (
         config,
         input_ids,
         token_type_ids,
         input_mask,
         sequence_labels,
         token_labels,
         choice_labels,
         encoder_hidden_states,
         encoder_attention_mask,
     )


   def     _SCREAMING_SNAKE_CASE    (     self			:					Union[str, Any],	_lowerCamelCase			:					Optional[Any],	_lowerCamelCase			:					Union[str, Any],	_lowerCamelCase			:					Dict,	_lowerCamelCase			:					Tuple,	_lowerCamelCase			:					Optional[int],	_lowerCamelCase			:					List[str],	_lowerCamelCase			:					str  ):

     '''simple docstring'''
     __A						      =			RemBertModel(config=_lowerCamelCase  )
     model.to(_lowerCamelCase  )
     model.eval()
     __A						      =			model(_lowerCamelCase,	attention_mask=_lowerCamelCase,	token_type_ids=_lowerCamelCase  )
     __A						      =			model(_lowerCamelCase,	token_type_ids=_lowerCamelCase  )
     __A						      =			model(_lowerCamelCase  )
     self.parent.assertEqual(result.last_hidden_state.shape,	(self.batch_size, self.seq_length, self.hidden_size)  )


   def     _SCREAMING_SNAKE_CASE    (     self			:					Tuple,	_lowerCamelCase			:					Optional[int],	_lowerCamelCase			:					Any,	_lowerCamelCase			:					Any,	_lowerCamelCase			:					Any,	_lowerCamelCase			:					Dict,	_lowerCamelCase			:					int,	_lowerCamelCase			:					List[str],	_lowerCamelCase			:					Optional[Any],	_lowerCamelCase			:					List[str],	):

     '''simple docstring'''
     __A						      =			True
     __A						      =			RemBertModel(_lowerCamelCase  )
     model.to(_lowerCamelCase  )
     model.eval()
     __A						      =			model(
         _lowerCamelCase,	attention_mask=_lowerCamelCase,	token_type_ids=_lowerCamelCase,	encoder_hidden_states=_lowerCamelCase,	encoder_attention_mask=_lowerCamelCase,	)
     __A						      =			model(
         _lowerCamelCase,	attention_mask=_lowerCamelCase,	token_type_ids=_lowerCamelCase,	encoder_hidden_states=_lowerCamelCase,	)
     __A						      =			model(_lowerCamelCase,	attention_mask=_lowerCamelCase,	token_type_ids=_lowerCamelCase  )
     self.parent.assertEqual(result.last_hidden_state.shape,	(self.batch_size, self.seq_length, self.hidden_size)  )


   def     _SCREAMING_SNAKE_CASE    (     self			:					Tuple,	_lowerCamelCase			:					Dict,	_lowerCamelCase			:					List[Any],	_lowerCamelCase			:					Union[str, Any],	_lowerCamelCase			:					Optional[int],	_lowerCamelCase			:					Dict,	_lowerCamelCase			:					str,	_lowerCamelCase			:					List[str],	_lowerCamelCase			:					Tuple,	_lowerCamelCase			:					Optional[Any],	):

     '''simple docstring'''
     __A						      =			RemBertForCausalLM(config=_lowerCamelCase  )
     model.to(_lowerCamelCase  )
     model.eval()
     __A						      =			model(_lowerCamelCase,	attention_mask=_lowerCamelCase,	token_type_ids=_lowerCamelCase,	labels=_lowerCamelCase  )
     self.parent.assertEqual(result.logits.shape,	(self.batch_size, self.seq_length, self.vocab_size)  )


   def     _SCREAMING_SNAKE_CASE    (     self			:					int,	_lowerCamelCase			:					List[Any],	_lowerCamelCase			:					Optional[Any],	_lowerCamelCase			:					List[str],	_lowerCamelCase			:					List[Any],	_lowerCamelCase			:					Union[str, Any],	_lowerCamelCase			:					Optional[int],	_lowerCamelCase			:					Union[str, Any]  ):

     '''simple docstring'''
     __A						      =			RemBertForMaskedLM(config=_lowerCamelCase  )
     model.to(_lowerCamelCase  )
     model.eval()
     __A						      =			model(_lowerCamelCase,	attention_mask=_lowerCamelCase,	token_type_ids=_lowerCamelCase,	labels=_lowerCamelCase  )
     self.parent.assertEqual(result.logits.shape,	(self.batch_size, self.seq_length, self.vocab_size)  )


   def     _SCREAMING_SNAKE_CASE    (     self			:					Dict,	_lowerCamelCase			:					str,	_lowerCamelCase			:					List[str],	_lowerCamelCase			:					Tuple,	_lowerCamelCase			:					Dict,	_lowerCamelCase			:					List[Any],	_lowerCamelCase			:					Optional[int],	_lowerCamelCase			:					Any,	_lowerCamelCase			:					List[str],	_lowerCamelCase			:					Union[str, Any],	):

     '''simple docstring'''
     __A						      =			True
     __A						      =			True
     __A						      =			RemBertForCausalLM(config=_lowerCamelCase  )
     model.to(_lowerCamelCase  )
     model.eval()

     # first forward pass
     __A						      =			model(
         _lowerCamelCase,	attention_mask=_lowerCamelCase,	encoder_hidden_states=_lowerCamelCase,	encoder_attention_mask=_lowerCamelCase,	use_cache=_lowerCamelCase,	)
     __A						      =			outputs.past_key_values

     # create hypothetical multiple next token and extent to next_input_ids
     __A						      =			ids_tensor((self.batch_size, 3),	config.vocab_size  )
     __A						      =			ids_tensor((self.batch_size, 3),	vocab_size=2  )

     # append to next input_ids and
     __A						      =			torch.cat([input_ids, next_tokens],	dim=-1  )
     __A						      =			torch.cat([input_mask, next_mask],	dim=-1  )

     __A						      =			model(
         _lowerCamelCase,	attention_mask=_lowerCamelCase,	encoder_hidden_states=_lowerCamelCase,	encoder_attention_mask=_lowerCamelCase,	output_hidden_states=_lowerCamelCase,	)['''hidden_states'''][0]
     __A						      =			model(
         _lowerCamelCase,	attention_mask=_lowerCamelCase,	encoder_hidden_states=_lowerCamelCase,	encoder_attention_mask=_lowerCamelCase,	past_key_values=_lowerCamelCase,	output_hidden_states=_lowerCamelCase,	)['''hidden_states'''][0]

     # select random slice
     __A						      =			ids_tensor((1,),	output_from_past.shape[-1]  ).item()
     __A						      =			output_from_no_past[:, -3:, random_slice_idx].detach()
     __A						      =			output_from_past[:, :, random_slice_idx].detach()

     self.parent.assertTrue(output_from_past_slice.shape[1] == next_tokens.shape[1]  )

     # test that outputs are equal for slice
     self.parent.assertTrue(torch.allclose(_lowerCamelCase,	_lowerCamelCase,	atol=1e-3  )  )


   def     _SCREAMING_SNAKE_CASE    (     self			:					List[Any],	_lowerCamelCase			:					Optional[Any],	_lowerCamelCase			:					List[str],	_lowerCamelCase			:					Optional[int],	_lowerCamelCase			:					Optional[Any],	_lowerCamelCase			:					Dict,	_lowerCamelCase			:					str,	_lowerCamelCase			:					Optional[Any]  ):

     '''simple docstring'''
     __A						      =			RemBertForQuestionAnswering(config=_lowerCamelCase  )
     model.to(_lowerCamelCase  )
     model.eval()
     __A						      =			model(
         _lowerCamelCase,	attention_mask=_lowerCamelCase,	token_type_ids=_lowerCamelCase,	start_positions=_lowerCamelCase,	end_positions=_lowerCamelCase,	)
     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     _SCREAMING_SNAKE_CASE    (     self			:					str,	_lowerCamelCase			:					str,	_lowerCamelCase			:					int,	_lowerCamelCase			:					List[str],	_lowerCamelCase			:					Optional[Any],	_lowerCamelCase			:					str,	_lowerCamelCase			:					Dict,	_lowerCamelCase			:					List[Any]  ):

     '''simple docstring'''
     __A						      =			self.num_labels
     __A						      =			RemBertForSequenceClassification(_lowerCamelCase  )
     model.to(_lowerCamelCase  )
     model.eval()
     __A						      =			model(_lowerCamelCase,	attention_mask=_lowerCamelCase,	token_type_ids=_lowerCamelCase,	labels=_lowerCamelCase  )
     self.parent.assertEqual(result.logits.shape,	(self.batch_size, self.num_labels)  )


   def     _SCREAMING_SNAKE_CASE    (     self			:					Dict,	_lowerCamelCase			:					Tuple,	_lowerCamelCase			:					Any,	_lowerCamelCase			:					str,	_lowerCamelCase			:					Dict,	_lowerCamelCase			:					Optional[Any],	_lowerCamelCase			:					Optional[int],	_lowerCamelCase			:					List[str]  ):

     '''simple docstring'''
     __A						      =			self.num_labels
     __A						      =			RemBertForTokenClassification(config=_lowerCamelCase  )
     model.to(_lowerCamelCase  )
     model.eval()
     __A						      =			model(_lowerCamelCase,	attention_mask=_lowerCamelCase,	token_type_ids=_lowerCamelCase,	labels=_lowerCamelCase  )
     self.parent.assertEqual(result.logits.shape,	(self.batch_size, self.seq_length, self.num_labels)  )


   def     _SCREAMING_SNAKE_CASE    (     self			:					Union[str, Any],	_lowerCamelCase			:					Optional[int],	_lowerCamelCase			:					Optional[int],	_lowerCamelCase			:					List[Any],	_lowerCamelCase			:					int,	_lowerCamelCase			:					Dict,	_lowerCamelCase			:					Union[str, Any],	_lowerCamelCase			:					Optional[Any]  ):

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






   def     _SCREAMING_SNAKE_CASE    (     self			:					int  ):

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





@require_torch
class  snake_case		(							lowercase_       ,     lowercase_       ,     unittest.TestCase						):


   '''simple docstring'''

   A_						:     Any									=			(
       (
           RemBertModel,
           RemBertForMaskedLM,
           RemBertForCausalLM,
           RemBertForMultipleChoice,
           RemBertForQuestionAnswering,
           RemBertForSequenceClassification,
           RemBertForTokenClassification,
       )
       if is_torch_available()
       else ()
   )
   A_						:     List[Any]									=			(RemBertForCausalLM,) if is_torch_available() else ()
   A_						:     Tuple									=			(
       {
           """feature-extraction""": RemBertModel,
           """fill-mask""": RemBertForMaskedLM,
           """question-answering""": RemBertForQuestionAnswering,
           """text-classification""": RemBertForSequenceClassification,
           """text-generation""": RemBertForCausalLM,
           """token-classification""": RemBertForTokenClassification,
           """zero-shot""": RemBertForSequenceClassification,
       }
       if is_torch_available()
       else {}
   )


   def     _SCREAMING_SNAKE_CASE    (     self			:					Optional[int]  ):

     '''simple docstring'''
     __A						      =			RemBertModelTester(self  )
     __A						      =			ConfigTester(self,	config_class=_lowerCamelCase,	hidden_size=37  )


   def     _SCREAMING_SNAKE_CASE    (     self			:					Union[str, Any]  ):

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


   def     _SCREAMING_SNAKE_CASE    (     self			:					int  ):

     '''simple docstring'''
     __A						      =			self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_model(*_lowerCamelCase  )


   def     _SCREAMING_SNAKE_CASE    (     self			:					List[str]  ):

     '''simple docstring'''
     __A						      =			self.model_tester.prepare_config_and_inputs()
     for type in ["absolute", "relative_key", "relative_key_query"]:
       __A						      =			type
       self.model_tester.create_and_check_model(*_lowerCamelCase  )


   def     _SCREAMING_SNAKE_CASE    (     self			:					int  ):

     '''simple docstring'''
     __A						      =			self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_masked_lm(*_lowerCamelCase  )


   def     _SCREAMING_SNAKE_CASE    (     self			:					Tuple  ):

     '''simple docstring'''
     __A						      =			self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_multiple_choice(*_lowerCamelCase  )


   def     _SCREAMING_SNAKE_CASE    (     self			:					int  ):

     '''simple docstring'''
     __A						      =			self.model_tester.prepare_config_and_inputs_for_decoder()
     self.model_tester.create_and_check_decoder_model_past_large_inputs(*_lowerCamelCase  )


   def     _SCREAMING_SNAKE_CASE    (     self			:					Union[str, Any]  ):

     '''simple docstring'''
     __A						      =			self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_question_answering(*_lowerCamelCase  )


   def     _SCREAMING_SNAKE_CASE    (     self			:					Optional[Any]  ):

     '''simple docstring'''
     __A						      =			self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_sequence_classification(*_lowerCamelCase  )


   def     _SCREAMING_SNAKE_CASE    (     self			:					Optional[Any]  ):

     '''simple docstring'''
     __A						      =			self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_token_classification(*_lowerCamelCase  )


   def     _SCREAMING_SNAKE_CASE    (     self			:					Optional[int]  ):

     '''simple docstring'''
     __A						      =			self.model_tester.prepare_config_and_inputs_for_decoder()
     self.model_tester.create_and_check_model_as_decoder(*_lowerCamelCase  )


   def     _SCREAMING_SNAKE_CASE    (     self			:					List[Any]  ):

     '''simple docstring'''
     (
         (
         __A
     )      ,   (
         __A
     )      ,   (
         __A
     )      ,   (
         __A
     )      ,   (
         __A
     )      ,   (
         __A
     )      ,   (
         __A
     )      ,   (
         __A
     )      ,   (
         __A
     )      ,   
     )						      =			self.model_tester.prepare_config_and_inputs_for_decoder()

     __A						      =			None

     self.model_tester.create_and_check_model_as_decoder(
         _lowerCamelCase,	_lowerCamelCase,	_lowerCamelCase,	_lowerCamelCase,	_lowerCamelCase,	_lowerCamelCase,	_lowerCamelCase,	_lowerCamelCase,	_lowerCamelCase,	)






   @slow
   def     _SCREAMING_SNAKE_CASE    (     self			:					int  ):

     '''simple docstring'''







     for model_name in REMBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
       __A						      =			RemBertModel.from_pretrained(_lowerCamelCase  )
       self.assertIsNotNone(_lowerCamelCase  )





@require_torch
class  snake_case		(							unittest.TestCase						):


   '''simple docstring'''

   @slow
   def     _SCREAMING_SNAKE_CASE    (     self			:					int  ):

     '''simple docstring'''
     __A						      =			RemBertModel.from_pretrained('''google/rembert'''  )
     __A						      =			torch.tensor([[3_12, 5_64_98, 3_13, 21_25, 3_13]]  )
     __A						      =			torch.tensor([[0, 0, 0, 1, 1]]  )
     with torch.no_grad():
       __A						      =			model(_lowerCamelCase,	token_type_ids=_lowerCamelCase,	output_hidden_states=_lowerCamelCase  )

     __A						      =			11_52

     __A						      =			torch.Size((1, 5, hidden_size)  )
     self.assertEqual(output['''last_hidden_state'''].shape,	_lowerCamelCase  )

     __A						      =			torch.tensor(
         [
             [
                 [0.07_54, -0.20_22, 0.19_04],
                 [-0.33_54, -0.36_92, -0.47_91],
                 [-0.23_14, -0.67_29, -0.07_49],
                 [-0.03_96, -0.31_05, -0.42_34],
                 [-0.15_71, -0.05_25, 0.53_53],
             ]
         ]  )

     # Running on the original tf implementation gives slightly different results here.
     # Not clear why this variations is present
     # TODO: Find reason for discrepancy
     # expected_original_implementation = [[
     #     [0.07630594074726105, -0.20146065950393677, 0.19107051193714142],
     #     [-0.3405614495277405, -0.36971670389175415, -0.4808273911476135],
     #     [-0.22587086260318756, -0.6656315922737122, -0.07844287157058716],
     #     [-0.04145475849509239, -0.3077218234539032, -0.42316967248916626],
     #     [-0.15887849032878876, -0.054529931396245956, 0.5356100797653198]
     # ]]

     self.assertTrue(torch.allclose(output['''last_hidden_state'''][:, :, :3],	_lowerCamelCase,	atol=1e-4  )  )












"""simple docstring"""


# This test is meant to be run in on an instance with TPUs like this:
#
#   python examples/pytorch/xla_spawn.py --num_cores=8 tests/test_trainer_tpu.py
#
# Replace 8 with the number of TPU cores you have.
#

import sys
from typing import Dict

from transformers import EvalPrediction, HfArgumentParser, TrainingArguments, is_torch_available
from transformers.utils import logging


lowercase_    					=					logging.get_logger(__name__)


if is_torch_available():
       import torch
       from torch import nn
       from torch.utils.data import Dataset

       from transformers import Trainer


       class  snake_case		(							_lowerCAmelCase						):


          '''simple docstring'''

          def __init__(     self			:					List[Any],	_lowerCamelCase			:					int = 1_01  ):

            '''simple docstring'''
            __A						      =			length


          def __len__(     self			:					Optional[int]  ):

            '''simple docstring'''
            return self.length






          def __getitem__(     self			:					List[str],	_lowerCamelCase			:					Union[str, Any]  ):

            '''simple docstring'''
            return i





       class  snake_case		:


          '''simple docstring'''

          def __call__(     self			:					Tuple,	_lowerCamelCase			:					Dict  ):

            '''simple docstring'''
            return {"input_ids": torch.tensor(_lowerCamelCase  ), "labels": torch.tensor(_lowerCamelCase  )}





       class  snake_case		(							nn.Module						):


          '''simple docstring'''

          def __init__(     self			:					List[str]  ):

            '''simple docstring'''
            super().__init__()
            # Add some (unused) params otherwise DDP will complain.
            __A						      =			nn.Linear(1_20,	80  )






          def     _SCREAMING_SNAKE_CASE    (     self			:					Optional[Any],	_lowerCamelCase			:					Tuple,	_lowerCamelCase			:					Optional[Any]=None  ):

            '''simple docstring'''
            if labels is not None:
              return torch.tensor(0.0,	device=input_ids.device  ), input_ids
            else:
              return input_ids



def  lowerCAmelCase    (   ):




  """simple docstring"""







  __A						      =			HfArgumentParser((TrainingArguments,)			)
  sys.argv += ["--output_dir", "./examples"]
  __A						      =			parser.parse_args_into_dataclasses()[0]

  logger.warning(
      f'Process rank: {training_args.local_rank}, device: {training_args.device}, '
      f'tpu_num_cores: {training_args.tpu_num_cores}'		,	)

  # Essentially, what we want to verify in the distributed case is
  # that we get all samples back, in the right order.
  # (this is crucial for prediction for instance)
  for dataset_length in [1_0_0_1, 2_5_6, 1_5]:
    __A						      =			DummyDataset(__UpperCamelCase			)

    def compute_metrics(__UpperCamelCase			) -> Dict:
      __A						      =			list(range(len(__UpperCamelCase			)			)			)
      __A						      =			p.predictions.tolist() == sequential and p.label_ids.tolist() == sequential
      return {"success": success}

    __A						      =			Trainer(
        model=DummyModel()		,	args=__UpperCamelCase		,	data_collator=DummyDataCollator()		,	eval_dataset=__UpperCamelCase		,	compute_metrics=__UpperCamelCase		,	)
    __A						      =			trainer.evaluate()
    logger.info(__UpperCamelCase			)
    if metrics["eval_success"] is not True:
      logger.error(__UpperCamelCase			)
      exit(1			)

    __A						      =			trainer.predict(__UpperCamelCase			)
    logger.info(p.metrics			)
    if p.metrics["test_success"] is not True:
      logger.error(p.metrics			)
      exit(1			)

    __A						      =			2

    __A						      =			trainer.evaluate()
    logger.info(__UpperCamelCase			)
    if metrics["eval_success"] is not True:
      logger.error(__UpperCamelCase			)
      exit(1			)

    __A						      =			trainer.predict(__UpperCamelCase			)
    logger.info(p.metrics			)
    if p.metrics["test_success"] is not True:
      logger.error(p.metrics			)
      exit(1			)

    __A						      =			None

  logger.info('''🔥 All distributed tests successful'''			)



def  lowerCAmelCase    (   __UpperCamelCase			):




  """simple docstring"""







  main()


if __name__ == "__main__":
       main()




