





from __future__ import annotations

import unittest

from transformers import is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
				import tensorflow as tf

				from transformers import (
				    TF_XLM_PRETRAINED_MODEL_ARCHIVE_LIST,
				    TFXLMForMultipleChoice,
				    TFXLMForQuestionAnsweringSimple,
				    TFXLMForSequenceClassification,
				    TFXLMForTokenClassification,
				    TFXLMModel,
				    TFXLMWithLMHeadModel,
				    XLMConfig,
				)
class 							snake_case_		:







		def __init__(		self      ,  __lowercase      ,  )					-> Union[str, Any]:
				lowerCamelCase :			Optional[Any]    		=parent
				lowerCamelCase :			Tuple    		=1_3
				lowerCamelCase :			Union[str, Any]    		=7
				lowerCamelCase :			Tuple    		=True
				lowerCamelCase :			List[str]    		=True
				lowerCamelCase :			str    		=True
				lowerCamelCase :			Optional[Any]    		=True
				lowerCamelCase :			Any    		=True
				lowerCamelCase :			List[str]    		=False
				lowerCamelCase :			Any    		=False
				lowerCamelCase :			Any    		=False
				lowerCamelCase :			Optional[Any]    		=2
				lowerCamelCase :			Union[str, Any]    		=9_9
				lowerCamelCase :			int    		=0
				lowerCamelCase :			Tuple    		=3_2
				lowerCamelCase :			int    		=2
				lowerCamelCase :			int    		=4
				lowerCamelCase :			Any    		=0.1
				lowerCamelCase :			str    		=0.1
				lowerCamelCase :			str    		=5_1_2
				lowerCamelCase :			Optional[int]    		=1_6
				lowerCamelCase :			Tuple    		=2
				lowerCamelCase :			List[str]    		=0.0_2
				lowerCamelCase :			Dict    		=3
				lowerCamelCase :			str    		=4
				lowerCamelCase :			List[str]    		='''last'''
				lowerCamelCase :			Union[str, Any]    		=True
				lowerCamelCase :			Dict    		=None
				lowerCamelCase :			List[str]    		=0







		def       __lowercase (		self )					-> List[str]:
				lowerCamelCase :			List[Any]    		=ids_tensor([self.batch_size, self.seq_length]      ,  self.vocab_size )
				lowerCamelCase :			str    		=random_attention_mask([self.batch_size, self.seq_length]      ,  dtype=tf.floataa )

				lowerCamelCase :			Dict    		=None
				if self.use_input_lengths:
						lowerCamelCase :			Union[str, Any]    		=(
						    ids_tensor([self.batch_size]      ,  vocab_size=2 ) + self.seq_length - 2
						)  # small variation of seq_length

				lowerCamelCase :			Tuple    		=None
				if self.use_token_type_ids:
						lowerCamelCase :			List[str]    		=ids_tensor([self.batch_size, self.seq_length]      ,  self.n_langs )

				lowerCamelCase :			Union[str, Any]    		=None
				lowerCamelCase :			str    		=None
				lowerCamelCase :			Tuple    		=None
				if self.use_labels:
						lowerCamelCase :			List[str]    		=ids_tensor([self.batch_size]      ,  self.type_sequence_label_size )
						lowerCamelCase :			str    		=ids_tensor([self.batch_size, self.seq_length]      ,  self.num_labels )
						lowerCamelCase :			Optional[int]    		=ids_tensor([self.batch_size]      ,  2      ,  dtype=tf.floataa )
						lowerCamelCase :			Dict    		=ids_tensor([self.batch_size]      ,  self.num_choices )

				lowerCamelCase :			List[str]    		=XLMConfig(
				    vocab_size=self.vocab_size      ,  n_special=self.n_special      ,  emb_dim=self.hidden_size      ,  n_layers=self.num_hidden_layers      ,  n_heads=self.num_attention_heads      ,  dropout=self.hidden_dropout_prob      ,  attention_dropout=self.attention_probs_dropout_prob      ,  gelu_activation=self.gelu_activation      ,  sinusoidal_embeddings=self.sinusoidal_embeddings      ,  asm=self.asm      ,  causal=self.causal      ,  n_langs=self.n_langs      ,  max_position_embeddings=self.max_position_embeddings      ,  initializer_range=self.initializer_range      ,  summary_type=self.summary_type      ,  use_proj=self.use_proj      ,  bos_token_id=self.bos_token_id      ,  )

				return (
				    config,
				    input_ids,
				    token_type_ids,
				    input_lengths,
				    sequence_labels,
				    token_labels,
				    is_impossible_labels,
				    choice_labels,
				    input_mask,
				)







		def       __lowercase (		self      ,  __lowercase      ,  __lowercase      ,  __lowercase      ,  __lowercase      ,  __lowercase      ,  __lowercase      ,  __lowercase      ,  __lowercase      ,  __lowercase      ,  )					-> Union[str, Any]:
				lowerCamelCase :			Dict    		=TFXLMModel(config=__lowercase )
				lowerCamelCase :			Optional[Any]    		={'''input_ids''': input_ids, '''lengths''': input_lengths, '''langs''': token_type_ids}
				lowerCamelCase :			Tuple    		=model(__lowercase )

				lowerCamelCase :			Optional[int]    		=[input_ids, input_mask]
				lowerCamelCase :			Any    		=model(__lowercase )
				self.parent.assertEqual(result.last_hidden_state.shape      ,  (self.batch_size, self.seq_length, self.hidden_size) )







		def       __lowercase (		self      ,  __lowercase      ,  __lowercase      ,  __lowercase      ,  __lowercase      ,  __lowercase      ,  __lowercase      ,  __lowercase      ,  __lowercase      ,  __lowercase      ,  )					-> Any:
				lowerCamelCase :			Any    		=TFXLMWithLMHeadModel(__lowercase )

				lowerCamelCase :			Union[str, Any]    		={'''input_ids''': input_ids, '''lengths''': input_lengths, '''langs''': token_type_ids}
				lowerCamelCase :			str    		=model(__lowercase )

				lowerCamelCase :			Union[str, Any]    		=outputs

				self.parent.assertEqual(result.logits.shape      ,  (self.batch_size, self.seq_length, self.vocab_size) )







		def       __lowercase (		self      ,  __lowercase      ,  __lowercase      ,  __lowercase      ,  __lowercase      ,  __lowercase      ,  __lowercase      ,  __lowercase      ,  __lowercase      ,  __lowercase      ,  )					-> Any:
				lowerCamelCase :			Any    		=TFXLMForQuestionAnsweringSimple(__lowercase )

				lowerCamelCase :			Optional[Any]    		={'''input_ids''': input_ids, '''lengths''': input_lengths}

				lowerCamelCase :			Optional[int]    		=model(__lowercase )

				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       __lowercase (		self      ,  __lowercase      ,  __lowercase      ,  __lowercase      ,  __lowercase      ,  __lowercase      ,  __lowercase      ,  __lowercase      ,  __lowercase      ,  __lowercase      ,  )					-> Dict:
				lowerCamelCase :			List[Any]    		=TFXLMForSequenceClassification(__lowercase )

				lowerCamelCase :			List[str]    		={'''input_ids''': input_ids, '''lengths''': input_lengths}

				lowerCamelCase :			Union[str, Any]    		=model(__lowercase )

				self.parent.assertEqual(result.logits.shape      ,  (self.batch_size, self.type_sequence_label_size) )







		def       __lowercase (		self      ,  __lowercase      ,  __lowercase      ,  __lowercase      ,  __lowercase      ,  __lowercase      ,  __lowercase      ,  __lowercase      ,  __lowercase      ,  __lowercase      ,  )					-> Any:
				lowerCamelCase :			str    		=self.num_labels
				lowerCamelCase :			List[str]    		=TFXLMForTokenClassification(config=__lowercase )
				lowerCamelCase :			List[Any]    		={'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
				lowerCamelCase :			Optional[Any]    		=model(__lowercase )
				self.parent.assertEqual(result.logits.shape      ,  (self.batch_size, self.seq_length, self.num_labels) )







		def       __lowercase (		self      ,  __lowercase      ,  __lowercase      ,  __lowercase      ,  __lowercase      ,  __lowercase      ,  __lowercase      ,  __lowercase      ,  __lowercase      ,  __lowercase      ,  )					-> List[Any]:
				lowerCamelCase :			str    		=self.num_choices
				lowerCamelCase :			Optional[Any]    		=TFXLMForMultipleChoice(config=__lowercase )
				lowerCamelCase :			Optional[Any]    		=tf.tile(tf.expand_dims(__lowercase      ,  1 )      ,  (1, self.num_choices, 1) )
				lowerCamelCase :			Dict    		=tf.tile(tf.expand_dims(__lowercase      ,  1 )      ,  (1, self.num_choices, 1) )
				lowerCamelCase :			Union[str, Any]    		=tf.tile(tf.expand_dims(__lowercase      ,  1 )      ,  (1, self.num_choices, 1) )
				lowerCamelCase :			Tuple    		={
				    '''input_ids''': multiple_choice_inputs_ids,
				    '''attention_mask''': multiple_choice_input_mask,
				    '''token_type_ids''': multiple_choice_token_type_ids,
				}
				lowerCamelCase :			Optional[int]    		=model(__lowercase )
				self.parent.assertEqual(result.logits.shape      ,  (self.batch_size, self.num_choices) )







		def       __lowercase (		self )					-> List[Any]:
				lowerCamelCase :			str    		=self.prepare_config_and_inputs()
				(
				    (
				    lowerCamelCase
				)		, (
				    lowerCamelCase
				)		, (
				    lowerCamelCase
				)		, (
				    lowerCamelCase
				)		, (
				    lowerCamelCase
				)		, (
				    lowerCamelCase
				)		, (
				    lowerCamelCase
				)		, (
				    lowerCamelCase
				)		, (
				    lowerCamelCase
				)		, 
				) :			Dict    		=config_and_inputs
				lowerCamelCase :			Optional[int]    		={
				    '''input_ids''': input_ids,
				    '''token_type_ids''': token_type_ids,
				    '''langs''': token_type_ids,
				    '''lengths''': input_lengths,
				}
				return config, inputs_dict

@require_tf
class 							snake_case_		(      _A    ,       _A    ,       unittest.TestCase):
		lowerCamelCase     :Any							=  (
		    (
		        TFXLMModel,
		        TFXLMWithLMHeadModel,
		        TFXLMForSequenceClassification,
		        TFXLMForQuestionAnsweringSimple,
		        TFXLMForTokenClassification,
		        TFXLMForMultipleChoice,
		    )
		    if is_tf_available()
		    else ()
		)
		lowerCamelCase     :Optional[Any]							=  (
		    (TFXLMWithLMHeadModel,) if is_tf_available() else ()
		)  # TODO (PVP): Check other models whether language generation is also applicable
		lowerCamelCase     :List[str]							=  (
		    {
		        "feature-extraction": TFXLMModel,
		        "fill-mask": TFXLMWithLMHeadModel,
		        "question-answering": TFXLMForQuestionAnsweringSimple,
		        "text-classification": TFXLMForSequenceClassification,
		        "text-generation": TFXLMWithLMHeadModel,
		        "token-classification": TFXLMForTokenClassification,
		        "zero-shot": TFXLMForSequenceClassification,
		    }
		    if is_tf_available()
		    else {}
		)
		lowerCamelCase     :Union[str, Any]							=  False
		lowerCamelCase     :int							=  False







		def       __lowercase (		self      ,  __lowercase      ,  __lowercase      ,  __lowercase      ,  __lowercase      ,  __lowercase )					-> int:
				if (
				    pipeline_test_casse_name == "QAPipelineTests"
				    and tokenizer_name is not None
				    and not tokenizer_name.endswith('''Fast''' )
				):
						# `QAPipelineTests` fails for a few models when the slower tokenizer are used.
						# (The slower tokenizers were never used for pipeline tests before the pipeline testing rework)
						# TODO: check (and possibly fix) the `QAPipelineTests` with slower tokenizer
						return True

				return False







		def       __lowercase (		self )					-> Optional[int]:
				lowerCamelCase :			Dict    		=TFXLMModelTester(self )
				lowerCamelCase :			Any    		=ConfigTester(self      ,  config_class=__lowercase      ,  emb_dim=3_7 )







		def       __lowercase (		self )					-> str:
				self.config_tester.run_common_tests()







		def       __lowercase (		self )					-> str:
				lowerCamelCase :			str    		=self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_xlm_model(*__lowercase )







		def       __lowercase (		self )					-> Optional[Any]:
				lowerCamelCase :			List[Any]    		=self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_xlm_lm_head(*__lowercase )







		def       __lowercase (		self )					-> Union[str, Any]:
				lowerCamelCase :			Optional[Any]    		=self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_xlm_qa(*__lowercase )







		def       __lowercase (		self )					-> List[Any]:
				lowerCamelCase :			str    		=self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_xlm_sequence_classif(*__lowercase )







		def       __lowercase (		self )					-> int:
				lowerCamelCase :			Dict    		=self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_xlm_for_token_classification(*__lowercase )







		def       __lowercase (		self )					-> List[Any]:
				lowerCamelCase :			List[Any]    		=self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_xlm_for_multiple_choice(*__lowercase )







		@slow
		def       __lowercase (		self )					-> Optional[Any]:
				for model_name in TF_XLM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
						lowerCamelCase :			Optional[int]    		=TFXLMModel.from_pretrained(__lowercase )
						self.assertIsNotNone(__lowercase )

@require_tf
class 							snake_case_		(      unittest.TestCase):







		@slow
		def       __lowercase (		self )					-> Optional[Any]:
				lowerCamelCase :			Optional[Any]    		=TFXLMWithLMHeadModel.from_pretrained('''xlm-mlm-en-2048''' )
				lowerCamelCase :			List[Any]    		=tf.convert_to_tensor([[1_4, 4_4_7]]      ,  dtype=tf.intaa )  # the president
				lowerCamelCase :			Dict    		=[
				    1_4,
				    4_4_7,
				    1_4,
				    4_4_7,
				    1_4,
				    4_4_7,
				    1_4,
				    4_4_7,
				    1_4,
				    4_4_7,
				    1_4,
				    4_4_7,
				    1_4,
				    4_4_7,
				    1_4,
				    4_4_7,
				    1_4,
				    4_4_7,
				    1_4,
				    4_4_7,
				]  # the president the president the president the president the president the president the president the president the president the president
				# TODO(PVP): this and other input_ids I tried for generation give pretty bad results. Not sure why. Model might just not be made for auto-regressive inference
				lowerCamelCase :			List[str]    		=model.generate(__lowercase      ,  do_sample=__lowercase )
				self.assertListEqual(output_ids[0].numpy().tolist()      ,  __lowercase )








import argparse

import torch
from flax.training.checkpoints import restore_checkpoint

from transformers import FNetConfig, FNetForPreTraining
from transformers.utils import logging


logging.set_verbosity_info()


def        A__						(							SCREAMING_SNAKE_CASE_				,							SCREAMING_SNAKE_CASE_				,							SCREAMING_SNAKE_CASE_		) ->  List[Any]:
		# Initialise PyTorch model
		lowerCamelCase :			Dict    		=FNetConfig.from_json_file(SCREAMING_SNAKE_CASE_		)
		print(F"Building PyTorch model from configuration: {config}"		)
		lowerCamelCase :			Tuple    		=FNetForPreTraining(SCREAMING_SNAKE_CASE_		)

		lowerCamelCase :			List[Any]    		=restore_checkpoint(SCREAMING_SNAKE_CASE_				,							SCREAMING_SNAKE_CASE_		)
		lowerCamelCase :			List[str]    		=checkpoint_dict['''target''']

		# Embeddings
		# Position IDs
		lowerCamelCase :			Optional[Any]    		=fnet_pretraining_model.state_dict()

		lowerCamelCase :			Tuple    		=state_dict['''fnet.embeddings.position_ids''']
		lowerCamelCase :			str    		={'''fnet.embeddings.position_ids''': position_ids}
		# Embedding Layers
		lowerCamelCase :			int    		=torch.tensor(
		    pretrained_model_params['''encoder''']['''embedder''']['''word''']['''embedding''']		)
		lowerCamelCase :			int    		=torch.tensor(
		    pretrained_model_params['''encoder''']['''embedder''']['''position''']['''embedding'''][0]		)
		lowerCamelCase :			str    		=torch.tensor(
		    pretrained_model_params['''encoder''']['''embedder''']['''type''']['''embedding''']		)
		lowerCamelCase :			str    		=torch.tensor(
		    pretrained_model_params['''encoder''']['''embedder''']['''hidden_mapping_in''']['''kernel''']		).T
		lowerCamelCase :			List[Any]    		=torch.tensor(
		    pretrained_model_params['''encoder''']['''embedder''']['''hidden_mapping_in''']['''bias''']		)
		lowerCamelCase :			List[str]    		=torch.tensor(
		    pretrained_model_params['''encoder''']['''embedder''']['''layer_norm''']['''scale''']		)
		lowerCamelCase :			Tuple    		=torch.tensor(
		    pretrained_model_params['''encoder''']['''embedder''']['''layer_norm''']['''bias''']		)

		# Encoder Layers
		for layer in range(config.num_hidden_layers		):
				lowerCamelCase :			int    		=torch.tensor(
				    pretrained_model_params['''encoder'''][F"encoder_{layer}"]['''mixing_layer_norm''']['''scale''']		)
				lowerCamelCase :			int    		=torch.tensor(
				    pretrained_model_params['''encoder'''][F"encoder_{layer}"]['''mixing_layer_norm''']['''bias''']		)

				lowerCamelCase :			List[Any]    		=torch.tensor(
				    pretrained_model_params['''encoder'''][F"feed_forward_{layer}"]['''intermediate''']['''kernel''']		).T
				lowerCamelCase :			List[Any]    		=torch.tensor(
				    pretrained_model_params['''encoder'''][F"feed_forward_{layer}"]['''intermediate''']['''bias''']		)

				lowerCamelCase :			str    		=torch.tensor(
				    pretrained_model_params['''encoder'''][F"feed_forward_{layer}"]['''output''']['''kernel''']		).T
				lowerCamelCase :			Tuple    		=torch.tensor(
				    pretrained_model_params['''encoder'''][F"feed_forward_{layer}"]['''output''']['''bias''']		)

				lowerCamelCase :			List[Any]    		=torch.tensor(
				    pretrained_model_params['''encoder'''][F"encoder_{layer}"]['''output_layer_norm''']['''scale''']		)
				lowerCamelCase :			Optional[Any]    		=torch.tensor(
				    pretrained_model_params['''encoder'''][F"encoder_{layer}"]['''output_layer_norm''']['''bias''']		)

		# Pooler Layers
		lowerCamelCase :			Any    		=torch.tensor(pretrained_model_params['''encoder''']['''pooler''']['''kernel''']		).T
		lowerCamelCase :			str    		=torch.tensor(pretrained_model_params['''encoder''']['''pooler''']['''bias''']		)

		# Masked LM Layers
		lowerCamelCase :			Optional[int]    		=torch.tensor(
		    pretrained_model_params['''predictions_dense''']['''kernel''']		).T
		lowerCamelCase :			str    		=torch.tensor(
		    pretrained_model_params['''predictions_dense''']['''bias''']		)
		lowerCamelCase :			Optional[Any]    		=torch.tensor(
		    pretrained_model_params['''predictions_layer_norm''']['''scale''']		)
		lowerCamelCase :			Union[str, Any]    		=torch.tensor(
		    pretrained_model_params['''predictions_layer_norm''']['''bias''']		)
		lowerCamelCase :			Optional[int]    		=torch.tensor(
		    pretrained_model_params['''encoder''']['''embedder''']['''word''']['''embedding''']		)
		lowerCamelCase :			Union[str, Any]    		=torch.tensor(
		    pretrained_model_params['''predictions_output''']['''output_bias''']		)
		lowerCamelCase :			Tuple    		=torch.tensor(pretrained_model_params['''predictions_output''']['''output_bias''']		)

		# Seq Relationship Layers
		lowerCamelCase :			Union[str, Any]    		=torch.tensor(
		    pretrained_model_params['''classification''']['''output_kernel''']		)
		lowerCamelCase :			Any    		=torch.tensor(
		    pretrained_model_params['''classification''']['''output_bias''']		)

		# Load State Dict
		fnet_pretraining_model.load_state_dict(SCREAMING_SNAKE_CASE_		)

		# Save PreTrained
		print(F"Saving pretrained model to {save_path}"		)
		fnet_pretraining_model.save_pretrained(SCREAMING_SNAKE_CASE_		)


if __name__ == "__main__":
				snake_case_				=      argparse.ArgumentParser()
				# Required parameters
				parser.add_argument(
				    '''--flax_checkpoint_path''', default=None, type=str, required=True, help='''Path to the TensorFlow checkpoint path.'''
				)
				parser.add_argument(
				    '''--fnet_config_file''',
				    default=None,
				    type=str,
				    required=True,
				    help=(
				        '''The config json file corresponding to the pre-trained FNet model. \n'''
				        '''This specifies the model architecture.'''
				    ),
				)
				parser.add_argument('''--save_path''', default=None, type=str, required=True, help='''Path to the output model.''')
				snake_case_				=      parser.parse_args()
				convert_flax_checkpoint_to_pytorch(args.flax_checkpoint_path, args.fnet_config_file, args.save_path)

