




"""simple docstring"""



import unittest

from transformers import YosoConfig, 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 (
          YosoForMaskedLM,
          YosoForMultipleChoice,
          YosoForQuestionAnswering,
          YosoForSequenceClassification,
          YosoForTokenClassification,
          YosoModel,
      )
      from transformers.models.yoso.modeling_yoso import YOSO_PRETRAINED_MODEL_ARCHIVE_LIST




class       UpperCAmelCase_    :




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




     '''simple docstring'''

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




    def 				__magic_name__				(       self		:			List[Any] ) ->				Optional[int]:




     '''simple docstring'''

     A__			=							ids_tensor([self.batch_size, self.seq_length]      ,					self.vocab_size )

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

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

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

     A__			=							self.get_config()

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




    def 				__magic_name__				(       self		:			Dict ) ->				Optional[int]:




     '''simple docstring'''

     return YosoConfig(
         vocab_size=self.vocab_size      ,					hidden_size=self.hidden_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=snake_case_      ,					initializer_range=self.initializer_range      ,					)




    def 				__magic_name__				(       self		:			Optional[int] ) ->				str:




     '''simple docstring'''

     A__			=							self.get_config()
     A__			=							300
     return config




    def 				__magic_name__				(       self		:			List[Any] ) ->				Optional[Any]:




     '''simple docstring'''

     (
         (
         A__
     ), (
         A__
     ), (
         A__
     ), (
         A__
     ), (
         A__
     ), (
         A__
     ), (
         A__
     ), 
     )			=							self.prepare_config_and_inputs()

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

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




    def 				__magic_name__				(       self		:			int      ,					snake_case_		:			Optional[int]      ,					snake_case_		:			Any      ,					snake_case_		:			List[str]      ,					snake_case_		:			Optional[int]      ,					snake_case_		:			str      ,					snake_case_		:			Tuple      ,					snake_case_		:			Tuple ) ->				List[str]:




     '''simple docstring'''

     A__			=							YosoModel(config=snake_case_ )
     model.to(snake_case_ )
     model.eval()
     A__			=							model(snake_case_      ,					attention_mask=snake_case_      ,					token_type_ids=snake_case_ )
     A__			=							model(snake_case_      ,					token_type_ids=snake_case_ )
     A__			=							model(snake_case_ )
     self.parent.assertEqual(result.last_hidden_state.shape      ,					(self.batch_size, self.seq_length, self.hidden_size) )




    def 				__magic_name__				(       self		:			Union[str, Any]      ,					snake_case_		:			int      ,					snake_case_		:			Union[str, Any]      ,					snake_case_		:			Tuple      ,					snake_case_		:			int      ,					snake_case_		:			Union[str, Any]      ,					snake_case_		:			Tuple      ,					snake_case_		:			str      ,					snake_case_		:			List[Any]      ,					snake_case_		:			Optional[int]      ,					) ->				List[Any]:




     '''simple docstring'''

     A__			=							True
     A__			=							YosoModel(snake_case_ )
     model.to(snake_case_ )
     model.eval()
     A__			=							model(
         snake_case_      ,					attention_mask=snake_case_      ,					token_type_ids=snake_case_      ,					encoder_hidden_states=snake_case_      ,					encoder_attention_mask=snake_case_      ,					)
     A__			=							model(
         snake_case_      ,					attention_mask=snake_case_      ,					token_type_ids=snake_case_      ,					encoder_hidden_states=snake_case_      ,					)
     A__			=							model(snake_case_      ,					attention_mask=snake_case_      ,					token_type_ids=snake_case_ )
     self.parent.assertEqual(result.last_hidden_state.shape      ,					(self.batch_size, self.seq_length, self.hidden_size) )




    def 				__magic_name__				(       self		:			Optional[int]      ,					snake_case_		:			Optional[int]      ,					snake_case_		:			Tuple      ,					snake_case_		:			Tuple      ,					snake_case_		:			Dict      ,					snake_case_		:			List[Any]      ,					snake_case_		:			List[str]      ,					snake_case_		:			List[Any] ) ->				Any:




     '''simple docstring'''

     A__			=							YosoForMaskedLM(config=snake_case_ )
     model.to(snake_case_ )
     model.eval()
     A__			=							model(snake_case_      ,					attention_mask=snake_case_      ,					token_type_ids=snake_case_      ,					labels=snake_case_ )
     self.parent.assertEqual(result.logits.shape      ,					(self.batch_size, self.seq_length, self.vocab_size) )




    def 				__magic_name__				(       self		:			Union[str, Any]      ,					snake_case_		:			int      ,					snake_case_		:			List[str]      ,					snake_case_		:			Union[str, Any]      ,					snake_case_		:			Union[str, Any]      ,					snake_case_		:			Optional[int]      ,					snake_case_		:			int      ,					snake_case_		:			Optional[Any] ) ->				Optional[Any]:




     '''simple docstring'''

     A__			=							YosoForQuestionAnswering(config=snake_case_ )
     model.to(snake_case_ )
     model.eval()
     A__			=							model(
         snake_case_      ,					attention_mask=snake_case_      ,					token_type_ids=snake_case_      ,					start_positions=snake_case_      ,					end_positions=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 				__magic_name__				(       self		:			Optional[Any]      ,					snake_case_		:			int      ,					snake_case_		:			Optional[Any]      ,					snake_case_		:			str      ,					snake_case_		:			Tuple      ,					snake_case_		:			List[str]      ,					snake_case_		:			Optional[int]      ,					snake_case_		:			List[Any] ) ->				List[str]:




     '''simple docstring'''

     A__			=							self.num_labels
     A__			=							YosoForSequenceClassification(snake_case_ )
     model.to(snake_case_ )
     model.eval()
     A__			=							model(snake_case_      ,					attention_mask=snake_case_      ,					token_type_ids=snake_case_      ,					labels=snake_case_ )
     self.parent.assertEqual(result.logits.shape      ,					(self.batch_size, self.num_labels) )




    def 				__magic_name__				(       self		:			Optional[int]      ,					snake_case_		:			Dict      ,					snake_case_		:			List[Any]      ,					snake_case_		:			List[Any]      ,					snake_case_		:			Union[str, Any]      ,					snake_case_		:			int      ,					snake_case_		:			int      ,					snake_case_		:			Tuple ) ->				Tuple:




     '''simple docstring'''

     A__			=							self.num_labels
     A__			=							YosoForTokenClassification(config=snake_case_ )
     model.to(snake_case_ )
     model.eval()
     A__			=							model(snake_case_      ,					attention_mask=snake_case_      ,					token_type_ids=snake_case_      ,					labels=snake_case_ )
     self.parent.assertEqual(result.logits.shape      ,					(self.batch_size, self.seq_length, self.num_labels) )




    def 				__magic_name__				(       self		:			Any      ,					snake_case_		:			Optional[int]      ,					snake_case_		:			int      ,					snake_case_		:			List[Any]      ,					snake_case_		:			int      ,					snake_case_		:			Optional[Any]      ,					snake_case_		:			Tuple      ,					snake_case_		:			Any ) ->				Tuple:




     '''simple docstring'''

     A__			=							self.num_choices
     A__			=							YosoForMultipleChoice(config=snake_case_ )
     model.to(snake_case_ )
     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(
         snake_case_      ,					attention_mask=snake_case_      ,					token_type_ids=snake_case_      ,					labels=snake_case_      ,					)
     self.parent.assertEqual(result.logits.shape      ,					(self.batch_size, self.num_choices) )




    def 				__magic_name__				(       self		:			Dict ) ->				Optional[int]:




     '''simple docstring'''

     A__			=							self.prepare_config_and_inputs()
     (
         (
         A__
     ), (
         A__
     ), (
         A__
     ), (
         A__
     ), (
         A__
     ), (
         A__
     ), (
         A__
     ), 
     )			=							config_and_inputs
     A__			=							{"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
     return config, inputs_dict






@require_torch
class       UpperCAmelCase_    (		A_,  A_,  unittest.TestCase				):
    lowercase__		   = (
        (
            YosoModel,
            YosoForMaskedLM,
            YosoForMultipleChoice,
            YosoForQuestionAnswering,
            YosoForSequenceClassification,
            YosoForTokenClassification,
        )
        if is_torch_available()
        else ()
    )
    lowercase__		   = False
    lowercase__		   = False
    lowercase__		   = False

    lowercase__		   = ()
    lowercase__		   = (
        {
            '''feature-extraction''': YosoModel,
            '''fill-mask''': YosoForMaskedLM,
            '''question-answering''': YosoForQuestionAnswering,
            '''text-classification''': YosoForSequenceClassification,
            '''token-classification''': YosoForTokenClassification,
            '''zero-shot''': YosoForSequenceClassification,
        }
        if is_torch_available()
        else {}
    )




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




     '''simple docstring'''

     A__			=							YosoModelTester(self )
     A__			=							ConfigTester(self      ,					config_class=snake_case_      ,					hidden_size=37 )




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




     '''simple docstring'''

     self.config_tester.run_common_tests()




    def 				__magic_name__				(       self		:			Optional[int] ) ->				str:




     '''simple docstring'''

     A__			=							self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_model(*snake_case_ )




    def 				__magic_name__				(       self		:			Optional[int] ) ->				Any:




     '''simple docstring'''

     A__			=							self.model_tester.prepare_config_and_inputs()
     for type in ["absolute", "relative_key", "relative_key_query"]:
      A__			=							type
      self.model_tester.create_and_check_model(*snake_case_ )




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




     '''simple docstring'''

     A__			=							self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_masked_lm(*snake_case_ )




    def 				__magic_name__				(       self		:			Dict ) ->				int:




     '''simple docstring'''

     A__			=							self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_multiple_choice(*snake_case_ )




    def 				__magic_name__				(       self		:			str ) ->				Optional[Any]:




     '''simple docstring'''

     A__			=							self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_question_answering(*snake_case_ )




    def 				__magic_name__				(       self		:			Optional[int] ) ->				Dict:




     '''simple docstring'''

     A__			=							self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_sequence_classification(*snake_case_ )




    def 				__magic_name__				(       self		:			Tuple ) ->				int:




     '''simple docstring'''

     A__			=							self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_token_classification(*snake_case_ )




    @slow
    def 				__magic_name__				(       self		:			str ) ->				Optional[int]:




     '''simple docstring'''

     for model_name in YOSO_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
      A__			=							YosoModel.from_pretrained(snake_case_ )
      self.assertIsNotNone(snake_case_ )




    def 				__magic_name__				(       self		:			Optional[int] ) ->				Optional[int]:




     '''simple docstring'''

     return






@require_torch
class       UpperCAmelCase_    (		unittest.TestCase				):




    @slow
    def 				__magic_name__				(       self		:			Any ) ->				Union[str, Any]:




     '''simple docstring'''

     A__			=							YosoModel.from_pretrained("uw-madison/yoso-4096" )
     A__			=							torch.tensor([[0, 1, 2, 3, 4, 5]] )

     with torch.no_grad():
      A__			=							model(snake_case_ )[0]

     A__			=							torch.Size((1, 6, 768) )
     self.assertEqual(output.shape      ,					snake_case_ )

     A__			=							torch.tensor(
         [[[-0.0611, 0.1242, 0.0840], [0.0280, -0.0048, 0.1125], [0.0106, 0.0226, 0.0751]]] )

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




    @slow
    def 				__magic_name__				(       self		:			List[str] ) ->				Optional[int]:




     '''simple docstring'''

     A__			=							YosoForMaskedLM.from_pretrained("uw-madison/yoso-4096" )
     A__			=							torch.tensor([[0, 1, 2, 3, 4, 5]] )

     with torch.no_grad():
      A__			=							model(snake_case_ )[0]

     A__			=							50_265

     A__			=							torch.Size((1, 6, vocab_size) )
     self.assertEqual(output.shape      ,					snake_case_ )

     A__			=							torch.tensor(
         [[[-2.1313, -3.7285, -2.2407], [-2.7047, -3.3314, -2.6408], [0.0629, -2.5166, -0.3356]]] )

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




    @slow
    def 				__magic_name__				(       self		:			List[Any] ) ->				List[str]:




     '''simple docstring'''

     A__			=							YosoForMaskedLM.from_pretrained("uw-madison/yoso-4096" )
     A__			=							torch.arange(4_096 ).unsqueeze(0 )

     with torch.no_grad():
      A__			=							model(snake_case_ )[0]

     A__			=							50_265

     A__			=							torch.Size((1, 4_096, vocab_size) )
     self.assertEqual(output.shape      ,					snake_case_ )

     A__			=							torch.tensor(
         [[[-2.3914, -4.3742, -5.0956], [-4.0988, -4.2384, -7.0406], [-3.1427, -3.7192, -6.6800]]] )

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











"""simple docstring"""



import argparse
import json

from transformers import GPTNeoConfig, GPTNeoForCausalLM, load_tf_weights_in_gpt_neo
from transformers.utils import logging


logging.set_verbosity_info()







def     _SCREAMING_SNAKE_CASE  (				lowercase_     ,					lowercase_     ,					lowercase_		)		->						Union[str, Any]:
 # Initialise PyTorch model
 A__			=							json.load(open(lowercase_     ,					"r"		)		)
 A__			=							GPTNeoConfig(
     hidden_size=config_json["n_embd"]     ,					num_layers=config_json["n_layer"]     ,					num_heads=config_json["n_head"]     ,					attention_types=config_json["attention_types"]     ,					max_position_embeddings=config_json["n_positions"]     ,					resid_dropout=config_json["res_dropout"]     ,					embed_dropout=config_json["embed_dropout"]     ,					attention_dropout=config_json["attn_dropout"]     ,					)
 print(f"""Building PyTorch model from configuration: {config}"""		)
 A__			=							GPTNeoForCausalLM(lowercase_		)

 # Load weights from tf checkpoint
 load_tf_weights_in_gpt_neo(lowercase_     ,					lowercase_     ,					lowercase_		)

 # Save pytorch-model
 print(f"""Save PyTorch model to {pytorch_dump_path}"""		)
 model.save_pretrained(lowercase_		)


if __name__ == "__main__":
      SCREAMING_SNAKE_CASE							=			argparse.ArgumentParser()
      # Required parameters
      parser.add_argument(
          "--tf_checkpoint_path", default=None, type=str, required=True, help="Path to the TensorFlow checkpoint path."
      )
      parser.add_argument(
          "--config_file",
          default=None,
          type=str,
          required=True,
          help=(
              "The config json file corresponding to the pre-trained mesh-tf model. \n"
              "This specifies the model architecture."
          ),
      )
      parser.add_argument(
          "--pytorch_dump_path", default=None, type=str, required=True, help="Path to the output PyTorch model."
      )
      SCREAMING_SNAKE_CASE							=			parser.parse_args()
      convert_tf_checkpoint_to_pytorch(args.tf_checkpoint_path, args.config_file, args.pytorch_dump_path)





