


from __future__ import annotations

import unittest

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

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 (
	    TFFunnelBaseModel,
	    TFFunnelForMaskedLM,
	    TFFunnelForMultipleChoice,
	    TFFunnelForPreTraining,
	    TFFunnelForQuestionAnswering,
	    TFFunnelForSequenceClassification,
	    TFFunnelForTokenClassification,
	    TFFunnelModel,
	)
class 	lowerCamelCase_       :



					'''simple docstring'''







					def __init__(				self ,       snake_case_ ,       snake_case_=1_3 ,       snake_case_=7 ,       snake_case_=True ,       snake_case_=True ,       snake_case_=True ,       snake_case_=True ,       snake_case_=9_9 ,       snake_case_=[1, 1, 2] ,       snake_case_=1 ,       snake_case_=3_2 ,       snake_case_=4 ,       snake_case_=8 ,       snake_case_=3_7 ,       snake_case_="gelu_new" ,       snake_case_=0.1 ,       snake_case_=0.1 ,       snake_case_=0.0 ,       snake_case_=5_1_2 ,       snake_case_=3 ,       snake_case_=0.0_2 ,       snake_case_=3 ,       snake_case_=4 ,       snake_case_=None ,       snake_case_=False ,       )   ->			Union[str, Any]:







								'''simple docstring'''

								__lowercase 					=   parent
								__lowercase 					=   batch_size
								__lowercase 					=   seq_length
								__lowercase 					=   is_training
								__lowercase 					=   use_input_mask
								__lowercase 					=   use_token_type_ids
								__lowercase 					=   use_labels
								__lowercase 					=   vocab_size
								__lowercase 					=   block_sizes
								__lowercase 					=   num_decoder_layers
								__lowercase 					=   d_model
								__lowercase 					=   n_head
								__lowercase 					=   d_head
								__lowercase 					=   d_inner
								__lowercase 					=   hidden_act
								__lowercase 					=   hidden_dropout
								__lowercase 					=   attention_dropout
								__lowercase 					=   activation_dropout
								__lowercase 					=   max_position_embeddings
								__lowercase 					=   type_vocab_size
								__lowercase 					=   2
								__lowercase 					=   num_labels
								__lowercase 					=   num_choices
								__lowercase 					=   scope
								__lowercase 					=   initializer_std

								# Used in the tests to check the size of the first attention layer
								__lowercase 					=   n_head
								# Used in the tests to check the size of the first hidden state
								__lowercase 					=   self.d_model
								# Used in the tests to check the number of output hidden states/attentions
								__lowercase 					=   sum(self.block_sizes						) + (0 if base else self.num_decoder_layers)
								# FunnelModel adds two hidden layers: input embeddings and the sum of the upsampled encoder hidden state with
								# the last hidden state of the first block (which is the first hidden state of the decoder).
								if not base:
											__lowercase 					=   self.num_hidden_layers + 2






					def A      (				self						)   ->			Optional[Any]:







								'''simple docstring'''

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

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

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

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

								__lowercase 					=   FunnelConfig(
								    vocab_size=self.vocab_size ,       block_sizes=self.block_sizes ,       num_decoder_layers=self.num_decoder_layers ,       d_model=self.d_model ,       n_head=self.n_head ,       d_head=self.d_head ,       d_inner=self.d_inner ,       hidden_act=self.hidden_act ,       hidden_dropout=self.hidden_dropout ,       attention_dropout=self.attention_dropout ,       activation_dropout=self.activation_dropout ,       max_position_embeddings=self.max_position_embeddings ,       type_vocab_size=self.type_vocab_size ,       initializer_std=self.initializer_std ,       )

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






					def A      (				self ,       snake_case_ ,       snake_case_ ,       snake_case_ ,       snake_case_ ,       snake_case_ ,       snake_case_ ,       snake_case_ ,       )   ->			int:







								'''simple docstring'''

								__lowercase 					=   TFFunnelModel(config=snake_case_						)
								__lowercase 					=   {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
								__lowercase 					=   model(snake_case_						)

								__lowercase 					=   [input_ids, input_mask]
								__lowercase 					=   model(snake_case_						)

								__lowercase 					=   model(snake_case_						)
								self.parent.assertEqual(result.last_hidden_state.shape ,       (self.batch_size, self.seq_length, self.d_model)						)

								__lowercase 					=   False
								__lowercase 					=   TFFunnelModel(config=snake_case_						)
								__lowercase 					=   model(snake_case_						)
								self.parent.assertEqual(result.last_hidden_state.shape ,       (self.batch_size, self.seq_length, self.d_model)						)

								__lowercase 					=   False
								__lowercase 					=   TFFunnelModel(config=snake_case_						)
								__lowercase 					=   model(snake_case_						)
								self.parent.assertEqual(result.last_hidden_state.shape ,       (self.batch_size, self.seq_length, self.d_model)						)






					def A      (				self ,       snake_case_ ,       snake_case_ ,       snake_case_ ,       snake_case_ ,       snake_case_ ,       snake_case_ ,       snake_case_ ,       )   ->			Union[str, Any]:







								'''simple docstring'''

								__lowercase 					=   TFFunnelBaseModel(config=snake_case_						)
								__lowercase 					=   {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
								__lowercase 					=   model(snake_case_						)

								__lowercase 					=   [input_ids, input_mask]
								__lowercase 					=   model(snake_case_						)

								__lowercase 					=   model(snake_case_						)
								self.parent.assertEqual(result.last_hidden_state.shape ,       (self.batch_size, 2, self.d_model)						)

								__lowercase 					=   False
								__lowercase 					=   TFFunnelBaseModel(config=snake_case_						)
								__lowercase 					=   model(snake_case_						)
								self.parent.assertEqual(result.last_hidden_state.shape ,       (self.batch_size, 3, self.d_model)						)

								__lowercase 					=   False
								__lowercase 					=   TFFunnelBaseModel(config=snake_case_						)
								__lowercase 					=   model(snake_case_						)
								self.parent.assertEqual(result.last_hidden_state.shape ,       (self.batch_size, 2, self.d_model)						)






					def A      (				self ,       snake_case_ ,       snake_case_ ,       snake_case_ ,       snake_case_ ,       snake_case_ ,       snake_case_ ,       snake_case_ ,       )   ->			int:







								'''simple docstring'''

								__lowercase 					=   TFFunnelForPreTraining(config=snake_case_						)
								__lowercase 					=   {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
								__lowercase 					=   model(snake_case_						)
								self.parent.assertEqual(result.logits.shape ,       (self.batch_size, self.seq_length)						)






					def A      (				self ,       snake_case_ ,       snake_case_ ,       snake_case_ ,       snake_case_ ,       snake_case_ ,       snake_case_ ,       snake_case_ ,       )   ->			Dict:







								'''simple docstring'''

								__lowercase 					=   TFFunnelForMaskedLM(config=snake_case_						)
								__lowercase 					=   {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
								__lowercase 					=   model(snake_case_						)
								self.parent.assertEqual(result.logits.shape ,       (self.batch_size, self.seq_length, self.vocab_size)						)






					def A      (				self ,       snake_case_ ,       snake_case_ ,       snake_case_ ,       snake_case_ ,       snake_case_ ,       snake_case_ ,       snake_case_ ,       )   ->			Any:







								'''simple docstring'''

								__lowercase 					=   self.num_labels
								__lowercase 					=   TFFunnelForSequenceClassification(config=snake_case_						)
								__lowercase 					=   {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
								__lowercase 					=   model(snake_case_						)
								self.parent.assertEqual(result.logits.shape ,       (self.batch_size, self.num_labels)						)






					def A      (				self ,       snake_case_ ,       snake_case_ ,       snake_case_ ,       snake_case_ ,       snake_case_ ,       snake_case_ ,       snake_case_ ,       )   ->			Any:







								'''simple docstring'''

								__lowercase 					=   self.num_choices
								__lowercase 					=   TFFunnelForMultipleChoice(config=snake_case_						)
								__lowercase 					=   tf.tile(tf.expand_dims(snake_case_ ,       1						) ,       (1, self.num_choices, 1)						)
								__lowercase 					=   tf.tile(tf.expand_dims(snake_case_ ,       1						) ,       (1, self.num_choices, 1)						)
								__lowercase 					=   tf.tile(tf.expand_dims(snake_case_ ,       1						) ,       (1, self.num_choices, 1)						)
								__lowercase 					=   {
								    '''input_ids''': multiple_choice_inputs_ids,
								    '''attention_mask''': multiple_choice_input_mask,
								    '''token_type_ids''': multiple_choice_token_type_ids,
								}
								__lowercase 					=   model(snake_case_						)
								self.parent.assertEqual(result.logits.shape ,       (self.batch_size, self.num_choices)						)






					def A      (				self ,       snake_case_ ,       snake_case_ ,       snake_case_ ,       snake_case_ ,       snake_case_ ,       snake_case_ ,       snake_case_ ,       )   ->			Union[str, Any]:







								'''simple docstring'''

								__lowercase 					=   self.num_labels
								__lowercase 					=   TFFunnelForTokenClassification(config=snake_case_						)
								__lowercase 					=   {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
								__lowercase 					=   model(snake_case_						)
								self.parent.assertEqual(result.logits.shape ,       (self.batch_size, self.seq_length, self.num_labels)						)






					def A      (				self ,       snake_case_ ,       snake_case_ ,       snake_case_ ,       snake_case_ ,       snake_case_ ,       snake_case_ ,       snake_case_ ,       )   ->			Any:







								'''simple docstring'''

								__lowercase 					=   TFFunnelForQuestionAnswering(config=snake_case_						)
								__lowercase 					=   {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
								__lowercase 					=   model(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 A      (				self						)   ->			List[str]:







								'''simple docstring'''

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




@require_tf
class 	lowerCamelCase_       (				lowerCAmelCase__     ,     lowerCAmelCase__     ,     unittest.TestCase				):



					'''simple docstring'''







					__UpperCAmelCase            =				(
					    (
					        TFFunnelModel,
					        TFFunnelForMaskedLM,
					        TFFunnelForPreTraining,
					        TFFunnelForQuestionAnswering,
					        TFFunnelForTokenClassification,
					    )
					    if is_tf_available()
					    else ()
					)
					__UpperCAmelCase            =				(
					    {
					        "feature-extraction": (TFFunnelBaseModel, TFFunnelModel),
					        "fill-mask": TFFunnelForMaskedLM,
					        "question-answering": TFFunnelForQuestionAnswering,
					        "text-classification": TFFunnelForSequenceClassification,
					        "token-classification": TFFunnelForTokenClassification,
					        "zero-shot": TFFunnelForSequenceClassification,
					    }
					    if is_tf_available()
					    else {}
					)
					__UpperCAmelCase            =				False
					__UpperCAmelCase            =				False






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







								'''simple docstring'''

								__lowercase 					=   TFFunnelModelTester(self						)
								__lowercase 					=   ConfigTester(self ,       config_class=snake_case_						)






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







								'''simple docstring'''

								self.config_tester.run_common_tests()






					def A      (				self						)   ->			List[str]:







								'''simple docstring'''

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






					def A      (				self						)   ->			Optional[Any]:







								'''simple docstring'''

								__lowercase 					=   self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_pretraining(*snake_case_						)






					def A      (				self						)   ->			List[str]:







								'''simple docstring'''

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






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







								'''simple docstring'''

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






					def A      (				self						)   ->			Tuple:







								'''simple docstring'''

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




@require_tf
class 	lowerCamelCase_       (				lowerCAmelCase__     ,     unittest.TestCase				):



					'''simple docstring'''







					__UpperCAmelCase            =				(
					    (TFFunnelBaseModel, TFFunnelForMultipleChoice, TFFunnelForSequenceClassification) if is_tf_available() else ()
					)
					__UpperCAmelCase            =				False
					__UpperCAmelCase            =				False






					def A      (				self						)   ->			List[Any]:







								'''simple docstring'''

								__lowercase 					=   TFFunnelModelTester(self ,       base=snake_case_						)
								__lowercase 					=   ConfigTester(self ,       config_class=snake_case_						)






					def A      (				self						)   ->			List[Any]:







								'''simple docstring'''

								self.config_tester.run_common_tests()






					def A      (				self						)   ->			List[Any]:







								'''simple docstring'''

								__lowercase 					=   self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_base_model(*snake_case_						)






					def A      (				self						)   ->			Tuple:







								'''simple docstring'''

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






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







								'''simple docstring'''

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






import numpy as np
import torch
from torch.utils.data import DataLoader

from accelerate.utils.dataclasses import DistributedType
class 	lowerCamelCase_       :



					'''simple docstring'''







					def __init__(				self ,       snake_case_=2 ,       snake_case_=3 ,       snake_case_=6_4 ,       snake_case_=None						)   ->			List[str]:







								'''simple docstring'''

								__lowercase 					=   np.random.default_rng(snake_case_						)
								__lowercase 					=   length
								__lowercase 					=   rng.normal(size=(length,)						).astype(np.floataa						)
								__lowercase 					=   a * self.x + b + rng.normal(scale=0.1 ,       size=(length,)						).astype(np.floataa						)






					def __len__(				self						)   ->			Dict:







								'''simple docstring'''

								return self.length






					def __getitem__(				self ,       snake_case_						)   ->			Union[str, Any]:







								'''simple docstring'''

								return {"x": self.x[i], "y": self.y[i]}




class 	lowerCamelCase_       (				torch.nn.Module				):



					'''simple docstring'''







					def __init__(				self ,       snake_case_=0 ,       snake_case_=0 ,       snake_case_=False						)   ->			Optional[Any]:







								'''simple docstring'''

								super().__init__()
								__lowercase 					=   torch.nn.Parameter(torch.tensor([2, 3]						).float()						)
								__lowercase 					=   torch.nn.Parameter(torch.tensor([2, 3]						).float()						)
								__lowercase 					=   True






					def A      (				self ,       snake_case_=None						)   ->			List[Any]:







								'''simple docstring'''

								if self.first_batch:
											print(F'Model dtype: {self.a.dtype}, {self.b.dtype}. Input dtype: {x.dtype}'						)
											__lowercase 					=   False
								return x * self.a[0] + self.b[0]




class 	lowerCamelCase_       (				torch.nn.Module				):



					'''simple docstring'''







					def __init__(				self ,       snake_case_=0 ,       snake_case_=0 ,       snake_case_=False						)   ->			List[str]:







								'''simple docstring'''

								super().__init__()
								__lowercase 					=   torch.nn.Parameter(torch.tensor(snake_case_						).float()						)
								__lowercase 					=   torch.nn.Parameter(torch.tensor(snake_case_						).float()						)
								__lowercase 					=   True






					def A      (				self ,       snake_case_=None						)   ->			str:







								'''simple docstring'''

								if self.first_batch:
											print(F'Model dtype: {self.a.dtype}, {self.b.dtype}. Input dtype: {x.dtype}'						)
											__lowercase 					=   False
								return x * self.a + self.b




def 		lowercase_  (     _UpperCamelCase       ,			_UpperCamelCase = 16  ):


			'''simple docstring'''







			from datasets import load_dataset
			from transformers import AutoTokenizer

			__lowercase 					=   AutoTokenizer.from_pretrained('''bert-base-cased'''  )
			__lowercase 					=   {'''train''': '''tests/test_samples/MRPC/train.csv''', '''validation''': '''tests/test_samples/MRPC/dev.csv'''}
			__lowercase 					=   load_dataset('''csv'''       ,			data_files=_UpperCamelCase  )
			__lowercase 					=   datasets['''train'''].unique('''label'''  )

			__lowercase 					=   {v: i for i, v in enumerate(_UpperCamelCase  )}

			def tokenize_function(_UpperCamelCase  ):
						# max_length=None => use the model max length (it's actually the default)
						__lowercase 					=   tokenizer(
						    examples['''sentence1''']       ,			examples['''sentence2''']       ,			truncation=_UpperCamelCase       ,			max_length=_UpperCamelCase       ,			padding='''max_length'''  )
						if "label" in examples:
									__lowercase 					=   [label_to_id[l] for l in examples['''label''']]
						return outputs

			# Apply the method we just defined to all the examples in all the splits of the dataset
			__lowercase 					=   datasets.map(
			    _UpperCamelCase       ,			batched=_UpperCamelCase       ,			remove_columns=['''sentence1''', '''sentence2''', '''label''']       ,			)

			def collate_fn(_UpperCamelCase  ):
						# On TPU it's best to pad everything to the same length or training will be very slow.
						if accelerator.distributed_type == DistributedType.TPU:
									return tokenizer.pad(_UpperCamelCase       ,			padding='''max_length'''       ,			max_length=1_28       ,			return_tensors='''pt'''  )
						return tokenizer.pad(_UpperCamelCase       ,			padding='''longest'''       ,			return_tensors='''pt'''  )

			# Instantiate dataloaders.
			__lowercase 					=   DataLoader(tokenized_datasets['''train''']       ,			shuffle=_UpperCamelCase       ,			collate_fn=_UpperCamelCase       ,			batch_size=2  )
			__lowercase 					=   DataLoader(tokenized_datasets['''validation''']       ,			shuffle=_UpperCamelCase       ,			collate_fn=_UpperCamelCase       ,			batch_size=1  )

			return train_dataloader, eval_dataloader


