code: 
  
  
  
  
  
  '''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 							__UpperCAmelCase  :
  
  
  
  
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		def __init__(self						:      List[str]      ,      _lowerCAmelCase						:      Tuple      ,      _lowerCAmelCase						:      Union[str, Any]=13      ,      _lowerCAmelCase						:      List[Any]=7      ,      _lowerCAmelCase						:      Tuple=True      ,      _lowerCAmelCase						:      str=True      ,      _lowerCAmelCase						:      Optional[int]=True      ,      _lowerCAmelCase						:      Tuple=True      ,      _lowerCAmelCase						:      Any=99      ,      _lowerCAmelCase						:      Optional[Any]=32      ,      _lowerCAmelCase						:      str=18      ,      _lowerCAmelCase						:      int=43      ,      _lowerCAmelCase						:      Optional[int]=5      ,      _lowerCAmelCase						:      List[str]=4      ,      _lowerCAmelCase						:      str=37      ,      _lowerCAmelCase						:      Optional[int]="gelu"      ,      _lowerCAmelCase						:      str=0.1      ,      _lowerCAmelCase						:      List[Any]=0.1      ,      _lowerCAmelCase						:      Union[str, Any]=512      ,      _lowerCAmelCase						:      Dict=16      ,      _lowerCAmelCase						:      Optional[int]=2      ,      _lowerCAmelCase						:      Dict=0.02      ,      _lowerCAmelCase						:      Any=3      ,      _lowerCAmelCase						:      Tuple=4      ,      _lowerCAmelCase						:      Any=None      ,      ):
  			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        A     (self						:      int      ):
  			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        A     (self						:      List[Any]      ):
  			(
  			    (
  			    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        A     (self						:      int      ,      _lowerCAmelCase						:      str      ,      _lowerCAmelCase						:      Optional[int]      ,      _lowerCAmelCase						:      Optional[int]      ,      _lowerCAmelCase						:      Optional[int]      ,      _lowerCAmelCase						:      Optional[int]      ,      _lowerCAmelCase						:      List[str]      ,      _lowerCAmelCase						:      Any      ):
  			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        A     (self						:      List[str]      ,      _lowerCAmelCase						:      Optional[int]      ,      _lowerCAmelCase						:      Union[str, Any]      ,      _lowerCAmelCase						:      Any      ,      _lowerCAmelCase						:      Optional[Any]      ,      _lowerCAmelCase						:      Optional[Any]      ,      _lowerCAmelCase						:      int      ,      _lowerCAmelCase						:      List[str]      ,      _lowerCAmelCase						:      Dict      ,      _lowerCAmelCase						:      int      ,      ):
  			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        A     (self						:      List[Any]      ,      _lowerCAmelCase						:      List[str]      ,      _lowerCAmelCase						:      Union[str, Any]      ,      _lowerCAmelCase						:      Dict      ,      _lowerCAmelCase						:      Union[str, Any]      ,      _lowerCAmelCase						:      Tuple      ,      _lowerCAmelCase						:      Optional[int]      ,      _lowerCAmelCase						:      Optional[Any]      ,      _lowerCAmelCase						:      str      ,      _lowerCAmelCase						:      Optional[int]      ,      ):
  			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        A     (self						:      str      ,      _lowerCAmelCase						:      str      ,      _lowerCAmelCase						:      Any      ,      _lowerCAmelCase						:      Optional[Any]      ,      _lowerCAmelCase						:      Any      ,      _lowerCAmelCase						:      str      ,      _lowerCAmelCase						:      Optional[int]      ,      _lowerCAmelCase						:      int      ):
  			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        A     (self						:      List[Any]      ,      _lowerCAmelCase						:      List[Any]      ,      _lowerCAmelCase						:      str      ,      _lowerCAmelCase						:      Dict      ,      _lowerCAmelCase						:      Any      ,      _lowerCAmelCase						:      Tuple      ,      _lowerCAmelCase						:      int      ,      _lowerCAmelCase						:      Any      ,      _lowerCAmelCase						:      Optional[Any]      ,      _lowerCAmelCase						:      List[str]      ,      ):
  			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        A     (self						:      str      ,      _lowerCAmelCase						:      Tuple      ,      _lowerCAmelCase						:      int      ,      _lowerCAmelCase						:      Optional[int]      ,      _lowerCAmelCase						:      Any      ,      _lowerCAmelCase						:      List[str]      ,      _lowerCAmelCase						:      Tuple      ,      _lowerCAmelCase						:      Union[str, Any]      ):
  			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        A     (self						:      Optional[int]      ,      _lowerCAmelCase						:      Any      ,      _lowerCAmelCase						:      Any      ,      _lowerCAmelCase						:      List[str]      ,      _lowerCAmelCase						:      int      ,      _lowerCAmelCase						:      str      ,      _lowerCAmelCase						:      List[str]      ,      _lowerCAmelCase						:      Dict      ):
  			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        A     (self						:      int      ,      _lowerCAmelCase						:      Optional[int]      ,      _lowerCAmelCase						:      str      ,      _lowerCAmelCase						:      str      ,      _lowerCAmelCase						:      int      ,      _lowerCAmelCase						:      int      ,      _lowerCAmelCase						:      Dict      ,      _lowerCAmelCase						:      str      ):
  			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        A     (self						:      str      ,      _lowerCAmelCase						:      str      ,      _lowerCAmelCase						:      Optional[Any]      ,      _lowerCAmelCase						:      List[Any]      ,      _lowerCAmelCase						:      List[Any]      ,      _lowerCAmelCase						:      int      ,      _lowerCAmelCase						:      int      ,      _lowerCAmelCase						:      str      ):
  			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        A     (self						:      Optional[Any]      ):
  			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 							__UpperCAmelCase  (       A__		,					A__		,					unittest.TestCase						):
  
  
  
  
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		__lowerCAmelCase 				=    (
  		    (
  		        RemBertModel,
  		        RemBertForMaskedLM,
  		        RemBertForCausalLM,
  		        RemBertForMultipleChoice,
  		        RemBertForQuestionAnswering,
  		        RemBertForSequenceClassification,
  		        RemBertForTokenClassification,
  		    )
  		    if is_torch_available()
  		    else ()
  		)
  		__lowerCAmelCase 				=    (RemBertForCausalLM,) if is_torch_available() else ()
  		__lowerCAmelCase 				=    (
  		    {
  		        '''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        A     (self						:      Tuple      ):
  			A							  =      RemBertModelTester(self      )
  			A							  =      ConfigTester(self      ,      config_class=_lowerCAmelCase      ,      hidden_size=37      )
  
  
  
  
  
  
  
  		def        A     (self						:      Tuple      ):
  			self.config_tester.run_common_tests()
  
  
  
  
  
  
  
  		def        A     (self						:      Dict      ):
  			A							  =      self.model_tester.prepare_config_and_inputs()
  			self.model_tester.create_and_check_model(*_lowerCAmelCase      )
  
  
  
  
  
  
  
  		def        A     (self						:      List[Any]      ):
  			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        A     (self						:      int      ):
  			A							  =      self.model_tester.prepare_config_and_inputs()
  			self.model_tester.create_and_check_for_masked_lm(*_lowerCAmelCase      )
  
  
  
  
  
  
  
  		def        A     (self						:      List[str]      ):
  			A							  =      self.model_tester.prepare_config_and_inputs()
  			self.model_tester.create_and_check_for_multiple_choice(*_lowerCAmelCase      )
  
  
  
  
  
  
  
  		def        A     (self						:      int      ):
  			A							  =      self.model_tester.prepare_config_and_inputs_for_decoder()
  			self.model_tester.create_and_check_decoder_model_past_large_inputs(*_lowerCAmelCase      )
  
  
  
  
  
  
  
  		def        A     (self						:      Tuple      ):
  			A							  =      self.model_tester.prepare_config_and_inputs()
  			self.model_tester.create_and_check_for_question_answering(*_lowerCAmelCase      )
  
  
  
  
  
  
  
  		def        A     (self						:      Any      ):
  			A							  =      self.model_tester.prepare_config_and_inputs()
  			self.model_tester.create_and_check_for_sequence_classification(*_lowerCAmelCase      )
  
  
  
  
  
  
  
  		def        A     (self						:      Tuple      ):
  			A							  =      self.model_tester.prepare_config_and_inputs()
  			self.model_tester.create_and_check_for_token_classification(*_lowerCAmelCase      )
  
  
  
  
  
  
  
  		def        A     (self						:      int      ):
  			A							  =      self.model_tester.prepare_config_and_inputs_for_decoder()
  			self.model_tester.create_and_check_model_as_decoder(*_lowerCAmelCase      )
  
  
  
  
  
  
  
  		def        A     (self						:      str      ):
  			# This regression test was failing with PyTorch < 1.3
  			(
  			    (
  			    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        A     (self						:      Optional[int]      ):
  			for model_name in REMBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
  				A							  =      RemBertModel.from_pretrained(_lowerCAmelCase      )
  				self.assertIsNotNone(_lowerCAmelCase      )
  
  
  
  
  
  
  @require_torch
  class 							__UpperCAmelCase  (       unittest.TestCase						):
  
  
  
  
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		@slow
  		def        A     (self						:      List[str]      ):
  			# Test exact values at the last hidden layer
  			A							  =      RemBertModel.from_pretrained("""google/rembert"""      )
  			A							  =      torch.tensor([[312, 5_6498, 313, 2125, 313]]      )
  			A							  =      torch.tensor([[0, 0, 0, 1, 1]]      )
  			with torch.no_grad():
  				A							  =      model(_lowerCAmelCase      ,      token_type_ids=_lowerCAmelCase      ,      output_hidden_states=_lowerCAmelCase      )
  
  			A							  =      1152
  
  			A							  =      torch.Size((1, 5, hidden_size)      )
  			self.assertEqual(output["""last_hidden_state"""].shape      ,      _lowerCAmelCase      )
  
  			A							  =      torch.tensor(
  			    [
  			        [
  			            [0.0_754, -0.2_022, 0.1_904],
  			            [-0.3_354, -0.3_692, -0.4_791],
  			            [-0.2_314, -0.6_729, -0.0_749],
  			            [-0.0_396, -0.3_105, -0.4_234],
  			            [-0.1_571, -0.0_525, 0.5_353],
  			        ]
  			    ]      )
  
  			# 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      )      )
  
  
style_context: 
  
  
  
  
  
  '''simple docstring'''
  
  
  import logging
  import random
  
  import ray
  
  from transformers import RagConfig, RagRetriever, RagTokenizer
  from transformers.models.rag.retrieval_rag import CustomHFIndex
  
  
  _lowerCamelCase				:	Tuple				    =       logging.getLogger(__name__)
  
  
  
  class 							__UpperCAmelCase  :
  
  
  
  
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		def __init__(self						:      Union[str, Any]      ):
  			A							  =      False
  
  
  
  
  
  
  
  		def        A     (self						:      Any      ,      _lowerCAmelCase						:      str      ,      _lowerCAmelCase						:      List[str]      ,      _lowerCAmelCase						:      str      ,      _lowerCAmelCase						:      Any      ):
  			if not self.initialized:
  				A							  =      RagRetriever(
  				    _lowerCAmelCase      ,      question_encoder_tokenizer=_lowerCAmelCase      ,      generator_tokenizer=_lowerCAmelCase      ,      index=_lowerCAmelCase      ,      init_retrieval=_lowerCAmelCase      ,      )
  				A							  =      True
  
  
  
  
  
  
  
  		def        A     (self						:      Optional[int]      ):
  			self.retriever.index.init_index()
  
  
  
  
  
  
  
  		def        A     (self						:      Optional[int]      ):
  			# delete the old self.retriever object before assigning the new index
  			del self.retriever
  			A							  =      False
  
  
  
  
  
  
  
  		def        A     (self						:      List[str]      ,      _lowerCAmelCase						:      int      ,      _lowerCAmelCase						:      List[str]      ):
  			A  ,    A							  =      self.retriever._main_retrieve(_lowerCAmelCase      ,      _lowerCAmelCase      )
  			A							  =      self.retriever.index.get_doc_dicts(_lowerCAmelCase      )
  			return doc_ids, retrieved_doc_embeds, doc_dicts
  
  
  
  
  
  
  class 							__UpperCAmelCase  (       A__						):
  
  
  
  
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		def __init__(self						:      Tuple      ,      _lowerCAmelCase						:      int      ,      _lowerCAmelCase						:      int      ,      _lowerCAmelCase						:      List[str]      ,      _lowerCAmelCase						:      Tuple      ,      _lowerCAmelCase						:      str=None      ):
  			if index is not None and index.is_initialized() and len(_lowerCAmelCase      ) > 0:
  				raise ValueError(
  				    """When using Ray for distributed fine-tuning, """
  				    """you'll need to provide the paths instead, """
  				    """as the dataset and the index are loaded """
  				    """separately. More info in examples/rag/use_own_knowledge_dataset.py """      )
  
  			super().__init__(
  			    _lowerCAmelCase      ,      question_encoder_tokenizer=_lowerCAmelCase      ,      generator_tokenizer=_lowerCAmelCase      ,      index=_lowerCAmelCase      ,      init_retrieval=_lowerCAmelCase      ,      )
  
  			A							  =      retrieval_workers
  			A							  =      question_encoder_tokenizer
  			A							  =      generator_tokenizer
  			if len(self.retrieval_workers      ) > 0:
  				ray.get(
  				    [
  				        worker.create_rag_retriever.remote(_lowerCAmelCase      ,      _lowerCAmelCase      ,      _lowerCAmelCase      ,      _lowerCAmelCase      )
  				        for worker in self.retrieval_workers
  				    ]      )
  
  
  
  
  
  
  
  		def        A     (self						:      List[Any]      ):
  			logger.info("""initializing retrieval"""      )
  
  			if len(self.retrieval_workers      ) > 0:
  				ray.get([worker.init_retrieval.remote() for worker in self.retrieval_workers]      )
  			else:
  				# Non-distributed training. Load index into this same process.
  				self.index.init_index()
  
  
  
  
  
  
  
  		def        A     (self						:      str      ,      _lowerCAmelCase						:      Dict      ,      _lowerCAmelCase						:      Tuple      ):
  			if len(self.retrieval_workers      ) > 0:
  				# Select a random retrieval actor.
  				A							  =      self.retrieval_workers[random.randint(0      ,      len(self.retrieval_workers      ) - 1      )]
  				A  ,    A  ,    A							  =      ray.get(
  				    random_worker.retrieve.remote(_lowerCAmelCase      ,      _lowerCAmelCase      )      )
  			else:
  				A  ,    A							  =      self._main_retrieve(_lowerCAmelCase      ,      _lowerCAmelCase      )
  				A							  =      self.index.get_doc_dicts(_lowerCAmelCase      )
  			return retrieved_doc_embeds, doc_ids, doc_dicts
  
  
  
  
  
  
  
  		@classmethod
  		def        A     (cls						:      Union[str, Any]      ,      _lowerCAmelCase						:      int      ,      _lowerCAmelCase						:      Optional[Any]=None      ,      **_lowerCAmelCase						:      Union[str, Any]      ):
  			return super(_lowerCAmelCase      ,      cls      ).get_tokenizers(_lowerCAmelCase      ,      _lowerCAmelCase      ,      **_lowerCAmelCase      )
  
  
  
  
  
  
  
  		@classmethod
  		def        A     (cls						:      List[Any]      ,      _lowerCAmelCase						:      Optional[int]      ,      _lowerCAmelCase						:      List[str]      ,      _lowerCAmelCase						:      str=None      ,      **_lowerCAmelCase						:      List[str]      ):
  			A							  =      kwargs.pop("""config"""      ,      _lowerCAmelCase      ) or RagConfig.from_pretrained(_lowerCAmelCase      ,      **_lowerCAmelCase      )
  			A							  =      RagTokenizer.from_pretrained(_lowerCAmelCase      ,      config=_lowerCAmelCase      )
  			A							  =      rag_tokenizer.question_encoder
  			A							  =      rag_tokenizer.generator
  
  			if indexed_dataset is not None:
  				A							  =      """custom"""
  				A							  =      CustomHFIndex(config.retrieval_vector_size      ,      _lowerCAmelCase      )
  			else:
  				A							  =      cls._build_index(_lowerCAmelCase      )
  
  			return cls(
  			    _lowerCAmelCase      ,      question_encoder_tokenizer=_lowerCAmelCase      ,      generator_tokenizer=_lowerCAmelCase      ,      retrieval_workers=_lowerCAmelCase      ,      index=_lowerCAmelCase      ,      )
  
  
  
  
  
  
  
  		def        A     (self						:      List[Any]      ):
  			logger.info("""re-loading the new dataset with embeddings"""      )
  			# access from the training loop
  
  			ray.get([worker.clear_object.remote() for worker in self.retrieval_workers]      )
  
  			# build the index object again
  			A							  =      self._build_index(self.config      )
  
  			ray.get(
  			    [
  			        worker.create_rag_retriever.remote(
  			            self.config      ,      self.question_encoder_tokenizer      ,      self.generator_tokenizer      ,      _lowerCAmelCase      )
  			        for worker in self.retrieval_workers
  			    ]      )
  
  