

'''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	:
   def __init__(	self     : Dict       ,SCREAMING_SNAKE_CASE__     : Optional[int]       ,SCREAMING_SNAKE_CASE__     : Tuple=13       ,SCREAMING_SNAKE_CASE__     : Optional[int]=7       ,SCREAMING_SNAKE_CASE__     : Optional[int]=True       ,SCREAMING_SNAKE_CASE__     : Dict=True       ,SCREAMING_SNAKE_CASE__     : Tuple=True       ,SCREAMING_SNAKE_CASE__     : Dict=True       ,SCREAMING_SNAKE_CASE__     : List[str]=99       ,SCREAMING_SNAKE_CASE__     : Optional[int]=32       ,SCREAMING_SNAKE_CASE__     : Tuple=18       ,SCREAMING_SNAKE_CASE__     : List[str]=43       ,SCREAMING_SNAKE_CASE__     : Dict=5       ,SCREAMING_SNAKE_CASE__     : int=4       ,SCREAMING_SNAKE_CASE__     : List[str]=37       ,SCREAMING_SNAKE_CASE__     : Any="gelu"       ,SCREAMING_SNAKE_CASE__     : Any=0.1       ,SCREAMING_SNAKE_CASE__     : Any=0.1       ,SCREAMING_SNAKE_CASE__     : Optional[int]=512       ,SCREAMING_SNAKE_CASE__     : Any=16       ,SCREAMING_SNAKE_CASE__     : int=2       ,SCREAMING_SNAKE_CASE__     : List[str]=0.02       ,SCREAMING_SNAKE_CASE__     : Dict=3       ,SCREAMING_SNAKE_CASE__     : Any=4       ,SCREAMING_SNAKE_CASE__     : List[Any]=None       ,):
     SCREAMING_SNAKE_CASE:List[str]   =   parent
     SCREAMING_SNAKE_CASE:List[Any]   =   batch_size
     SCREAMING_SNAKE_CASE:Tuple   =   seq_length
     SCREAMING_SNAKE_CASE:Any   =   is_training
     SCREAMING_SNAKE_CASE:str   =   use_input_mask
     SCREAMING_SNAKE_CASE:Union[str, Any]   =   use_token_type_ids
     SCREAMING_SNAKE_CASE:Optional[int]   =   use_labels
     SCREAMING_SNAKE_CASE:List[Any]   =   vocab_size
     SCREAMING_SNAKE_CASE:Any   =   hidden_size
     SCREAMING_SNAKE_CASE:Dict   =   input_embedding_size
     SCREAMING_SNAKE_CASE:Tuple   =   output_embedding_size
     SCREAMING_SNAKE_CASE:Any   =   num_hidden_layers
     SCREAMING_SNAKE_CASE:str   =   num_attention_heads
     SCREAMING_SNAKE_CASE:Optional[Any]   =   intermediate_size
     SCREAMING_SNAKE_CASE:List[Any]   =   hidden_act
     SCREAMING_SNAKE_CASE:Optional[Any]   =   hidden_dropout_prob
     SCREAMING_SNAKE_CASE:List[str]   =   attention_probs_dropout_prob
     SCREAMING_SNAKE_CASE:Tuple   =   max_position_embeddings
     SCREAMING_SNAKE_CASE:Optional[Any]   =   type_vocab_size
     SCREAMING_SNAKE_CASE:List[Any]   =   type_sequence_label_size
     SCREAMING_SNAKE_CASE:Optional[Any]   =   initializer_range
     SCREAMING_SNAKE_CASE:List[str]   =   num_labels
     SCREAMING_SNAKE_CASE:List[Any]   =   num_choices
     SCREAMING_SNAKE_CASE:Dict   =   scope
   def       __UpperCamelCase  (	self     : List[Any]							):
     SCREAMING_SNAKE_CASE:int   =   ids_tensor([self.batch_size, self.seq_length]       ,self.vocab_size							)

     SCREAMING_SNAKE_CASE:Any   =   None
     if self.use_input_mask:
       SCREAMING_SNAKE_CASE:List[str]   =   random_attention_mask([self.batch_size, self.seq_length]							)

     SCREAMING_SNAKE_CASE:List[str]   =   None
     if self.use_token_type_ids:
       SCREAMING_SNAKE_CASE:Any   =   ids_tensor([self.batch_size, self.seq_length]       ,self.type_vocab_size							)

     SCREAMING_SNAKE_CASE:str   =   None
     SCREAMING_SNAKE_CASE:int   =   None
     SCREAMING_SNAKE_CASE:List[str]   =   None
     if self.use_labels:
       SCREAMING_SNAKE_CASE:Union[str, Any]   =   ids_tensor([self.batch_size]       ,self.type_sequence_label_size							)
       SCREAMING_SNAKE_CASE:str   =   ids_tensor([self.batch_size, self.seq_length]       ,self.num_labels							)
       SCREAMING_SNAKE_CASE:Union[str, Any]   =   ids_tensor([self.batch_size]       ,self.num_choices							)

     SCREAMING_SNAKE_CASE:Union[str, Any]   =   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=SCREAMING_SNAKE_CASE__       ,initializer_range=self.initializer_range       ,)

     return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels
   def       __UpperCamelCase  (	self     : List[Any]							):
     (
         (
         SCREAMING_SNAKE_CASE
     ) ,      (
         SCREAMING_SNAKE_CASE
     ) ,      (
         SCREAMING_SNAKE_CASE
     ) ,      (
         SCREAMING_SNAKE_CASE
     ) ,      (
         SCREAMING_SNAKE_CASE
     ) ,      (
         SCREAMING_SNAKE_CASE
     ) ,      (
         SCREAMING_SNAKE_CASE
     ) ,      
     ):List[str]   =   self.prepare_config_and_inputs()

     SCREAMING_SNAKE_CASE:Union[str, Any]   =   True
     SCREAMING_SNAKE_CASE:List[Any]   =   floats_tensor([self.batch_size, self.seq_length, self.hidden_size]							)
     SCREAMING_SNAKE_CASE:int   =   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     : Any       ,SCREAMING_SNAKE_CASE__     : Dict       ,SCREAMING_SNAKE_CASE__     : Optional[int]       ,SCREAMING_SNAKE_CASE__     : Any       ,SCREAMING_SNAKE_CASE__     : str       ,SCREAMING_SNAKE_CASE__     : Tuple       ,SCREAMING_SNAKE_CASE__     : Any       ,SCREAMING_SNAKE_CASE__     : Dict							):
     SCREAMING_SNAKE_CASE:List[Any]   =   RemBertModel(config=SCREAMING_SNAKE_CASE__							)
     model.to(SCREAMING_SNAKE_CASE__							)
     model.eval()
     SCREAMING_SNAKE_CASE:Optional[int]   =   model(SCREAMING_SNAKE_CASE__       ,attention_mask=SCREAMING_SNAKE_CASE__       ,token_type_ids=SCREAMING_SNAKE_CASE__							)
     SCREAMING_SNAKE_CASE:List[Any]   =   model(SCREAMING_SNAKE_CASE__       ,token_type_ids=SCREAMING_SNAKE_CASE__							)
     SCREAMING_SNAKE_CASE:Dict   =   model(SCREAMING_SNAKE_CASE__							)
     self.parent.assertEqual(result.last_hidden_state.shape       ,(self.batch_size, self.seq_length, self.hidden_size)							)
   def       __UpperCamelCase  (	self     : List[Any]       ,SCREAMING_SNAKE_CASE__     : Any       ,SCREAMING_SNAKE_CASE__     : Any       ,SCREAMING_SNAKE_CASE__     : Optional[Any]       ,SCREAMING_SNAKE_CASE__     : List[str]       ,SCREAMING_SNAKE_CASE__     : Optional[Any]       ,SCREAMING_SNAKE_CASE__     : Dict       ,SCREAMING_SNAKE_CASE__     : Union[str, Any]       ,SCREAMING_SNAKE_CASE__     : int       ,SCREAMING_SNAKE_CASE__     : int       ,):
     SCREAMING_SNAKE_CASE:Optional[int]   =   True
     SCREAMING_SNAKE_CASE:str   =   RemBertModel(SCREAMING_SNAKE_CASE__							)
     model.to(SCREAMING_SNAKE_CASE__							)
     model.eval()
     SCREAMING_SNAKE_CASE:Optional[int]   =   model(
         SCREAMING_SNAKE_CASE__       ,attention_mask=SCREAMING_SNAKE_CASE__       ,token_type_ids=SCREAMING_SNAKE_CASE__       ,encoder_hidden_states=SCREAMING_SNAKE_CASE__       ,encoder_attention_mask=SCREAMING_SNAKE_CASE__       ,)
     SCREAMING_SNAKE_CASE:Tuple   =   model(
         SCREAMING_SNAKE_CASE__       ,attention_mask=SCREAMING_SNAKE_CASE__       ,token_type_ids=SCREAMING_SNAKE_CASE__       ,encoder_hidden_states=SCREAMING_SNAKE_CASE__       ,)
     SCREAMING_SNAKE_CASE:Optional[Any]   =   model(SCREAMING_SNAKE_CASE__       ,attention_mask=SCREAMING_SNAKE_CASE__       ,token_type_ids=SCREAMING_SNAKE_CASE__							)
     self.parent.assertEqual(result.last_hidden_state.shape       ,(self.batch_size, self.seq_length, self.hidden_size)							)
   def       __UpperCamelCase  (	self     : Optional[Any]       ,SCREAMING_SNAKE_CASE__     : Optional[int]       ,SCREAMING_SNAKE_CASE__     : Any       ,SCREAMING_SNAKE_CASE__     : Optional[Any]       ,SCREAMING_SNAKE_CASE__     : Union[str, Any]       ,SCREAMING_SNAKE_CASE__     : Tuple       ,SCREAMING_SNAKE_CASE__     : Any       ,SCREAMING_SNAKE_CASE__     : Any       ,SCREAMING_SNAKE_CASE__     : List[str]       ,SCREAMING_SNAKE_CASE__     : Any       ,):
     SCREAMING_SNAKE_CASE:List[str]   =   RemBertForCausalLM(config=SCREAMING_SNAKE_CASE__							)
     model.to(SCREAMING_SNAKE_CASE__							)
     model.eval()
     SCREAMING_SNAKE_CASE:Tuple   =   model(SCREAMING_SNAKE_CASE__       ,attention_mask=SCREAMING_SNAKE_CASE__       ,token_type_ids=SCREAMING_SNAKE_CASE__       ,labels=SCREAMING_SNAKE_CASE__							)
     self.parent.assertEqual(result.logits.shape       ,(self.batch_size, self.seq_length, self.vocab_size)							)
   def       __UpperCamelCase  (	self     : Optional[int]       ,SCREAMING_SNAKE_CASE__     : Tuple       ,SCREAMING_SNAKE_CASE__     : Dict       ,SCREAMING_SNAKE_CASE__     : Dict       ,SCREAMING_SNAKE_CASE__     : Tuple       ,SCREAMING_SNAKE_CASE__     : Union[str, Any]       ,SCREAMING_SNAKE_CASE__     : List[str]       ,SCREAMING_SNAKE_CASE__     : Optional[int]							):
     SCREAMING_SNAKE_CASE:Dict   =   RemBertForMaskedLM(config=SCREAMING_SNAKE_CASE__							)
     model.to(SCREAMING_SNAKE_CASE__							)
     model.eval()
     SCREAMING_SNAKE_CASE:Optional[int]   =   model(SCREAMING_SNAKE_CASE__       ,attention_mask=SCREAMING_SNAKE_CASE__       ,token_type_ids=SCREAMING_SNAKE_CASE__       ,labels=SCREAMING_SNAKE_CASE__							)
     self.parent.assertEqual(result.logits.shape       ,(self.batch_size, self.seq_length, self.vocab_size)							)
   def       __UpperCamelCase  (	self     : Optional[int]       ,SCREAMING_SNAKE_CASE__     : Optional[Any]       ,SCREAMING_SNAKE_CASE__     : List[str]       ,SCREAMING_SNAKE_CASE__     : int       ,SCREAMING_SNAKE_CASE__     : int       ,SCREAMING_SNAKE_CASE__     : int       ,SCREAMING_SNAKE_CASE__     : Dict       ,SCREAMING_SNAKE_CASE__     : str       ,SCREAMING_SNAKE_CASE__     : List[Any]       ,SCREAMING_SNAKE_CASE__     : Union[str, Any]       ,):
     SCREAMING_SNAKE_CASE:Union[str, Any]   =   True
     SCREAMING_SNAKE_CASE:Dict   =   True
     SCREAMING_SNAKE_CASE:Optional[int]   =   RemBertForCausalLM(config=SCREAMING_SNAKE_CASE__							)
     model.to(SCREAMING_SNAKE_CASE__							)
     model.eval()

     # first forward pass
     SCREAMING_SNAKE_CASE:List[str]   =   model(
         SCREAMING_SNAKE_CASE__       ,attention_mask=SCREAMING_SNAKE_CASE__       ,encoder_hidden_states=SCREAMING_SNAKE_CASE__       ,encoder_attention_mask=SCREAMING_SNAKE_CASE__       ,use_cache=SCREAMING_SNAKE_CASE__       ,)
     SCREAMING_SNAKE_CASE:int   =   outputs.past_key_values

     # create hypothetical multiple next token and extent to next_input_ids
     SCREAMING_SNAKE_CASE:List[Any]   =   ids_tensor((self.batch_size, 3)       ,config.vocab_size							)
     SCREAMING_SNAKE_CASE:Tuple   =   ids_tensor((self.batch_size, 3)       ,vocab_size=2							)

     # append to next input_ids and
     SCREAMING_SNAKE_CASE:List[str]   =   torch.cat([input_ids, next_tokens]       ,dim=-1							)
     SCREAMING_SNAKE_CASE:Union[str, Any]   =   torch.cat([input_mask, next_mask]       ,dim=-1							)

     SCREAMING_SNAKE_CASE:str   =   model(
         SCREAMING_SNAKE_CASE__       ,attention_mask=SCREAMING_SNAKE_CASE__       ,encoder_hidden_states=SCREAMING_SNAKE_CASE__       ,encoder_attention_mask=SCREAMING_SNAKE_CASE__       ,output_hidden_states=SCREAMING_SNAKE_CASE__       ,)["hidden_states"][0]
     SCREAMING_SNAKE_CASE:Tuple   =   model(
         SCREAMING_SNAKE_CASE__       ,attention_mask=SCREAMING_SNAKE_CASE__       ,encoder_hidden_states=SCREAMING_SNAKE_CASE__       ,encoder_attention_mask=SCREAMING_SNAKE_CASE__       ,past_key_values=SCREAMING_SNAKE_CASE__       ,output_hidden_states=SCREAMING_SNAKE_CASE__       ,)["hidden_states"][0]

     # select random slice
     SCREAMING_SNAKE_CASE:Dict   =   ids_tensor((1,)       ,output_from_past.shape[-1]							).item()
     SCREAMING_SNAKE_CASE:int   =   output_from_no_past[:, -3:, random_slice_idx].detach()
     SCREAMING_SNAKE_CASE:Optional[Any]   =   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(SCREAMING_SNAKE_CASE__       ,SCREAMING_SNAKE_CASE__       ,atol=1e-3							)							)
   def       __UpperCamelCase  (	self     : Optional[int]       ,SCREAMING_SNAKE_CASE__     : Tuple       ,SCREAMING_SNAKE_CASE__     : str       ,SCREAMING_SNAKE_CASE__     : str       ,SCREAMING_SNAKE_CASE__     : List[str]       ,SCREAMING_SNAKE_CASE__     : Union[str, Any]       ,SCREAMING_SNAKE_CASE__     : Optional[int]       ,SCREAMING_SNAKE_CASE__     : Tuple							):
     SCREAMING_SNAKE_CASE:Dict   =   RemBertForQuestionAnswering(config=SCREAMING_SNAKE_CASE__							)
     model.to(SCREAMING_SNAKE_CASE__							)
     model.eval()
     SCREAMING_SNAKE_CASE:Tuple   =   model(
         SCREAMING_SNAKE_CASE__       ,attention_mask=SCREAMING_SNAKE_CASE__       ,token_type_ids=SCREAMING_SNAKE_CASE__       ,start_positions=SCREAMING_SNAKE_CASE__       ,end_positions=SCREAMING_SNAKE_CASE__       ,)
     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     : Dict       ,SCREAMING_SNAKE_CASE__     : List[Any]       ,SCREAMING_SNAKE_CASE__     : str       ,SCREAMING_SNAKE_CASE__     : List[str]       ,SCREAMING_SNAKE_CASE__     : Optional[int]       ,SCREAMING_SNAKE_CASE__     : Union[str, Any]       ,SCREAMING_SNAKE_CASE__     : Optional[Any]       ,SCREAMING_SNAKE_CASE__     : List[str]							):
     SCREAMING_SNAKE_CASE:Any   =   self.num_labels
     SCREAMING_SNAKE_CASE:Tuple   =   RemBertForSequenceClassification(SCREAMING_SNAKE_CASE__							)
     model.to(SCREAMING_SNAKE_CASE__							)
     model.eval()
     SCREAMING_SNAKE_CASE:Any   =   model(SCREAMING_SNAKE_CASE__       ,attention_mask=SCREAMING_SNAKE_CASE__       ,token_type_ids=SCREAMING_SNAKE_CASE__       ,labels=SCREAMING_SNAKE_CASE__							)
     self.parent.assertEqual(result.logits.shape       ,(self.batch_size, self.num_labels)							)
   def       __UpperCamelCase  (	self     : Optional[Any]       ,SCREAMING_SNAKE_CASE__     : Optional[int]       ,SCREAMING_SNAKE_CASE__     : Any       ,SCREAMING_SNAKE_CASE__     : Tuple       ,SCREAMING_SNAKE_CASE__     : List[str]       ,SCREAMING_SNAKE_CASE__     : Dict       ,SCREAMING_SNAKE_CASE__     : Optional[int]       ,SCREAMING_SNAKE_CASE__     : Any							):
     SCREAMING_SNAKE_CASE:Optional[int]   =   self.num_labels
     SCREAMING_SNAKE_CASE:int   =   RemBertForTokenClassification(config=SCREAMING_SNAKE_CASE__							)
     model.to(SCREAMING_SNAKE_CASE__							)
     model.eval()
     SCREAMING_SNAKE_CASE:List[Any]   =   model(SCREAMING_SNAKE_CASE__       ,attention_mask=SCREAMING_SNAKE_CASE__       ,token_type_ids=SCREAMING_SNAKE_CASE__       ,labels=SCREAMING_SNAKE_CASE__							)
     self.parent.assertEqual(result.logits.shape       ,(self.batch_size, self.seq_length, self.num_labels)							)
   def       __UpperCamelCase  (	self     : Tuple       ,SCREAMING_SNAKE_CASE__     : Tuple       ,SCREAMING_SNAKE_CASE__     : Any       ,SCREAMING_SNAKE_CASE__     : str       ,SCREAMING_SNAKE_CASE__     : str       ,SCREAMING_SNAKE_CASE__     : List[Any]       ,SCREAMING_SNAKE_CASE__     : List[Any]       ,SCREAMING_SNAKE_CASE__     : List[str]							):
     SCREAMING_SNAKE_CASE:List[str]   =   self.num_choices
     SCREAMING_SNAKE_CASE:List[Any]   =   RemBertForMultipleChoice(config=SCREAMING_SNAKE_CASE__							)
     model.to(SCREAMING_SNAKE_CASE__							)
     model.eval()
     SCREAMING_SNAKE_CASE:str   =   input_ids.unsqueeze(1							).expand(-1       ,self.num_choices       ,-1							).contiguous()
     SCREAMING_SNAKE_CASE:Optional[Any]   =   token_type_ids.unsqueeze(1							).expand(-1       ,self.num_choices       ,-1							).contiguous()
     SCREAMING_SNAKE_CASE:Any   =   input_mask.unsqueeze(1							).expand(-1       ,self.num_choices       ,-1							).contiguous()
     SCREAMING_SNAKE_CASE:Optional[int]   =   model(
         SCREAMING_SNAKE_CASE__       ,attention_mask=SCREAMING_SNAKE_CASE__       ,token_type_ids=SCREAMING_SNAKE_CASE__       ,labels=SCREAMING_SNAKE_CASE__       ,)
     self.parent.assertEqual(result.logits.shape       ,(self.batch_size, self.num_choices)							)







   def       __UpperCamelCase  (	self     : Any							):
     SCREAMING_SNAKE_CASE:List[Any]   =   self.prepare_config_and_inputs()
     (
         (
         SCREAMING_SNAKE_CASE
     ) ,      (
         SCREAMING_SNAKE_CASE
     ) ,      (
         SCREAMING_SNAKE_CASE
     ) ,      (
         SCREAMING_SNAKE_CASE
     ) ,      (
         SCREAMING_SNAKE_CASE
     ) ,      (
         SCREAMING_SNAKE_CASE
     ) ,      (
         SCREAMING_SNAKE_CASE
     ) ,      
     ):Any   =   config_and_inputs
     SCREAMING_SNAKE_CASE:List[str]   =   {"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
     return config, inputs_dict





@require_torch
class 							_snake_case	(    _a		,					_a		,					unittest.TestCase  ):
   _A    :       Optional[int]      				=   (
       (
           RemBertModel,
           RemBertForMaskedLM,
           RemBertForCausalLM,
           RemBertForMultipleChoice,
           RemBertForQuestionAnswering,
           RemBertForSequenceClassification,
           RemBertForTokenClassification,
       )
       if is_torch_available()
       else ()
   )
   _A    :       List[str]      				=   (RemBertForCausalLM,) if is_torch_available() else ()
   _A    :       str      				=   (
       {
           '''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     : Tuple							):
     SCREAMING_SNAKE_CASE:Dict   =   RemBertModelTester(self							)
     SCREAMING_SNAKE_CASE:Union[str, Any]   =   ConfigTester(self       ,config_class=SCREAMING_SNAKE_CASE__       ,hidden_size=37							)
   def       __UpperCamelCase  (	self     : str							):
     self.config_tester.run_common_tests()
   def       __UpperCamelCase  (	self     : Any							):
     SCREAMING_SNAKE_CASE:Optional[int]   =   self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE__							)
   def       __UpperCamelCase  (	self     : Union[str, Any]							):
     SCREAMING_SNAKE_CASE:List[str]   =   self.model_tester.prepare_config_and_inputs()
     for type in ["absolute", "relative_key", "relative_key_query"]:
       SCREAMING_SNAKE_CASE:str   =   type
       self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE__							)
   def       __UpperCamelCase  (	self     : Union[str, Any]							):
     SCREAMING_SNAKE_CASE:str   =   self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_masked_lm(*SCREAMING_SNAKE_CASE__							)
   def       __UpperCamelCase  (	self     : Tuple							):
     SCREAMING_SNAKE_CASE:List[Any]   =   self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_multiple_choice(*SCREAMING_SNAKE_CASE__							)
   def       __UpperCamelCase  (	self     : Optional[int]							):
     SCREAMING_SNAKE_CASE:List[str]   =   self.model_tester.prepare_config_and_inputs_for_decoder()
     self.model_tester.create_and_check_decoder_model_past_large_inputs(*SCREAMING_SNAKE_CASE__							)
   def       __UpperCamelCase  (	self     : Any							):
     SCREAMING_SNAKE_CASE:Optional[Any]   =   self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_question_answering(*SCREAMING_SNAKE_CASE__							)
   def       __UpperCamelCase  (	self     : int							):
     SCREAMING_SNAKE_CASE:int   =   self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_sequence_classification(*SCREAMING_SNAKE_CASE__							)
   def       __UpperCamelCase  (	self     : Union[str, Any]							):
     SCREAMING_SNAKE_CASE:List[Any]   =   self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_token_classification(*SCREAMING_SNAKE_CASE__							)
   def       __UpperCamelCase  (	self     : Dict							):
     SCREAMING_SNAKE_CASE:Union[str, Any]   =   self.model_tester.prepare_config_and_inputs_for_decoder()
     self.model_tester.create_and_check_model_as_decoder(*SCREAMING_SNAKE_CASE__							)
   def       __UpperCamelCase  (	self     : str							):
     # This regression test was failing with PyTorch < 1.3
     (
         (
         SCREAMING_SNAKE_CASE
     ) ,      (
         SCREAMING_SNAKE_CASE
     ) ,      (
         SCREAMING_SNAKE_CASE
     ) ,      (
         SCREAMING_SNAKE_CASE
     ) ,      (
         SCREAMING_SNAKE_CASE
     ) ,      (
         SCREAMING_SNAKE_CASE
     ) ,      (
         SCREAMING_SNAKE_CASE
     ) ,      (
         SCREAMING_SNAKE_CASE
     ) ,      (
         SCREAMING_SNAKE_CASE
     ) ,      
     ):str   =   self.model_tester.prepare_config_and_inputs_for_decoder()

     SCREAMING_SNAKE_CASE:Optional[int]   =   None

     self.model_tester.create_and_check_model_as_decoder(
         SCREAMING_SNAKE_CASE__       ,SCREAMING_SNAKE_CASE__       ,SCREAMING_SNAKE_CASE__       ,SCREAMING_SNAKE_CASE__       ,SCREAMING_SNAKE_CASE__       ,SCREAMING_SNAKE_CASE__       ,SCREAMING_SNAKE_CASE__       ,SCREAMING_SNAKE_CASE__       ,SCREAMING_SNAKE_CASE__       ,)







   @slow
   def       __UpperCamelCase  (	self     : Any							):
     for model_name in REMBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
       SCREAMING_SNAKE_CASE:Tuple   =   RemBertModel.from_pretrained(SCREAMING_SNAKE_CASE__							)
       self.assertIsNotNone(SCREAMING_SNAKE_CASE__							)





@require_torch
class 							_snake_case	(    unittest.TestCase  ):
   @slow
   def       __UpperCamelCase  (	self     : Tuple							):
     # Test exact values at the last hidden layer
     SCREAMING_SNAKE_CASE:Optional[int]   =   RemBertModel.from_pretrained("google/rembert"							)
     SCREAMING_SNAKE_CASE:Any   =   torch.tensor([[312, 56_498, 313, 2_125, 313]]							)
     SCREAMING_SNAKE_CASE:int   =   torch.tensor([[0, 0, 0, 1, 1]]							)
     with torch.no_grad():
       SCREAMING_SNAKE_CASE:str   =   model(SCREAMING_SNAKE_CASE__       ,token_type_ids=SCREAMING_SNAKE_CASE__       ,output_hidden_states=SCREAMING_SNAKE_CASE__							)

     SCREAMING_SNAKE_CASE:Dict   =   1_152

     SCREAMING_SNAKE_CASE:List[str]   =   torch.Size((1, 5, hidden_size)							)
     self.assertEqual(output["last_hidden_state"].shape       ,SCREAMING_SNAKE_CASE__							)

     SCREAMING_SNAKE_CASE:List[str]   =   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]       ,SCREAMING_SNAKE_CASE__       ,atol=1e-4							)							)









'''simple docstring'''

import argparse
import json
import re
from pathlib import Path

import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import (
    MobileNetVaConfig,
    MobileNetVaForImageClassification,
    MobileNetVaForSemanticSegmentation,
    MobileNetVaImageProcessor,
    load_tf_weights_in_mobilenet_va,
)
from transformers.utils import logging


logging.set_verbosity_info()
A_      =							logging.get_logger(__name__)





def 				A_					(	snake_case	):
  SCREAMING_SNAKE_CASE:Optional[int]   =   MobileNetVaConfig(layer_norm_eps=0.001	)

  if "quant" in model_name:
    raise ValueError("Quantized models are not supported."	)

  SCREAMING_SNAKE_CASE:Any   =   re.match(r"^.*mobilenet_v2_([^_]*)_([^_]*)$" ,    snake_case	)
  if matches:
    SCREAMING_SNAKE_CASE:str   =   float(matches[1]	)
    SCREAMING_SNAKE_CASE:Optional[int]   =   int(matches[2]	)

  if model_name.startswith("deeplabv3_"	):
    SCREAMING_SNAKE_CASE:str   =   8
    SCREAMING_SNAKE_CASE:Optional[int]   =   21
    SCREAMING_SNAKE_CASE:Dict   =   "pascal-voc-id2label.json"
  else:
    # The TensorFlow version of MobileNetV2 predicts 1001 classes instead
    # of the usual 1000. The first class (index 0) is "background".
    SCREAMING_SNAKE_CASE:Optional[int]   =   1001
    SCREAMING_SNAKE_CASE:Tuple   =   "imagenet-1k-id2label.json"

  SCREAMING_SNAKE_CASE:List[str]   =   "huggingface/label-files"
  SCREAMING_SNAKE_CASE:List[Any]   =   json.load(open(hf_hub_download(snake_case ,    snake_case ,    repo_type="dataset"	) ,    "r"	)	)

  if config.num_labels == 1001:
    SCREAMING_SNAKE_CASE:Optional[Any]   =   {int(snake_case	) + 1: v for k, v in idalabel.items()}
    SCREAMING_SNAKE_CASE:Dict   =   "background"
  else:
    SCREAMING_SNAKE_CASE:Union[str, Any]   =   {int(snake_case	): v for k, v in idalabel.items()}

  SCREAMING_SNAKE_CASE:str   =   idalabel
  SCREAMING_SNAKE_CASE:int   =   {v: k for k, v in idalabel.items()}

  return config





def 				A_					(	):
  SCREAMING_SNAKE_CASE:Optional[Any]   =   "http://images.cocodataset.org/val2017/000000039769.jpg"
  SCREAMING_SNAKE_CASE:str   =   Image.open(requests.get(snake_case ,    stream=snake_case	).raw	)
  return im





@torch.no_grad()
def 				A_					(	snake_case ,    snake_case ,    snake_case ,    snake_case=False	):
  SCREAMING_SNAKE_CASE:Any   =   get_mobilenet_va_config(snake_case	)

  # Load 🤗 model
  if model_name.startswith("deeplabv3_"	):
    SCREAMING_SNAKE_CASE:Tuple   =   MobileNetVaForSemanticSegmentation(snake_case	).eval()
  else:
    SCREAMING_SNAKE_CASE:List[str]   =   MobileNetVaForImageClassification(snake_case	).eval()

  # Load weights from TensorFlow checkpoint
  load_tf_weights_in_mobilenet_va(snake_case ,    snake_case ,    snake_case	)

  # Check outputs on an image, prepared by MobileNetV2ImageProcessor
  SCREAMING_SNAKE_CASE:List[str]   =   MobileNetVaImageProcessor(
      crop_size={"width": config.image_size, "height": config.image_size} ,    size={"shortest_edge": config.image_size + 32} ,    )
  SCREAMING_SNAKE_CASE:List[Any]   =   image_processor(images=prepare_img() ,    return_tensors="pt"	)
  SCREAMING_SNAKE_CASE:Union[str, Any]   =   model(**snake_case	)
  SCREAMING_SNAKE_CASE:str   =   outputs.logits

  if model_name.startswith("deeplabv3_"	):
    assert logits.shape == (1, 21, 65, 65)

    if model_name == "deeplabv3_mobilenet_v2_1.0_513":
      SCREAMING_SNAKE_CASE:int   =   torch.tensor(
          [
              [[17.5790, 17.7581, 18.3355], [18.3257, 18.4230, 18.8973], [18.6169, 18.8650, 19.2187]],
              [[-2.1595, -2.0977, -2.3741], [-2.4226, -2.3028, -2.6835], [-2.7819, -2.5991, -2.7706]],
              [[4.2058, 4.8317, 4.7638], [4.4136, 5.0361, 4.9383], [4.5028, 4.9644, 4.8734]],
          ]	)

    else:
      raise ValueError(F'''Unknown model name: {model_name}'''	)

    assert torch.allclose(logits[0, :3, :3, :3] ,    snake_case ,    atol=1e-4	)
  else:
    assert logits.shape == (1, 1001)

    if model_name == "mobilenet_v2_1.4_224":
      SCREAMING_SNAKE_CASE:List[str]   =   torch.tensor([0.0181, -1.0015, 0.4688]	)
    elif model_name == "mobilenet_v2_1.0_224":
      SCREAMING_SNAKE_CASE:Optional[int]   =   torch.tensor([0.2445, -1.1993, 0.1905]	)
    elif model_name == "mobilenet_v2_0.75_160":
      SCREAMING_SNAKE_CASE:Optional[Any]   =   torch.tensor([0.2482, 0.4136, 0.6669]	)
    elif model_name == "mobilenet_v2_0.35_96":
      SCREAMING_SNAKE_CASE:int   =   torch.tensor([0.1451, -0.4624, 0.7192]	)
    else:
      SCREAMING_SNAKE_CASE:Optional[Any]   =   None

    if expected_logits is not None:
      assert torch.allclose(logits[0, :3] ,    snake_case ,    atol=1e-4	)

  Path(snake_case	).mkdir(exist_ok=snake_case	)
  print(F'''Saving model {model_name} to {pytorch_dump_folder_path}'''	)
  model.save_pretrained(snake_case	)
  print(F'''Saving image processor to {pytorch_dump_folder_path}'''	)
  image_processor.save_pretrained(snake_case	)

  if push_to_hub:
    print("Pushing to the hub..."	)
    SCREAMING_SNAKE_CASE:Tuple   =   "google/" + model_name
    image_processor.push_to_hub(snake_case	)
    model.push_to_hub(snake_case	)


if __name__ == "__main__":
    A_      =							argparse.ArgumentParser()
    # Required parameters
    parser.add_argument(
        "--model_name",
        default="mobilenet_v2_1.0_224",
        type=str,
        help="Name of the MobileNetV2 model you'd like to convert. Should in the form 'mobilenet_v2_<depth>_<size>'.",
    )
    parser.add_argument(
        "--checkpoint_path", required=True, type=str, help="Path to the original TensorFlow checkpoint (.ckpt file)."
    )
    parser.add_argument(
        "--pytorch_dump_folder_path", required=True, type=str, help="Path to the output PyTorch model directory."
    )
    parser.add_argument(
        "--push_to_hub", action="store_true", help="Whether or not to push the converted model to the 🤗 hub."
    )

    A_      =							parser.parse_args()
    convert_movilevit_checkpoint(
        args.model_name, args.checkpoint_path, args.pytorch_dump_folder_path, args.push_to_hub
    )






