






"""simple docstring"""



import itertools
import os
import subprocess
from os.path import dirname

from parameterized import parameterized

from tests.trainer.test_trainer import TrainerIntegrationCommon  # noqa
from transformers import is_torch_available
from transformers.testing_utils import (
    TestCasePlus,
    execute_subprocess_async,
    get_gpu_count,
    get_tests_dir,
    require_deepspeed,
    require_torch_gpu,
    slow,
)
from transformers.trainer_utils import set_seed


if is_torch_available():
    from tests.trainer.test_trainer import (  # noqa
        RegressionModelConfig,
        RegressionPreTrainedModel,
        get_regression_trainer,
    )


set_seed(4_2)

_UpperCamelCase		:  Dict						=		get_tests_dir('fixtures')
_UpperCamelCase		:  Dict						=		os.path.join(dirname(get_tests_dir()))
_UpperCamelCase		:  List[Any]						=		dirname(os.path.abspath(__file__))

# default torch.distributed port
_UpperCamelCase		:  List[str]						=		'10999'

_UpperCamelCase		:  Union[str, Any]						=		't5-small'

# *** Working Models ***
_UpperCamelCase		:  Optional[int]						=		'hf-internal-testing/tiny-albert'
_UpperCamelCase		:  Dict						=		'sshleifer/bart-tiny-random'
_UpperCamelCase		:  List[Any]						=		'hf-internal-testing/tiny-bert'
_UpperCamelCase		:  Union[str, Any]						=		'hf-internal-testing/tiny-random-bigbird_pegasus'
_UpperCamelCase		:  Optional[int]						=		'hf-internal-testing/tiny-random-big_bird'
_UpperCamelCase		:  Optional[int]						=		'hf-internal-testing/tiny-random-blenderbot'
_UpperCamelCase		:  Dict						=		'bigscience/bigscience-small-testing'
_UpperCamelCase		:  List[Any]						=		'hf-internal-testing/tiny-random-deberta'
_UpperCamelCase		:  Optional[Any]						=		'hf-internal-testing/tiny-random-deberta-v2'
_UpperCamelCase		:  int						=		'sshleifer/tiny-distilbert-base-cased'
_UpperCamelCase		:  Dict						=		'hf-internal-testing/tiny-electra'
_UpperCamelCase		:  List[str]						=		'hf-internal-testing/tiny-random-flaubert'
_UpperCamelCase		:  int						=		'stas/tiny-wmt19-en-de'
_UpperCamelCase		:  int						=		'hf-internal-testing/tiny-random-funnel'
_UpperCamelCase		:  Tuple						=		'sshleifer/tiny-gpt2'
_UpperCamelCase		:  List[Any]						=		'hf-internal-testing/tiny-random-gptj'
_UpperCamelCase		:  Optional[Any]						=		'hf-internal-testing/tiny-random-gpt_neo'
_UpperCamelCase		:  Tuple						=		'hf-internal-testing/tiny-layoutlm'
_UpperCamelCase		:  List[Any]						=		'hf-internal-testing/tiny-random-led'
_UpperCamelCase		:  Optional[int]						=		'hf-internal-testing/tiny-random-longformer'
_UpperCamelCase		:  Any						=		'stas/tiny-m2m_100'  # hf tiny model is unsuitable
_UpperCamelCase		:  Union[str, Any]						=		'sshleifer/tiny-marian-en-de'
_UpperCamelCase		:  Optional[Any]						=		'sshleifer/tiny-mbart'
_UpperCamelCase		:  Optional[Any]						=		'hf-internal-testing/tiny-random-mobilebert'
_UpperCamelCase		:  Union[str, Any]						=		'hf-internal-testing/tiny-random-mpnet'
_UpperCamelCase		:  Dict						=		'stas/pegasus-cnn_dailymail-tiny-random'
_UpperCamelCase		:  Optional[Any]						=		'hf-internal-testing/tiny-random-prophetnet'
_UpperCamelCase		:  Dict						=		'sshleifer/tiny-distilroberta-base'
_UpperCamelCase		:  Tuple						=		'hf-internal-testing/tiny-random-squeezebert'
_UpperCamelCase		:  Union[str, Any]						=		'patrickvonplaten/t5-tiny-random'
_UpperCamelCase		:  List[Any]						=		'hf-internal-testing/tiny-random-t5-v1.1'
_UpperCamelCase		:  List[str]						=		'hf-internal-testing/tiny-random-vit'
_UpperCamelCase		:  Any						=		'hf-internal-testing/tiny-xlm-roberta'
_UpperCamelCase		:  Optional[int]						=		'sshleifer/tiny-xlnet-base-cased'


# *** To Fix ***


# *** tiny model issues ***
# missing model files:
_UpperCamelCase		:  Dict						=		'hf-internal-testing/tiny-random-mt5'
_UpperCamelCase		:  str						=		'hf-internal-testing/tiny-random-camembert'
_UpperCamelCase		:  Union[str, Any]						=		'hf-internal-testing/tiny-random-openai-gpt'

# missing tokenizer files
_UpperCamelCase		:  Dict						=		'hf-internal-testing/tiny-random-convbert'
_UpperCamelCase		:  Dict						=		'hf-internal-testing/tiny-random-layoutlmv2'
_UpperCamelCase		:  Tuple						=		'hf-internal-testing/tiny-random-hubert'

# issues with tokenizer
_UpperCamelCase		:  Optional[Any]						=		'hf-internal-testing/tiny-random-ctrl'
_UpperCamelCase		:  Tuple						=		'hf-internal-testing/tiny-random-transfo-xl'  # same as ctrl

# other issues with tiny models
_UpperCamelCase		:  List[Any]						=		'hf-internal-testing/tiny-random-ibert'  # multiple issues with either mlm/qa/clas
_UpperCamelCase		:  int						=		'hf-internal-testing/tiny-random-reformer'  # multiple issues with either mlm/qa/clas

# *** Lacking official examples to test with ***
# or not working with examples
_UpperCamelCase		:  Optional[Any]						=		'hf-internal-testing/tiny-random-dpr'
# - "dpr"  examples/research_projects/rag-end2end-retriever/
_UpperCamelCase		:  int						=		'hf-internal-testing/tiny-random-rag'
# - "rag" research_projects
_UpperCamelCase		:  Optional[int]						=		''
# - "luke" Entities classes - no plan to make such example
_UpperCamelCase		:  List[str]						=		'hf-internal-testing/tiny-random-lxmert'
# - "lxmert" doesn't work with run_qa.py
_UpperCamelCase		:  List[Any]						=		'hf-internal-testing/tiny-random-clip'
# - "clip" nothing under pytorch examples - XXX: Suraj is working on adding some - check by end of Sep
_UpperCamelCase		:  Optional[Any]						=		'hf-internal-testing/tiny-random-speech_to_text'
# - "speech_to_text", nothing under pytorch examples


# *** Reactive mode ***
# models with low usage, unstable API, things about to change - do nothing about the following until someone runs into a problem
_UpperCamelCase		:  int						=		'hf-internal-testing/tiny-random-tapas'
def 						_SCREAMING_SNAKE_CASE			(   __snake_case      :  List[str]=False   ):





  '''simple docstring'''





  lowercase							   =	min(2		, get_gpu_count()   ) if distributed else 1
  lowercase							   =	os.environ.get('DS_TEST_PORT'		, __snake_case   )
  return f'deepspeed --num_nodes 1 --num_gpus {num_gpus} --master_port {master_port}'.split()
def 						_SCREAMING_SNAKE_CASE			(   ):





  '''simple docstring'''





  lowercase							   =	f'{FIXTURE_DIRECTORY}/tests_samples'
  lowercase							   =	f'{data_dir_samples}/wmt_en_ro'
  lowercase							   =	f'{data_dir_samples}/xsum'
  lowercase							   =	'\n        --do_train\n        --max_train_samples 4\n        --per_device_train_batch_size 2\n        --num_train_epochs 1\n        --fp16\n        --report_to none\n        --overwrite_output_dir\n        '.split()

  # try to cover as many models as possible once (it's enough to run on one task per model)
  # but need a tiny model for each
  #
  # should have "{model_type.upper()}_TINY" corresponding vars defined, e.g., T5_TINY, etc.
  lowercase							   =	{
      'trans': [
          'bart',
          'fsmt',
          'm2m_100',
          'marian',
          'mbart',
          't5',
          't5_v1',
          # "mt5", missing model files
      ],
      'sum': [
          'pegasus',
      ],
      'clm': [
          'big_bird',
          'bigbird_pegasus',
          'blenderbot',
          'bloom',
          'gpt2',
          'gpt_neo',
          'gptj',
          'xlm-roberta',
          'prophetnet',
          # "camembert", missing model files
      ],
      'mlm': [
          'albert',
          'deberta',
          'deberta-v2',
          'distilbert',
          'electra',
          'flaubert',
          'funnel',
          'layoutlm',
          # "reformer", # multiple issues with either mlm/qa/clas
      ],
      'qa': [
          'led',
          'longformer',
          'mobilebert',
          'mpnet',
          'roberta',
          'squeezebert',
          # "convbert", # missing tokenizer files
          # "layoutlmv2", missing model files
      ],
      'clas': [
          'bert',
          'xlnet',
          # "hubert", # missing tokenizer files
          # "ibert", # multiple issues with either mlm/qa/clas
          # "transfo-xl", # tokenizer issues as ctrl
          # "ctrl", # tokenizer issues
          # "openai-gpt", missing model files
          # "tapas", multiple issues
      ],
      'img_clas': [
          'vit',
      ],
  }

  lowercase							   =	f'{ROOT_DIRECTORY}/examples/pytorch'

  lowercase							   =	{
      'trans': f'\n        {scripts_dir}/translation/run_translation.py\n        --train_file {data_dir_wmt}/train.json\n        --source_lang en\n        --target_lang ro\n        ',
      'sum': f'\n        {scripts_dir}/summarization/run_summarization.py\n        --train_file {data_dir_xsum}/sample.json\n        --max_source_length 12\n        --max_target_length 12\n        --lang en\n        ',
      'clm': f'\n        {scripts_dir}/language-modeling/run_clm.py\n        --train_file {FIXTURE_DIRECTORY}/sample_text.txt\n        --block_size 8\n        ',
      'mlm': f'\n        {scripts_dir}/language-modeling/run_mlm.py\n        --train_file {FIXTURE_DIRECTORY}/sample_text.txt\n        ',
      'qa': f'\n        {scripts_dir}/question-answering/run_qa.py\n        --train_file {data_dir_samples}/SQUAD/sample.json\n        ',
      'clas': f'\n        {scripts_dir}/text-classification/run_glue.py\n        --train_file {data_dir_samples}/MRPC/train.csv\n        --max_seq_length 12\n        --task_name MRPC\n        ',
      'img_clas': f'\n        {scripts_dir}/image-classification/run_image_classification.py\n            --dataset_name hf-internal-testing/cats_vs_dogs_sample\n            --remove_unused_columns False\n            --max_steps 10\n            --image_processor_name {DS_TESTS_DIRECTORY}/vit_feature_extractor.json\n        ',
  }

  lowercase							   =	get_launcher(distributed=__snake_case   )

  lowercase							   =	{}
  for task, args in tasks.items():
    lowercase							   =	args.split()
    for model in tasksamodels[task]:
      lowercase							   =	globals()[f'{model.upper().replace("-"		, "_"   )}_TINY']
      lowercase							   =	f'--model_name_or_path {model_name}'.split()
      lowercase							   =	launcher + args + args_model + args_main

      # # generation special case
      # if task == "gen":
      #     launcher = f"deepspeed --num_nodes 1 --num_gpus 1".split()
      #     args_model += f"--model_type {model}".split()
      #     cmds[f"{task}_{model}"] = launcher + args + args_model
      # else:

  return cmds


_UpperCamelCase		:  str						=		make_task_cmds()

_UpperCamelCase		:  str						=		'zero2'
_UpperCamelCase		:  Any						=		'zero3'

_UpperCamelCase		:  Any						=		[ZEROa, ZEROa]
def 						_SCREAMING_SNAKE_CASE			(   __snake_case      :  Any		, __snake_case      :  int		, __snake_case      :  int   ):





  '''simple docstring'''





  lowercase							   =	parameterized.to_safe_name('_'.join(str(__snake_case   ) for x in param.args   )   )
  return f'{func.__name__}_{param_based_name}'


# Cartesian-product of zero stages with models to test
_UpperCamelCase		:  List[Any]						=		list(itertools.product(stages, task_cmds.keys()))



@slow
@require_deepspeed
@require_torch_gpu
class 	a					(       a_     ):







    def       UpperCamelCase_				(      self							,		_lowerCamelCase							,		_lowerCamelCase	):
      # return a ready to run train cmd
      if task not in task_cmds:
        raise ValueError(F'don\'t know of task {task}, have {task_cmds.keys()}'	)

      lowercase							   =	task_cmds[task]
      lowercase							   =	F'--deepspeed {self.test_file_dir_str}/ds_config_{stage}.json'.split()

      lowercase							   =	self.get_auto_remove_tmp_dir()
      lowercase							   =	F'--output_dir {output_dir}'.split()

      cmd += args_ds + args_out

      return cmd, output_dir







    @parameterized.expand(_lowerCamelCase							,		name_func=_lowerCamelCase	)
    def       UpperCamelCase_				(      self							,		_lowerCamelCase							,		_lowerCamelCase	):
      # testing the ability to do a run followed by recovery of full fp32 weights

      lowercase   ,					lowercase							   =	self.get_task_cmd(_lowerCamelCase							,		_lowerCamelCase	)

      # 1. generate the checkpoint
      cmd += "--save_steps 1".split()
      # keep for quick debug
      # print(" ".join([f"\nPYTHONPATH={self.src_dir_str}"] + cmd)); die
      execute_subprocess_async(_lowerCamelCase							,		env=self.get_env()	)

      # 2. test that the fp32 weights get reconsolidated
      lowercase							   =	F'{output_dir}/checkpoint-1'
      lowercase							   =	F'{chkpt_dir}/out.bin'
      lowercase							   =	F'{chkpt_dir}/zero_to_fp32.py {chkpt_dir} {recovered_model_path}'
      # keep for quick debug
      # print(" ".join([f"\nPYTHONPATH={self.src_dir_str}"] +cmd)); die
      subprocess.check_call(_lowerCamelCase							,		shell=_lowerCamelCase	)
      assert os.path.exists(_lowerCamelCase	), F'{recovered_model_path} was not found'

      # possibly could also test that the resulting saved model is usable but given that we use
      # random models we won't know if it's any good










"""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 	a					:







    def __init__(      self							,		_lowerCamelCase							,		_lowerCamelCase=1_3							,		_lowerCamelCase=7							,		_lowerCamelCase=True							,		_lowerCamelCase=True							,		_lowerCamelCase=True							,		_lowerCamelCase=True							,		_lowerCamelCase=9_9							,		_lowerCamelCase=3_2							,		_lowerCamelCase=1_8							,		_lowerCamelCase=4_3							,		_lowerCamelCase=5							,		_lowerCamelCase=4							,		_lowerCamelCase=3_7							,		_lowerCamelCase="gelu"							,		_lowerCamelCase=0.1							,		_lowerCamelCase=0.1							,		_lowerCamelCase=5_1_2							,		_lowerCamelCase=1_6							,		_lowerCamelCase=2							,		_lowerCamelCase=0.0_2							,		_lowerCamelCase=3							,		_lowerCamelCase=4							,		_lowerCamelCase=None							,		):
      lowercase							   =	parent
      lowercase							   =	batch_size
      lowercase							   =	seq_length
      lowercase							   =	is_training
      lowercase							   =	use_input_mask
      lowercase							   =	use_token_type_ids
      lowercase							   =	use_labels
      lowercase							   =	vocab_size
      lowercase							   =	hidden_size
      lowercase							   =	input_embedding_size
      lowercase							   =	output_embedding_size
      lowercase							   =	num_hidden_layers
      lowercase							   =	num_attention_heads
      lowercase							   =	intermediate_size
      lowercase							   =	hidden_act
      lowercase							   =	hidden_dropout_prob
      lowercase							   =	attention_probs_dropout_prob
      lowercase							   =	max_position_embeddings
      lowercase							   =	type_vocab_size
      lowercase							   =	type_sequence_label_size
      lowercase							   =	initializer_range
      lowercase							   =	num_labels
      lowercase							   =	num_choices
      lowercase							   =	scope







    def       UpperCamelCase_				(      self	):
      lowercase							   =	ids_tensor([self.batch_size, self.seq_length]							,		self.vocab_size	)

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

      lowercase							   =	None
      if self.use_token_type_ids:
        lowercase							   =	ids_tensor([self.batch_size, self.seq_length]							,		self.type_vocab_size	)

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

      lowercase							   =	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       UpperCamelCase_				(      self	):
      (
          (
          lowercase
      )   ,					(
          lowercase
      )   ,					(
          lowercase
      )   ,					(
          lowercase
      )   ,					(
          lowercase
      )   ,					(
          lowercase
      )   ,					(
          lowercase
      )   ,					
      )							   =	self.prepare_config_and_inputs()

      lowercase							   =	True
      lowercase							   =	floats_tensor([self.batch_size, self.seq_length, self.hidden_size]	)
      lowercase							   =	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       UpperCamelCase_				(      self							,		_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase	):
      lowercase							   =	RemBertModel(config=_lowerCamelCase	)
      model.to(_lowerCamelCase	)
      model.eval()
      lowercase							   =	model(_lowerCamelCase							,		attention_mask=_lowerCamelCase							,		token_type_ids=_lowerCamelCase	)
      lowercase							   =	model(_lowerCamelCase							,		token_type_ids=_lowerCamelCase	)
      lowercase							   =	model(_lowerCamelCase	)
      self.parent.assertEqual(result.last_hidden_state.shape							,		(self.batch_size, self.seq_length, self.hidden_size)	)







    def       UpperCamelCase_				(      self							,		_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase							,		):
      lowercase							   =	True
      lowercase							   =	RemBertModel(_lowerCamelCase	)
      model.to(_lowerCamelCase	)
      model.eval()
      lowercase							   =	model(
          _lowerCamelCase							,		attention_mask=_lowerCamelCase							,		token_type_ids=_lowerCamelCase							,		encoder_hidden_states=_lowerCamelCase							,		encoder_attention_mask=_lowerCamelCase							,		)
      lowercase							   =	model(
          _lowerCamelCase							,		attention_mask=_lowerCamelCase							,		token_type_ids=_lowerCamelCase							,		encoder_hidden_states=_lowerCamelCase							,		)
      lowercase							   =	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       UpperCamelCase_				(      self							,		_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase							,		):
      lowercase							   =	RemBertForCausalLM(config=_lowerCamelCase	)
      model.to(_lowerCamelCase	)
      model.eval()
      lowercase							   =	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       UpperCamelCase_				(      self							,		_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase	):
      lowercase							   =	RemBertForMaskedLM(config=_lowerCamelCase	)
      model.to(_lowerCamelCase	)
      model.eval()
      lowercase							   =	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       UpperCamelCase_				(      self							,		_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase							,		):
      lowercase							   =	True
      lowercase							   =	True
      lowercase							   =	RemBertForCausalLM(config=_lowerCamelCase	)
      model.to(_lowerCamelCase	)
      model.eval()

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

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

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

      lowercase							   =	model(
          _lowerCamelCase							,		attention_mask=_lowerCamelCase							,		encoder_hidden_states=_lowerCamelCase							,		encoder_attention_mask=_lowerCamelCase							,		output_hidden_states=_lowerCamelCase							,		)['hidden_states'][0]
      lowercase							   =	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
      lowercase							   =	ids_tensor((1,)							,		output_from_past.shape[-1]	).item()
      lowercase							   =	output_from_no_past[:, -3:, random_slice_idx].detach()
      lowercase							   =	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       UpperCamelCase_				(      self							,		_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase	):
      lowercase							   =	RemBertForQuestionAnswering(config=_lowerCamelCase	)
      model.to(_lowerCamelCase	)
      model.eval()
      lowercase							   =	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       UpperCamelCase_				(      self							,		_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase	):
      lowercase							   =	self.num_labels
      lowercase							   =	RemBertForSequenceClassification(_lowerCamelCase	)
      model.to(_lowerCamelCase	)
      model.eval()
      lowercase							   =	model(_lowerCamelCase							,		attention_mask=_lowerCamelCase							,		token_type_ids=_lowerCamelCase							,		labels=_lowerCamelCase	)
      self.parent.assertEqual(result.logits.shape							,		(self.batch_size, self.num_labels)	)







    def       UpperCamelCase_				(      self							,		_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase	):
      lowercase							   =	self.num_labels
      lowercase							   =	RemBertForTokenClassification(config=_lowerCamelCase	)
      model.to(_lowerCamelCase	)
      model.eval()
      lowercase							   =	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       UpperCamelCase_				(      self							,		_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase	):
      lowercase							   =	self.num_choices
      lowercase							   =	RemBertForMultipleChoice(config=_lowerCamelCase	)
      model.to(_lowerCamelCase	)
      model.eval()
      lowercase							   =	input_ids.unsqueeze(1	).expand(-1							,		self.num_choices							,		-1	).contiguous()
      lowercase							   =	token_type_ids.unsqueeze(1	).expand(-1							,		self.num_choices							,		-1	).contiguous()
      lowercase							   =	input_mask.unsqueeze(1	).expand(-1							,		self.num_choices							,		-1	).contiguous()
      lowercase							   =	model(
          _lowerCamelCase							,		attention_mask=_lowerCamelCase							,		token_type_ids=_lowerCamelCase							,		labels=_lowerCamelCase							,		)
      self.parent.assertEqual(result.logits.shape							,		(self.batch_size, self.num_choices)	)







    def       UpperCamelCase_				(      self	):
      lowercase							   =	self.prepare_config_and_inputs()
      (
          (
          lowercase
      )   ,					(
          lowercase
      )   ,					(
          lowercase
      )   ,					(
          lowercase
      )   ,					(
          lowercase
      )   ,					(
          lowercase
      )   ,					(
          lowercase
      )   ,					
      )							   =	config_and_inputs
      lowercase							   =	{'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': input_mask}
      return config, inputs_dict


@require_torch
class 	a					(       a_,    a_,    unittest.TestCase     ):
    UpperCAmelCase_			:					int     					=(
        (
            RemBertModel,
            RemBertForMaskedLM,
            RemBertForCausalLM,
            RemBertForMultipleChoice,
            RemBertForQuestionAnswering,
            RemBertForSequenceClassification,
            RemBertForTokenClassification,
        )
        if is_torch_available()
        else ()
    )
    UpperCAmelCase_			:					Optional[int]     					=(RemBertForCausalLM,) if is_torch_available() else ()
    UpperCAmelCase_			:					Dict     					=(
        {
            "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       UpperCamelCase_				(      self	):
      lowercase							   =	RemBertModelTester(self	)
      lowercase							   =	ConfigTester(self							,		config_class=_lowerCamelCase							,		hidden_size=3_7	)







    def       UpperCamelCase_				(      self	):
      self.config_tester.run_common_tests()







    def       UpperCamelCase_				(      self	):
      lowercase							   =	self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_model(*_lowerCamelCase	)







    def       UpperCamelCase_				(      self	):
      lowercase							   =	self.model_tester.prepare_config_and_inputs()
      for type in ["absolute", "relative_key", "relative_key_query"]:
        lowercase							   =	type
        self.model_tester.create_and_check_model(*_lowerCamelCase	)







    def       UpperCamelCase_				(      self	):
      lowercase							   =	self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_masked_lm(*_lowerCamelCase	)







    def       UpperCamelCase_				(      self	):
      lowercase							   =	self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_multiple_choice(*_lowerCamelCase	)







    def       UpperCamelCase_				(      self	):
      lowercase							   =	self.model_tester.prepare_config_and_inputs_for_decoder()
      self.model_tester.create_and_check_decoder_model_past_large_inputs(*_lowerCamelCase	)







    def       UpperCamelCase_				(      self	):
      lowercase							   =	self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_question_answering(*_lowerCamelCase	)







    def       UpperCamelCase_				(      self	):
      lowercase							   =	self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_sequence_classification(*_lowerCamelCase	)







    def       UpperCamelCase_				(      self	):
      lowercase							   =	self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_token_classification(*_lowerCamelCase	)







    def       UpperCamelCase_				(      self	):
      lowercase							   =	self.model_tester.prepare_config_and_inputs_for_decoder()
      self.model_tester.create_and_check_model_as_decoder(*_lowerCamelCase	)







    def       UpperCamelCase_				(      self	):
      # This regression test was failing with PyTorch < 1.3
      (
          (
          lowercase
      )   ,					(
          lowercase
      )   ,					(
          lowercase
      )   ,					(
          lowercase
      )   ,					(
          lowercase
      )   ,					(
          lowercase
      )   ,					(
          lowercase
      )   ,					(
          lowercase
      )   ,					(
          lowercase
      )   ,					
      )							   =	self.model_tester.prepare_config_and_inputs_for_decoder()

      lowercase							   =	None

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







    @slow
    def       UpperCamelCase_				(      self	):
      for model_name in REMBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
        lowercase							   =	RemBertModel.from_pretrained(_lowerCamelCase	)
        self.assertIsNotNone(_lowerCamelCase	)


@require_torch
class 	a					(       unittest.TestCase     ):







    @slow
    def       UpperCamelCase_				(      self	):
      # Test exact values at the last hidden layer
      lowercase							   =	RemBertModel.from_pretrained('google/rembert'	)
      lowercase							   =	torch.tensor([[3_1_2, 5_6_4_9_8, 3_1_3, 2_1_2_5, 3_1_3]]	)
      lowercase							   =	torch.tensor([[0, 0, 0, 1, 1]]	)
      with torch.no_grad():
        lowercase							   =	model(_lowerCamelCase							,		token_type_ids=_lowerCamelCase							,		output_hidden_states=_lowerCamelCase	)

      lowercase							   =	1_1_5_2

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

      lowercase							   =	torch.tensor(
          [
              [
                  [0.0_7_5_4, -0.2_0_2_2, 0.1_9_0_4],
                  [-0.3_3_5_4, -0.3_6_9_2, -0.4_7_9_1],
                  [-0.2_3_1_4, -0.6_7_2_9, -0.0_7_4_9],
                  [-0.0_3_9_6, -0.3_1_0_5, -0.4_2_3_4],
                  [-0.1_5_7_1, -0.0_5_2_5, 0.5_3_5_3],
              ]
          ]	)

      # 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	)	)


