






import argparse
import os

# New Code #
import evaluate
import torch
from datasets import load_dataset
from torch.optim import AdamW
from torch.utils.data import DataLoader
from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed

from accelerate import Accelerator, DistributedType
from accelerate.utils import find_executable_batch_size


########################################################################
# This is a fully working simple example to use Accelerate,
# specifically showcasing how to ensure out-of-memory errors never
# interrupt training, and builds off the `nlp_example.py` script.
#
# This example trains a Bert base model on GLUE MRPC
# in any of the following settings (with the same script):
#   - single CPU or single GPU
#   - multi GPUS (using PyTorch distributed mode)
#   - (multi) TPUs
#   - fp16 (mixed-precision) or fp32 (normal precision)
#
# New additions from the base script can be found quickly by
# looking for the # New Code # tags
#
# To run it in each of these various modes, follow the instructions
# in the readme for examples:
# https://github.com/huggingface/accelerate/tree/main/examples
#
########################################################################


_UpperCAmelCase			:		Optional[int]   				=    16
_UpperCAmelCase			:		Union[str, Any]   				=    32





def 	lowerCAmelCase_		(lowercase__						:       Accelerator     ,					lowercase__						:       int = 16       )		->    Any:


				'''simple docstring'''




				lowerCAmelCase__      			=						AutoTokenizer.from_pretrained('''bert-base-cased'''       )
				lowerCAmelCase__      			=						load_dataset('''glue'''     ,					'''mrpc'''       )

				def tokenize_function(lowercase__						:       str       ):
								# max_length=None => use the model max length (it's actually the default)
								lowerCAmelCase__      			=						tokenizer(examples['''sentence1''']     ,					examples['''sentence2''']     ,					truncation=lowercase__     ,					max_length=lowercase__       )
								return outputs

				# Apply the method we just defined to all the examples in all the splits of the dataset
				# starting with the main process first:
				with accelerator.main_process_first():
								lowerCAmelCase__      			=						datasets.map(
								    lowercase__     ,					batched=lowercase__     ,					remove_columns=['''idx''', '''sentence1''', '''sentence2''']     ,					)

				# We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
				# transformers library
				lowerCAmelCase__      			=						tokenized_datasets.rename_column('''label'''     ,					'''labels'''       )

				def collate_fn(lowercase__						:       Any       ):
								# On TPU it's best to pad everything to the same length or training will be very slow.
								lowerCAmelCase__      			=						1_28 if accelerator.distributed_type == DistributedType.TPU else None
								# When using mixed precision we want round multiples of 8/16
								if accelerator.mixed_precision == "fp8":
												lowerCAmelCase__      			=						16
								elif accelerator.mixed_precision != "no":
												lowerCAmelCase__      			=						8
								else:
												lowerCAmelCase__      			=						None

								return tokenizer.pad(
								    lowercase__     ,					padding='''longest'''     ,					max_length=lowercase__     ,					pad_to_multiple_of=lowercase__     ,					return_tensors='''pt'''     ,					)

				# Instantiate dataloaders.
				lowerCAmelCase__      			=						DataLoader(
				    tokenized_datasets['''train''']     ,					shuffle=lowercase__     ,					collate_fn=lowercase__     ,					batch_size=lowercase__       )
				lowerCAmelCase__      			=						DataLoader(
				    tokenized_datasets['''validation''']     ,					shuffle=lowercase__     ,					collate_fn=lowercase__     ,					batch_size=lowercase__       )

				return train_dataloader, eval_dataloader


# For testing only
if os.environ.get("TESTING_MOCKED_DATALOADERS", None) == "1":
	from accelerate.test_utils.training import mocked_dataloaders

	_UpperCAmelCase			:		List[Any]   				=    mocked_dataloaders  # noqa: F811





def 	lowerCAmelCase_		(lowercase__						:       str     ,					lowercase__						:       Any       )		->    Dict:


				'''simple docstring'''




				if os.environ.get('''TESTING_MOCKED_DATALOADERS'''     ,					lowercase__       ) == "1":
								lowerCAmelCase__      			=						2
				# Initialize accelerator
				lowerCAmelCase__      			=						Accelerator(cpu=args.cpu     ,					mixed_precision=args.mixed_precision       )
				# Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
				lowerCAmelCase__      			=						config['''lr''']
				lowerCAmelCase__      			=						int(config['''num_epochs''']       )
				lowerCAmelCase__      			=						int(config['''seed''']       )
				lowerCAmelCase__      			=						int(config['''batch_size''']       )

				lowerCAmelCase__      			=						evaluate.load('''glue'''     ,					'''mrpc'''       )

				# New Code #
				# We now can define an inner training loop function. It should take a batch size as the only parameter,
				# and build the dataloaders in there.
				# It also gets our decorator
				@find_executable_batch_size(starting_batch_size=lowercase__       )
				def inner_training_loop(lowercase__						:       Tuple       ):
								# And now just move everything below under this function
								# We need to bring in the Accelerator object from earlier
								nonlocal accelerator
								# And reset all of its attributes that could hold onto any memory:
								accelerator.free_memory()

								# Then we can declare the model, optimizer, and everything else:
								set_seed(lowercase__       )

								# Instantiate the model (we build the model here so that the seed also control new weights initialization)
								lowerCAmelCase__      			=						AutoModelForSequenceClassification.from_pretrained('''bert-base-cased'''     ,					return_dict=lowercase__       )

								# We could avoid this line since the accelerator is set with `device_placement=True` (default value).
								# Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer
								# creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).
								lowerCAmelCase__      			=						model.to(accelerator.device       )

								# Instantiate optimizer
								lowerCAmelCase__      			=						AdamW(params=model.parameters()     ,					lr=lowercase__       )
								lowerCAmelCase__			,		lowerCAmelCase__      			=						get_dataloaders(lowercase__     ,					lowercase__       )

								# Instantiate scheduler
								lowerCAmelCase__      			=						get_linear_schedule_with_warmup(
								    optimizer=lowercase__     ,					num_warmup_steps=1_00     ,					num_training_steps=(len(lowercase__       ) * num_epochs)     ,					)

								# Prepare everything
								# There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
								# prepare method.
								lowerCAmelCase__			,		lowerCAmelCase__			,		lowerCAmelCase__			,		lowerCAmelCase__			,		lowerCAmelCase__      			=						accelerator.prepare(
								    lowercase__     ,					lowercase__     ,					lowercase__     ,					lowercase__     ,					lowercase__       )

								# Now we train the model
								for epoch in range(lowercase__       ):
												model.train()
												for step, batch in enumerate(lowercase__       ):
																# We could avoid this line since we set the accelerator with `device_placement=True`.
																batch.to(accelerator.device       )
																lowerCAmelCase__      			=						model(**lowercase__       )
																lowerCAmelCase__      			=						outputs.loss
																accelerator.backward(lowercase__       )
																optimizer.step()
																lr_scheduler.step()
																optimizer.zero_grad()

												model.eval()
												for step, batch in enumerate(lowercase__       ):
																# We could avoid this line since we set the accelerator with `device_placement=True`.
																batch.to(accelerator.device       )
																with torch.no_grad():
																				lowerCAmelCase__      			=						model(**lowercase__       )
																lowerCAmelCase__      			=						outputs.logits.argmax(dim=-1       )
																lowerCAmelCase__			,		lowerCAmelCase__      			=						accelerator.gather_for_metrics((predictions, batch['''labels'''])       )
																metric.add_batch(
																    predictions=lowercase__     ,					references=lowercase__     ,					)

												lowerCAmelCase__      			=						metric.compute()
												# Use accelerator.print to print only on the main process.
												accelerator.print(f'epoch {epoch}:'     ,					lowercase__       )

    # New Code #
    # And call it at the end with no arguments
    # Note: You could also refactor this outside of your training loop function
				inner_training_loop()





def 	lowerCAmelCase_		()		->    Dict:


				'''simple docstring'''




				lowerCAmelCase__      			=						argparse.ArgumentParser(description='''Simple example of training script.'''       )
				parser.add_argument(
				    '''--mixed_precision'''     ,					type=lowercase__     ,					default=lowercase__     ,					choices=['''no''', '''fp16''', '''bf16''', '''fp8''']     ,					help='''Whether to use mixed precision. Choose'''
				    '''between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10.'''
				    '''and an Nvidia Ampere GPU.'''     ,					)
				parser.add_argument('''--cpu'''     ,					action='''store_true'''     ,					help='''If passed, will train on the CPU.'''       )
				lowerCAmelCase__      			=						parser.parse_args()
				lowerCAmelCase__      			=						{'''lr''': 2e-5, '''num_epochs''': 3, '''seed''': 42, '''batch_size''': 16}
				training_function(lowercase__     ,					lowercase__       )


if __name__ == "__main__":
	main()













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




							def __init__(   self     :   Dict ,							SCREAMING_SNAKE_CASE_     :   str ,							SCREAMING_SNAKE_CASE_     :   Union[str, Any]=13 ,							SCREAMING_SNAKE_CASE_     :   str=7 ,							SCREAMING_SNAKE_CASE_     :   Tuple=True ,							SCREAMING_SNAKE_CASE_     :   Union[str, Any]=True ,							SCREAMING_SNAKE_CASE_     :   List[Any]=True ,							SCREAMING_SNAKE_CASE_     :   Optional[int]=True ,							SCREAMING_SNAKE_CASE_     :   Optional[int]=99 ,							SCREAMING_SNAKE_CASE_     :   Optional[Any]=[1, 1, 2] ,							SCREAMING_SNAKE_CASE_     :   List[str]=1 ,							SCREAMING_SNAKE_CASE_     :   str=32 ,							SCREAMING_SNAKE_CASE_     :   Dict=4 ,							SCREAMING_SNAKE_CASE_     :   Optional[Any]=8 ,							SCREAMING_SNAKE_CASE_     :   Dict=37 ,							SCREAMING_SNAKE_CASE_     :   Tuple="gelu_new" ,							SCREAMING_SNAKE_CASE_     :   Dict=0.1 ,							SCREAMING_SNAKE_CASE_     :   Any=0.1 ,							SCREAMING_SNAKE_CASE_     :   Any=0.0 ,							SCREAMING_SNAKE_CASE_     :   Tuple=512 ,							SCREAMING_SNAKE_CASE_     :   str=3 ,							SCREAMING_SNAKE_CASE_     :   int=0.02 ,							SCREAMING_SNAKE_CASE_     :   Optional[int]=3 ,							SCREAMING_SNAKE_CASE_     :   Any=4 ,							SCREAMING_SNAKE_CASE_     :   Optional[int]=None ,							SCREAMING_SNAKE_CASE_     :   Optional[Any]=False ,							):
											lowerCAmelCase__      			=						parent
											lowerCAmelCase__      			=						batch_size
											lowerCAmelCase__      			=						seq_length
											lowerCAmelCase__      			=						is_training
											lowerCAmelCase__      			=						use_input_mask
											lowerCAmelCase__      			=						use_token_type_ids
											lowerCAmelCase__      			=						use_labels
											lowerCAmelCase__      			=						vocab_size
											lowerCAmelCase__      			=						block_sizes
											lowerCAmelCase__      			=						num_decoder_layers
											lowerCAmelCase__      			=						d_model
											lowerCAmelCase__      			=						n_head
											lowerCAmelCase__      			=						d_head
											lowerCAmelCase__      			=						d_inner
											lowerCAmelCase__      			=						hidden_act
											lowerCAmelCase__      			=						hidden_dropout
											lowerCAmelCase__      			=						attention_dropout
											lowerCAmelCase__      			=						activation_dropout
											lowerCAmelCase__      			=						max_position_embeddings
											lowerCAmelCase__      			=						type_vocab_size
											lowerCAmelCase__      			=						2
											lowerCAmelCase__      			=						num_labels
											lowerCAmelCase__      			=						num_choices
											lowerCAmelCase__      			=						scope
											lowerCAmelCase__      			=						initializer_std

											# Used in the tests to check the size of the first attention layer
											lowerCAmelCase__      			=						n_head
											# Used in the tests to check the size of the first hidden state
											lowerCAmelCase__      			=						self.d_model
											# Used in the tests to check the number of output hidden states/attentions
											lowerCAmelCase__      			=						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:
															lowerCAmelCase__      			=						self.num_hidden_layers + 2




							def 						__snake_case					(   self     :   Optional[int]	):
											lowerCAmelCase__      			=						ids_tensor([self.batch_size, self.seq_length] ,							self.vocab_size	)

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

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

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

											lowerCAmelCase__      			=						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 						__snake_case					(   self     :   List[str] ,							SCREAMING_SNAKE_CASE_     :   str ,							SCREAMING_SNAKE_CASE_     :   int ,							SCREAMING_SNAKE_CASE_     :   int ,							SCREAMING_SNAKE_CASE_     :   int ,							SCREAMING_SNAKE_CASE_     :   int ,							SCREAMING_SNAKE_CASE_     :   List[Any] ,							SCREAMING_SNAKE_CASE_     :   Dict ,							):
											lowerCAmelCase__      			=						TFFunnelModel(config=SCREAMING_SNAKE_CASE_	)
											lowerCAmelCase__      			=						{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
											lowerCAmelCase__      			=						model(SCREAMING_SNAKE_CASE_	)

											lowerCAmelCase__      			=						[input_ids, input_mask]
											lowerCAmelCase__      			=						model(SCREAMING_SNAKE_CASE_	)

											lowerCAmelCase__      			=						model(SCREAMING_SNAKE_CASE_	)
											self.parent.assertEqual(result.last_hidden_state.shape ,							(self.batch_size, self.seq_length, self.d_model)	)

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

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




							def 						__snake_case					(   self     :   Any ,							SCREAMING_SNAKE_CASE_     :   str ,							SCREAMING_SNAKE_CASE_     :   Optional[int] ,							SCREAMING_SNAKE_CASE_     :   Tuple ,							SCREAMING_SNAKE_CASE_     :   List[Any] ,							SCREAMING_SNAKE_CASE_     :   Union[str, Any] ,							SCREAMING_SNAKE_CASE_     :   Any ,							SCREAMING_SNAKE_CASE_     :   Optional[int] ,							):
											lowerCAmelCase__      			=						TFFunnelBaseModel(config=SCREAMING_SNAKE_CASE_	)
											lowerCAmelCase__      			=						{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
											lowerCAmelCase__      			=						model(SCREAMING_SNAKE_CASE_	)

											lowerCAmelCase__      			=						[input_ids, input_mask]
											lowerCAmelCase__      			=						model(SCREAMING_SNAKE_CASE_	)

											lowerCAmelCase__      			=						model(SCREAMING_SNAKE_CASE_	)
											self.parent.assertEqual(result.last_hidden_state.shape ,							(self.batch_size, 2, self.d_model)	)

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

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




							def 						__snake_case					(   self     :   Union[str, Any] ,							SCREAMING_SNAKE_CASE_     :   Optional[int] ,							SCREAMING_SNAKE_CASE_     :   List[str] ,							SCREAMING_SNAKE_CASE_     :   Union[str, Any] ,							SCREAMING_SNAKE_CASE_     :   str ,							SCREAMING_SNAKE_CASE_     :   Optional[int] ,							SCREAMING_SNAKE_CASE_     :   Any ,							SCREAMING_SNAKE_CASE_     :   Any ,							):
											lowerCAmelCase__      			=						TFFunnelForPreTraining(config=SCREAMING_SNAKE_CASE_	)
											lowerCAmelCase__      			=						{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
											lowerCAmelCase__      			=						model(SCREAMING_SNAKE_CASE_	)
											self.parent.assertEqual(result.logits.shape ,							(self.batch_size, self.seq_length)	)




							def 						__snake_case					(   self     :   Tuple ,							SCREAMING_SNAKE_CASE_     :   Tuple ,							SCREAMING_SNAKE_CASE_     :   Optional[Any] ,							SCREAMING_SNAKE_CASE_     :   Tuple ,							SCREAMING_SNAKE_CASE_     :   List[str] ,							SCREAMING_SNAKE_CASE_     :   Tuple ,							SCREAMING_SNAKE_CASE_     :   List[str] ,							SCREAMING_SNAKE_CASE_     :   int ,							):
											lowerCAmelCase__      			=						TFFunnelForMaskedLM(config=SCREAMING_SNAKE_CASE_	)
											lowerCAmelCase__      			=						{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
											lowerCAmelCase__      			=						model(SCREAMING_SNAKE_CASE_	)
											self.parent.assertEqual(result.logits.shape ,							(self.batch_size, self.seq_length, self.vocab_size)	)




							def 						__snake_case					(   self     :   List[str] ,							SCREAMING_SNAKE_CASE_     :   Tuple ,							SCREAMING_SNAKE_CASE_     :   int ,							SCREAMING_SNAKE_CASE_     :   Any ,							SCREAMING_SNAKE_CASE_     :   Optional[int] ,							SCREAMING_SNAKE_CASE_     :   List[Any] ,							SCREAMING_SNAKE_CASE_     :   Optional[int] ,							SCREAMING_SNAKE_CASE_     :   Any ,							):
											lowerCAmelCase__      			=						self.num_labels
											lowerCAmelCase__      			=						TFFunnelForSequenceClassification(config=SCREAMING_SNAKE_CASE_	)
											lowerCAmelCase__      			=						{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
											lowerCAmelCase__      			=						model(SCREAMING_SNAKE_CASE_	)
											self.parent.assertEqual(result.logits.shape ,							(self.batch_size, self.num_labels)	)




							def 						__snake_case					(   self     :   int ,							SCREAMING_SNAKE_CASE_     :   Any ,							SCREAMING_SNAKE_CASE_     :   Tuple ,							SCREAMING_SNAKE_CASE_     :   str ,							SCREAMING_SNAKE_CASE_     :   List[Any] ,							SCREAMING_SNAKE_CASE_     :   Any ,							SCREAMING_SNAKE_CASE_     :   Union[str, Any] ,							SCREAMING_SNAKE_CASE_     :   Tuple ,							):
											lowerCAmelCase__      			=						self.num_choices
											lowerCAmelCase__      			=						TFFunnelForMultipleChoice(config=SCREAMING_SNAKE_CASE_	)
											lowerCAmelCase__      			=						tf.tile(tf.expand_dims(SCREAMING_SNAKE_CASE_ ,							1	) ,							(1, self.num_choices, 1)	)
											lowerCAmelCase__      			=						tf.tile(tf.expand_dims(SCREAMING_SNAKE_CASE_ ,							1	) ,							(1, self.num_choices, 1)	)
											lowerCAmelCase__      			=						tf.tile(tf.expand_dims(SCREAMING_SNAKE_CASE_ ,							1	) ,							(1, self.num_choices, 1)	)
											lowerCAmelCase__      			=						{
											    '''input_ids''': multiple_choice_inputs_ids,
											    '''attention_mask''': multiple_choice_input_mask,
											    '''token_type_ids''': multiple_choice_token_type_ids,
											}
											lowerCAmelCase__      			=						model(SCREAMING_SNAKE_CASE_	)
											self.parent.assertEqual(result.logits.shape ,							(self.batch_size, self.num_choices)	)




							def 						__snake_case					(   self     :   Optional[Any] ,							SCREAMING_SNAKE_CASE_     :   List[Any] ,							SCREAMING_SNAKE_CASE_     :   Optional[Any] ,							SCREAMING_SNAKE_CASE_     :   List[str] ,							SCREAMING_SNAKE_CASE_     :   Tuple ,							SCREAMING_SNAKE_CASE_     :   Tuple ,							SCREAMING_SNAKE_CASE_     :   Tuple ,							SCREAMING_SNAKE_CASE_     :   Optional[int] ,							):
											lowerCAmelCase__      			=						self.num_labels
											lowerCAmelCase__      			=						TFFunnelForTokenClassification(config=SCREAMING_SNAKE_CASE_	)
											lowerCAmelCase__      			=						{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
											lowerCAmelCase__      			=						model(SCREAMING_SNAKE_CASE_	)
											self.parent.assertEqual(result.logits.shape ,							(self.batch_size, self.seq_length, self.num_labels)	)




							def 						__snake_case					(   self     :   str ,							SCREAMING_SNAKE_CASE_     :   Union[str, Any] ,							SCREAMING_SNAKE_CASE_     :   str ,							SCREAMING_SNAKE_CASE_     :   List[str] ,							SCREAMING_SNAKE_CASE_     :   Optional[Any] ,							SCREAMING_SNAKE_CASE_     :   Any ,							SCREAMING_SNAKE_CASE_     :   Optional[Any] ,							SCREAMING_SNAKE_CASE_     :   Any ,							):
											lowerCAmelCase__      			=						TFFunnelForQuestionAnswering(config=SCREAMING_SNAKE_CASE_	)
											lowerCAmelCase__      			=						{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
											lowerCAmelCase__      			=						model(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 						__snake_case					(   self     :   List[str]	):
											lowerCAmelCase__      			=						self.prepare_config_and_inputs()
											(
											    (
											    lowerCAmelCase__
											)			,		(
											    lowerCAmelCase__
											)			,		(
											    lowerCAmelCase__
											)			,		(
											    lowerCAmelCase__
											)			,		(
											    lowerCAmelCase__
											)			,		(
											    lowerCAmelCase__
											)			,		(
											    lowerCAmelCase__
											)			,		
											)      			=						config_and_inputs
											lowerCAmelCase__      			=						{'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
											return config, inputs_dict



@require_tf
class  lowerCAmelCase_       (				snake_case__     ,    snake_case__     ,    unittest.TestCase					):
							UpperCamelCase_		:Dict		      =    (
							    (
							        TFFunnelModel,
							        TFFunnelForMaskedLM,
							        TFFunnelForPreTraining,
							        TFFunnelForQuestionAnswering,
							        TFFunnelForTokenClassification,
							    )
							    if is_tf_available()
							    else ()
							)
							UpperCamelCase_		:Dict		      =    (
							    {
							        '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_		:Tuple		      =    False
							UpperCamelCase_		:Tuple		      =    False




							def 						__snake_case					(   self     :   str	):
											lowerCAmelCase__      			=						TFFunnelModelTester(self	)
											lowerCAmelCase__      			=						ConfigTester(self ,							config_class=SCREAMING_SNAKE_CASE_	)




							def 						__snake_case					(   self     :   Dict	):
											self.config_tester.run_common_tests()




							def 						__snake_case					(   self     :   Union[str, Any]	):
											lowerCAmelCase__      			=						self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE_	)




							def 						__snake_case					(   self     :   Union[str, Any]	):
											lowerCAmelCase__      			=						self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_pretraining(*SCREAMING_SNAKE_CASE_	)




							def 						__snake_case					(   self     :   int	):
											lowerCAmelCase__      			=						self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_masked_lm(*SCREAMING_SNAKE_CASE_	)




							def 						__snake_case					(   self     :   str	):
											lowerCAmelCase__      			=						self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_token_classification(*SCREAMING_SNAKE_CASE_	)




							def 						__snake_case					(   self     :   Dict	):
											lowerCAmelCase__      			=						self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_question_answering(*SCREAMING_SNAKE_CASE_	)



@require_tf
class  lowerCAmelCase_       (				snake_case__     ,    unittest.TestCase					):
							UpperCamelCase_		:Optional[Any]		      =    (
							    (TFFunnelBaseModel, TFFunnelForMultipleChoice, TFFunnelForSequenceClassification) if is_tf_available() else ()
							)
							UpperCamelCase_		:int		      =    False
							UpperCamelCase_		:Any		      =    False




							def 						__snake_case					(   self     :   List[str]	):
											lowerCAmelCase__      			=						TFFunnelModelTester(self ,							base=SCREAMING_SNAKE_CASE_	)
											lowerCAmelCase__      			=						ConfigTester(self ,							config_class=SCREAMING_SNAKE_CASE_	)




							def 						__snake_case					(   self     :   Tuple	):
											self.config_tester.run_common_tests()




							def 						__snake_case					(   self     :   Union[str, Any]	):
											lowerCAmelCase__      			=						self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_base_model(*SCREAMING_SNAKE_CASE_	)




							def 						__snake_case					(   self     :   int	):
											lowerCAmelCase__      			=						self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_sequence_classification(*SCREAMING_SNAKE_CASE_	)




							def 						__snake_case					(   self     :   List[str]	):
											lowerCAmelCase__      			=						self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_multiple_choice(*SCREAMING_SNAKE_CASE_	)





