


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    _a				:




	def __init__(   self		:						Dict						,   _SCREAMING_SNAKE_CASE		:						List[Any]						,   _SCREAMING_SNAKE_CASE		:						str=13						,   _SCREAMING_SNAKE_CASE		:						Optional[int]=7						,   _SCREAMING_SNAKE_CASE		:						Any=True						,   _SCREAMING_SNAKE_CASE		:						int=True						,   _SCREAMING_SNAKE_CASE		:						Union[str, Any]=True						,   _SCREAMING_SNAKE_CASE		:						int=True						,   _SCREAMING_SNAKE_CASE		:						Any=99						,   _SCREAMING_SNAKE_CASE		:						List[Any]=32						,   _SCREAMING_SNAKE_CASE		:						Any=5						,   _SCREAMING_SNAKE_CASE		:						int=4						,   _SCREAMING_SNAKE_CASE		:						List[Any]=37						,   _SCREAMING_SNAKE_CASE		:						str="gelu"						,   _SCREAMING_SNAKE_CASE		:						Tuple=0.1						,   _SCREAMING_SNAKE_CASE		:						Any=0.1						,   _SCREAMING_SNAKE_CASE		:						Optional[Any]=512						,   _SCREAMING_SNAKE_CASE		:						Any=16						,   _SCREAMING_SNAKE_CASE		:						Optional[int]=2						,   _SCREAMING_SNAKE_CASE		:						Optional[Any]=0.02						,   _SCREAMING_SNAKE_CASE		:						Tuple=3						,   _SCREAMING_SNAKE_CASE		:						str=4						,   _SCREAMING_SNAKE_CASE		:						Union[str, Any]=None						,   )->			Optional[int]:
			lowerCAmelCase__  :     Any       	=   parent
			lowerCAmelCase__  :     Tuple       	=   batch_size
			lowerCAmelCase__  :     Union[str, Any]       	=   seq_length
			lowerCAmelCase__  :     List[str]       	=   is_training
			lowerCAmelCase__  :     Dict       	=   use_input_mask
			lowerCAmelCase__  :     Any       	=   use_token_type_ids
			lowerCAmelCase__  :     int       	=   use_labels
			lowerCAmelCase__  :     Optional[int]       	=   vocab_size
			lowerCAmelCase__  :     Any       	=   hidden_size
			lowerCAmelCase__  :     str       	=   num_hidden_layers
			lowerCAmelCase__  :     Union[str, Any]       	=   num_attention_heads
			lowerCAmelCase__  :     Optional[int]       	=   intermediate_size
			lowerCAmelCase__  :     List[str]       	=   hidden_act
			lowerCAmelCase__  :     Optional[int]       	=   hidden_dropout_prob
			lowerCAmelCase__  :     int       	=   attention_probs_dropout_prob
			lowerCAmelCase__  :     int       	=   max_position_embeddings
			lowerCAmelCase__  :     Dict       	=   type_vocab_size
			lowerCAmelCase__  :     Optional[int]       	=   type_sequence_label_size
			lowerCAmelCase__  :     List[str]       	=   initializer_range
			lowerCAmelCase__  :     Any       	=   num_labels
			lowerCAmelCase__  :     Optional[int]       	=   num_choices
			lowerCAmelCase__  :     Optional[Any]       	=   scope




	def 				UpperCAmelCase__(   self		:						int						)->			Optional[int]:
			lowerCAmelCase__  :     List[Any]       	=   ids_tensor([self.batch_size, self.seq_length]						,   self.vocab_size						)

			lowerCAmelCase__  :     List[Any]       	=   None
			if self.use_input_mask:
					lowerCAmelCase__  :     Optional[int]       	=   random_attention_mask([self.batch_size, self.seq_length]						)

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

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

			lowerCAmelCase__  :     int       	=   self.get_config()

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




	def 				UpperCAmelCase__(   self		:						Optional[Any]						)->			int:
			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=A_						,   initializer_range=self.initializer_range						,   )




	def 				UpperCAmelCase__(   self		:						Union[str, Any]						)->			Union[str, Any]:
			lowerCAmelCase__  :     List[str]       	=   self.get_config()
			lowerCAmelCase__  :     List[Any]       	=   300
			return config




	def 				UpperCAmelCase__(   self		:						Dict						)->			Union[str, Any]:
			(
			    (
			    lowerCAmelCase__     
			)	,      (
			    lowerCAmelCase__     
			)	,      (
			    lowerCAmelCase__     
			)	,      (
			    lowerCAmelCase__     
			)	,      (
			    lowerCAmelCase__     
			)	,      (
			    lowerCAmelCase__     
			)	,      (
			    lowerCAmelCase__     
			)	,           
			)  :     str       	=   self.prepare_config_and_inputs()

			lowerCAmelCase__  :     Any       	=   True
			lowerCAmelCase__  :     List[str]       	=   floats_tensor([self.batch_size, self.seq_length, self.hidden_size]						)
			lowerCAmelCase__  :     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		:						int						,   _SCREAMING_SNAKE_CASE		:						int						,   _SCREAMING_SNAKE_CASE		:						Union[str, Any]						,   _SCREAMING_SNAKE_CASE		:						Optional[int]						,   _SCREAMING_SNAKE_CASE		:						Union[str, Any]						,   _SCREAMING_SNAKE_CASE		:						Optional[Any]						,   _SCREAMING_SNAKE_CASE		:						Optional[int]						,   _SCREAMING_SNAKE_CASE		:						int						)->			Optional[int]:
			lowerCAmelCase__  :     List[str]       	=   YosoModel(config=A_						)
			model.to(A_						)
			model.eval()
			lowerCAmelCase__  :     str       	=   model(A_						,   attention_mask=A_						,   token_type_ids=A_						)
			lowerCAmelCase__  :     Optional[Any]       	=   model(A_						,   token_type_ids=A_						)
			lowerCAmelCase__  :     List[str]       	=   model(A_						)
			self.parent.assertEqual(result.last_hidden_state.shape						,   (self.batch_size, self.seq_length, self.hidden_size)						)




	def 				UpperCAmelCase__(   self		:						List[str]						,   _SCREAMING_SNAKE_CASE		:						Optional[int]						,   _SCREAMING_SNAKE_CASE		:						Union[str, Any]						,   _SCREAMING_SNAKE_CASE		:						Optional[Any]						,   _SCREAMING_SNAKE_CASE		:						Optional[int]						,   _SCREAMING_SNAKE_CASE		:						Tuple						,   _SCREAMING_SNAKE_CASE		:						str						,   _SCREAMING_SNAKE_CASE		:						str						,   _SCREAMING_SNAKE_CASE		:						List[str]						,   _SCREAMING_SNAKE_CASE		:						int						,   )->			int:
			lowerCAmelCase__  :     Optional[int]       	=   True
			lowerCAmelCase__  :     List[Any]       	=   YosoModel(A_						)
			model.to(A_						)
			model.eval()
			lowerCAmelCase__  :     str       	=   model(
			    A_						,   attention_mask=A_						,   token_type_ids=A_						,   encoder_hidden_states=A_						,   encoder_attention_mask=A_						,   )
			lowerCAmelCase__  :     Tuple       	=   model(
			    A_						,   attention_mask=A_						,   token_type_ids=A_						,   encoder_hidden_states=A_						,   )
			lowerCAmelCase__  :     str       	=   model(A_						,   attention_mask=A_						,   token_type_ids=A_						)
			self.parent.assertEqual(result.last_hidden_state.shape						,   (self.batch_size, self.seq_length, self.hidden_size)						)




	def 				UpperCAmelCase__(   self		:						Union[str, Any]						,   _SCREAMING_SNAKE_CASE		:						Optional[Any]						,   _SCREAMING_SNAKE_CASE		:						Optional[int]						,   _SCREAMING_SNAKE_CASE		:						Optional[int]						,   _SCREAMING_SNAKE_CASE		:						Tuple						,   _SCREAMING_SNAKE_CASE		:						int						,   _SCREAMING_SNAKE_CASE		:						Optional[Any]						,   _SCREAMING_SNAKE_CASE		:						Union[str, Any]						)->			Optional[Any]:
			lowerCAmelCase__  :     Any       	=   YosoForMaskedLM(config=A_						)
			model.to(A_						)
			model.eval()
			lowerCAmelCase__  :     Tuple       	=   model(A_						,   attention_mask=A_						,   token_type_ids=A_						,   labels=A_						)
			self.parent.assertEqual(result.logits.shape						,   (self.batch_size, self.seq_length, self.vocab_size)						)




	def 				UpperCAmelCase__(   self		:						Union[str, Any]						,   _SCREAMING_SNAKE_CASE		:						Tuple						,   _SCREAMING_SNAKE_CASE		:						List[Any]						,   _SCREAMING_SNAKE_CASE		:						Any						,   _SCREAMING_SNAKE_CASE		:						Tuple						,   _SCREAMING_SNAKE_CASE		:						Dict						,   _SCREAMING_SNAKE_CASE		:						Optional[int]						,   _SCREAMING_SNAKE_CASE		:						Union[str, Any]						)->			Optional[Any]:
			lowerCAmelCase__  :     Dict       	=   YosoForQuestionAnswering(config=A_						)
			model.to(A_						)
			model.eval()
			lowerCAmelCase__  :     List[Any]       	=   model(
			    A_						,   attention_mask=A_						,   token_type_ids=A_						,   start_positions=A_						,   end_positions=A_						,   )
			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		:						List[Any]						,   _SCREAMING_SNAKE_CASE		:						Any						,   _SCREAMING_SNAKE_CASE		:						int						,   _SCREAMING_SNAKE_CASE		:						Union[str, Any]						,   _SCREAMING_SNAKE_CASE		:						Optional[Any]						,   _SCREAMING_SNAKE_CASE		:						Optional[Any]						,   _SCREAMING_SNAKE_CASE		:						int						,   _SCREAMING_SNAKE_CASE		:						Union[str, Any]						)->			Union[str, Any]:
			lowerCAmelCase__  :     int       	=   self.num_labels
			lowerCAmelCase__  :     Any       	=   YosoForSequenceClassification(A_						)
			model.to(A_						)
			model.eval()
			lowerCAmelCase__  :     Union[str, Any]       	=   model(A_						,   attention_mask=A_						,   token_type_ids=A_						,   labels=A_						)
			self.parent.assertEqual(result.logits.shape						,   (self.batch_size, self.num_labels)						)




	def 				UpperCAmelCase__(   self		:						Optional[int]						,   _SCREAMING_SNAKE_CASE		:						List[Any]						,   _SCREAMING_SNAKE_CASE		:						List[str]						,   _SCREAMING_SNAKE_CASE		:						Optional[Any]						,   _SCREAMING_SNAKE_CASE		:						Union[str, Any]						,   _SCREAMING_SNAKE_CASE		:						List[str]						,   _SCREAMING_SNAKE_CASE		:						Any						,   _SCREAMING_SNAKE_CASE		:						Union[str, Any]						)->			str:
			lowerCAmelCase__  :     Union[str, Any]       	=   self.num_labels
			lowerCAmelCase__  :     Dict       	=   YosoForTokenClassification(config=A_						)
			model.to(A_						)
			model.eval()
			lowerCAmelCase__  :     str       	=   model(A_						,   attention_mask=A_						,   token_type_ids=A_						,   labels=A_						)
			self.parent.assertEqual(result.logits.shape						,   (self.batch_size, self.seq_length, self.num_labels)						)




	def 				UpperCAmelCase__(   self		:						Union[str, Any]						,   _SCREAMING_SNAKE_CASE		:						str						,   _SCREAMING_SNAKE_CASE		:						Tuple						,   _SCREAMING_SNAKE_CASE		:						int						,   _SCREAMING_SNAKE_CASE		:						Optional[Any]						,   _SCREAMING_SNAKE_CASE		:						Optional[Any]						,   _SCREAMING_SNAKE_CASE		:						str						,   _SCREAMING_SNAKE_CASE		:						Optional[Any]						)->			Any:
			lowerCAmelCase__  :     Union[str, Any]       	=   self.num_choices
			lowerCAmelCase__  :     Dict       	=   YosoForMultipleChoice(config=A_						)
			model.to(A_						)
			model.eval()
			lowerCAmelCase__  :     Tuple       	=   input_ids.unsqueeze(1						).expand(-1						,   self.num_choices						,   -1						).contiguous()
			lowerCAmelCase__  :     Any       	=   token_type_ids.unsqueeze(1						).expand(-1						,   self.num_choices						,   -1						).contiguous()
			lowerCAmelCase__  :     Tuple       	=   input_mask.unsqueeze(1						).expand(-1						,   self.num_choices						,   -1						).contiguous()
			lowerCAmelCase__  :     Dict       	=   model(
			    A_						,   attention_mask=A_						,   token_type_ids=A_						,   labels=A_						,   )
			self.parent.assertEqual(result.logits.shape						,   (self.batch_size, self.num_choices)						)




	def 				UpperCAmelCase__(   self		:						Optional[Any]						)->			Optional[int]:
			lowerCAmelCase__  :     Union[str, Any]       	=   self.prepare_config_and_inputs()
			(
			    (
			    lowerCAmelCase__     
			)	,      (
			    lowerCAmelCase__     
			)	,      (
			    lowerCAmelCase__     
			)	,      (
			    lowerCAmelCase__     
			)	,      (
			    lowerCAmelCase__     
			)	,      (
			    lowerCAmelCase__     
			)	,      (
			    lowerCAmelCase__     
			)	,           
			)  :     Optional[int]       	=   config_and_inputs
			lowerCAmelCase__  :     Optional[Any]       	=   {'''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):
	_a    :    int        =				(
	    (
	        YosoModel,
	        YosoForMaskedLM,
	        YosoForMultipleChoice,
	        YosoForQuestionAnswering,
	        YosoForSequenceClassification,
	        YosoForTokenClassification,
	    )
	    if is_torch_available()
	    else ()
	)
	_a    :    Optional[int]        =				False
	_a    :    Union[str, Any]        =				False
	_a    :    Optional[int]        =				False

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




	def 				UpperCAmelCase__(   self		:						List[Any]						)->			Union[str, Any]:
			lowerCAmelCase__  :     Tuple       	=   YosoModelTester(self						)
			lowerCAmelCase__  :     Optional[Any]       	=   ConfigTester(self						,   config_class=A_						,   hidden_size=37						)




	def 				UpperCAmelCase__(   self		:						str						)->			List[str]:
			self.config_tester.run_common_tests()




	def 				UpperCAmelCase__(   self		:						Union[str, Any]						)->			Optional[int]:
			lowerCAmelCase__  :     Dict       	=   self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_model(*A_						)




	def 				UpperCAmelCase__(   self		:						Union[str, Any]						)->			Optional[int]:
			lowerCAmelCase__  :     str       	=   self.model_tester.prepare_config_and_inputs()
			for type in ["absolute", "relative_key", "relative_key_query"]:
					lowerCAmelCase__  :     List[Any]       	=   type
					self.model_tester.create_and_check_model(*A_						)




	def 				UpperCAmelCase__(   self		:						Dict						)->			Tuple:
			lowerCAmelCase__  :     int       	=   self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_for_masked_lm(*A_						)




	def 				UpperCAmelCase__(   self		:						Optional[Any]						)->			Any:
			lowerCAmelCase__  :     List[Any]       	=   self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_for_multiple_choice(*A_						)




	def 				UpperCAmelCase__(   self		:						List[Any]						)->			List[str]:
			lowerCAmelCase__  :     int       	=   self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_for_question_answering(*A_						)




	def 				UpperCAmelCase__(   self		:						List[Any]						)->			Optional[Any]:
			lowerCAmelCase__  :     Tuple       	=   self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_for_sequence_classification(*A_						)




	def 				UpperCAmelCase__(   self		:						Optional[Any]						)->			Dict:
			lowerCAmelCase__  :     Optional[Any]       	=   self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_for_token_classification(*A_						)




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



			for model_name in YOSO_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
					lowerCAmelCase__  :     Dict       	=   YosoModel.from_pretrained(A_						)
					self.assertIsNotNone(A_						)




	def 				UpperCAmelCase__(   self		:						int						)->			List[str]:
			return





@require_torch
class    _a				(							unittest.TestCase):




	@slow
	def 				UpperCAmelCase__(   self		:						int						)->			Optional[int]:
			lowerCAmelCase__  :     Tuple       	=   YosoModel.from_pretrained('''uw-madison/yoso-4096'''						)
			lowerCAmelCase__  :     Optional[int]       	=   torch.tensor([[0, 1, 2, 3, 4, 5]]						)

			with torch.no_grad():
					lowerCAmelCase__  :     Optional[Any]       	=   model(A_						)[0]

			lowerCAmelCase__  :     Optional[int]       	=   torch.Size((1, 6, 768)						)
			self.assertEqual(output.shape						,   A_						)

			lowerCAmelCase__  :     int       	=   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]						,   A_						,   atol=1E-4						)						)




	@slow
	def 				UpperCAmelCase__(   self		:						List[Any]						)->			Optional[int]:
			lowerCAmelCase__  :     Tuple       	=   YosoForMaskedLM.from_pretrained('''uw-madison/yoso-4096'''						)
			lowerCAmelCase__  :     Optional[Any]       	=   torch.tensor([[0, 1, 2, 3, 4, 5]]						)

			with torch.no_grad():
					lowerCAmelCase__  :     Tuple       	=   model(A_						)[0]

			lowerCAmelCase__  :     int       	=   5_0265

			lowerCAmelCase__  :     Union[str, Any]       	=   torch.Size((1, 6, vocab_size)						)
			self.assertEqual(output.shape						,   A_						)

			lowerCAmelCase__  :     List[Any]       	=   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]						,   A_						,   atol=1E-4						)						)




	@slow
	def 				UpperCAmelCase__(   self		:						Optional[Any]						)->			List[Any]:
			lowerCAmelCase__  :     List[str]       	=   YosoForMaskedLM.from_pretrained('''uw-madison/yoso-4096'''						)
			lowerCAmelCase__  :     Optional[Any]       	=   torch.arange(4096						).unsqueeze(0						)

			with torch.no_grad():
					lowerCAmelCase__  :     Any       	=   model(A_						)[0]

			lowerCAmelCase__  :     List[str]       	=   5_0265

			lowerCAmelCase__  :     int       	=   torch.Size((1, 4096, vocab_size)						)
			self.assertEqual(output.shape						,   A_						)

			lowerCAmelCase__  :     List[Any]       	=   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]						,   A_						,   atol=1E-4						)						)









"""simple docstring"""






import argparse

from transformers import CanineConfig, CanineModel, CanineTokenizer, load_tf_weights_in_canine
from transformers.utils import logging


logging.set_verbosity_info()


def 							A__ (     UpperCamelCase__       ,							UpperCamelCase__    ):


  '''simple docstring'''


  _SCREAMING_SNAKE_CASE        =			CanineConfig()
  _SCREAMING_SNAKE_CASE        =			CanineModel(UpperCamelCase__    )
  model.eval()

  print(F'''Building PyTorch model from configuration: {config}'''    )

  # Load weights from tf checkpoint
  load_tf_weights_in_canine(UpperCamelCase__       ,							UpperCamelCase__       ,							UpperCamelCase__    )

  # Save pytorch-model (weights and configuration)
  print(F'''Save PyTorch model to {pytorch_dump_path}'''    )
  model.save_pretrained(UpperCamelCase__    )

  # Save tokenizer files
  _SCREAMING_SNAKE_CASE        =			CanineTokenizer()
  print(F'''Save tokenizer files to {pytorch_dump_path}'''    )
  tokenizer.save_pretrained(UpperCamelCase__    )


if __name__ == "__main__":
    lowerCamelCase				:     Dict								=  argparse.ArgumentParser()
    # Required parameters
    parser.add_argument(
        """--tf_checkpoint_path""",
        default=None,
        type=str,
        required=True,
        help="""Path to the TensorFlow checkpoint. Should end with model.ckpt""",
    )
    parser.add_argument(
        """--pytorch_dump_path""",
        default=None,
        type=str,
        required=True,
        help="""Path to a folder where the PyTorch model will be placed.""",
    )
    lowerCamelCase				:     str								=  parser.parse_args()
    convert_tf_checkpoint_to_pytorch(args.tf_checkpoint_path, args.pytorch_dump_path)
