




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



		def __init__(     self,	UpperCamelCase__,	UpperCamelCase__=13,	UpperCamelCase__=7,	UpperCamelCase__=True,	UpperCamelCase__=True,	UpperCamelCase__=True,	UpperCamelCase__=True,	UpperCamelCase__=99,	UpperCamelCase__=[1, 1, 2],	UpperCamelCase__=1,	UpperCamelCase__=32,	UpperCamelCase__=4,	UpperCamelCase__=8,	UpperCamelCase__=37,	UpperCamelCase__="gelu_new",	UpperCamelCase__=0.1,	UpperCamelCase__=0.1,	UpperCamelCase__=0.0,	UpperCamelCase__=512,	UpperCamelCase__=3,	UpperCamelCase__=0.02,	UpperCamelCase__=3,	UpperCamelCase__=4,	UpperCamelCase__=None,	UpperCamelCase__=False,	):







								"""simple docstring"""



								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        SCREAMING_SNAKE_CASE__			(     self   ):







								"""simple docstring"""



								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        SCREAMING_SNAKE_CASE__			(     self,	UpperCamelCase__,	UpperCamelCase__,	UpperCamelCase__,	UpperCamelCase__,	UpperCamelCase__,	UpperCamelCase__,	UpperCamelCase__,	):







								"""simple docstring"""



								lowerCAmelCase_      			=  TFFunnelModel(config=a_   )
								lowerCAmelCase_      			=  {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
								lowerCAmelCase_      			=  model(a_   )

								lowerCAmelCase_      			=  [input_ids, input_mask]
								lowerCAmelCase_      			=  model(a_   )

								lowerCAmelCase_      			=  model(a_   )
								self.parent.assertEqual(result.last_hidden_state.shape,	(self.batch_size, self.seq_length, self.d_model)   )

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

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



		def        SCREAMING_SNAKE_CASE__			(     self,	UpperCamelCase__,	UpperCamelCase__,	UpperCamelCase__,	UpperCamelCase__,	UpperCamelCase__,	UpperCamelCase__,	UpperCamelCase__,	):







								"""simple docstring"""



								lowerCAmelCase_      			=  TFFunnelBaseModel(config=a_   )
								lowerCAmelCase_      			=  {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
								lowerCAmelCase_      			=  model(a_   )

								lowerCAmelCase_      			=  [input_ids, input_mask]
								lowerCAmelCase_      			=  model(a_   )

								lowerCAmelCase_      			=  model(a_   )
								self.parent.assertEqual(result.last_hidden_state.shape,	(self.batch_size, 2, self.d_model)   )

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

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



		def        SCREAMING_SNAKE_CASE__			(     self,	UpperCamelCase__,	UpperCamelCase__,	UpperCamelCase__,	UpperCamelCase__,	UpperCamelCase__,	UpperCamelCase__,	UpperCamelCase__,	):







								"""simple docstring"""



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



		def        SCREAMING_SNAKE_CASE__			(     self,	UpperCamelCase__,	UpperCamelCase__,	UpperCamelCase__,	UpperCamelCase__,	UpperCamelCase__,	UpperCamelCase__,	UpperCamelCase__,	):







								"""simple docstring"""



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



		def        SCREAMING_SNAKE_CASE__			(     self,	UpperCamelCase__,	UpperCamelCase__,	UpperCamelCase__,	UpperCamelCase__,	UpperCamelCase__,	UpperCamelCase__,	UpperCamelCase__,	):







								"""simple docstring"""



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



		def        SCREAMING_SNAKE_CASE__			(     self,	UpperCamelCase__,	UpperCamelCase__,	UpperCamelCase__,	UpperCamelCase__,	UpperCamelCase__,	UpperCamelCase__,	UpperCamelCase__,	):







								"""simple docstring"""



								lowerCAmelCase_      			=  self.num_choices
								lowerCAmelCase_      			=  TFFunnelForMultipleChoice(config=a_   )
								lowerCAmelCase_      			=  tf.tile(tf.expand_dims(a_,	1   ),	(1, self.num_choices, 1)   )
								lowerCAmelCase_      			=  tf.tile(tf.expand_dims(a_,	1   ),	(1, self.num_choices, 1)   )
								lowerCAmelCase_      			=  tf.tile(tf.expand_dims(a_,	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(a_   )
								self.parent.assertEqual(result.logits.shape,	(self.batch_size, self.num_choices)   )



		def        SCREAMING_SNAKE_CASE__			(     self,	UpperCamelCase__,	UpperCamelCase__,	UpperCamelCase__,	UpperCamelCase__,	UpperCamelCase__,	UpperCamelCase__,	UpperCamelCase__,	):







								"""simple docstring"""



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



		def        SCREAMING_SNAKE_CASE__			(     self,	UpperCamelCase__,	UpperCamelCase__,	UpperCamelCase__,	UpperCamelCase__,	UpperCamelCase__,	UpperCamelCase__,	UpperCamelCase__,	):







								"""simple docstring"""



								lowerCAmelCase_      			=  TFFunnelForQuestionAnswering(config=a_   )
								lowerCAmelCase_      			=  {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
								lowerCAmelCase_      			=  model(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        SCREAMING_SNAKE_CASE__			(     self   ):







								"""simple docstring"""



								lowerCAmelCase_      			=  self.prepare_config_and_inputs()
								(
								    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  A		(  __UpperCamelCase      ,				__UpperCamelCase      ,				unittest.TestCase     ):
		__snake_case									=   (
		    (
		        TFFunnelModel,
		        TFFunnelForMaskedLM,
		        TFFunnelForPreTraining,
		        TFFunnelForQuestionAnswering,
		        TFFunnelForTokenClassification,
		    )
		    if is_tf_available()
		    else ()
		)
		__snake_case									=   (
		    {
		        'feature-extraction': (TFFunnelBaseModel, TFFunnelModel),
		        'fill-mask': TFFunnelForMaskedLM,
		        'question-answering': TFFunnelForQuestionAnswering,
		        'text-classification': TFFunnelForSequenceClassification,
		        'token-classification': TFFunnelForTokenClassification,
		        'zero-shot': TFFunnelForSequenceClassification,
		    }
		    if is_tf_available()
		    else {}
		)
		__snake_case									=   False
		__snake_case									=   False



		def        SCREAMING_SNAKE_CASE__			(     self   ):







								"""simple docstring"""



								lowerCAmelCase_      			=  TFFunnelModelTester(self   )
								lowerCAmelCase_      			=  ConfigTester(self,	config_class=a_   )



		def        SCREAMING_SNAKE_CASE__			(     self   ):







								"""simple docstring"""



								self.config_tester.run_common_tests()



		def        SCREAMING_SNAKE_CASE__			(     self   ):







								"""simple docstring"""



								lowerCAmelCase_      			=  self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_model(*a_   )



		def        SCREAMING_SNAKE_CASE__			(     self   ):







								"""simple docstring"""



								lowerCAmelCase_      			=  self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_pretraining(*a_   )



		def        SCREAMING_SNAKE_CASE__			(     self   ):







								"""simple docstring"""



								lowerCAmelCase_      			=  self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_masked_lm(*a_   )



		def        SCREAMING_SNAKE_CASE__			(     self   ):







								"""simple docstring"""



								lowerCAmelCase_      			=  self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_token_classification(*a_   )







		def        SCREAMING_SNAKE_CASE__			(     self   ):







								"""simple docstring"""



								lowerCAmelCase_      			=  self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_question_answering(*a_   )







@require_tf
class  A		(  __UpperCamelCase      ,				unittest.TestCase     ):
		__snake_case									=   (
		    (TFFunnelBaseModel, TFFunnelForMultipleChoice, TFFunnelForSequenceClassification) if is_tf_available() else ()
		)
		__snake_case									=   False
		__snake_case									=   False



		def        SCREAMING_SNAKE_CASE__			(     self   ):







								"""simple docstring"""



								lowerCAmelCase_      			=  TFFunnelModelTester(self,	base=a_   )
								lowerCAmelCase_      			=  ConfigTester(self,	config_class=a_   )



		def        SCREAMING_SNAKE_CASE__			(     self   ):







								"""simple docstring"""



								self.config_tester.run_common_tests()



		def        SCREAMING_SNAKE_CASE__			(     self   ):







								"""simple docstring"""



								lowerCAmelCase_      			=  self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_base_model(*a_   )



		def        SCREAMING_SNAKE_CASE__			(     self   ):







								"""simple docstring"""



								lowerCAmelCase_      			=  self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_sequence_classification(*a_   )







		def        SCREAMING_SNAKE_CASE__			(     self   ):







								"""simple docstring"""



								lowerCAmelCase_      			=  self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_multiple_choice(*a_   )









import argparse
import shlex

import runhouse as rh


if __name__ == "__main__":
				# Refer to https://runhouse-docs.readthedocs-hosted.com/en/latest/api/python/cluster.html#hardware-setup for cloud access
				# setup instructions, if using on-demand hardware

				# If user passes --user <user> --host <host> --key_path <key_path> <example> <args>, fill them in as BYO cluster
				# If user passes --instance <instance> --provider <provider> <example> <args>, fill them in as on-demand cluster
				# Throw an error if user passes both BYO and on-demand cluster args
				# Otherwise, use default values
				__A   			=argparse.ArgumentParser()
				parser.add_argument("--user", type=str, default="ubuntu")
				parser.add_argument("--host", type=str, default="localhost")
				parser.add_argument("--key_path", type=str, default=None)
				parser.add_argument("--instance", type=str, default="V100:1")
				parser.add_argument("--provider", type=str, default="cheapest")
				parser.add_argument("--use_spot", type=bool, default=False)
				parser.add_argument("--example", type=str, default="pytorch/text-generation/run_generation.py")
				__A			,			__A   			=parser.parse_known_args()
				if args.host != "localhost":
								if args.instance != "V100:1" or args.provider != "cheapest":
												raise ValueError("Cannot specify both BYO and on-demand cluster args")
								__A   			=rh.cluster(
								    name="rh-cluster", ips=[args.host], ssh_creds={"ssh_user": args.user, "ssh_private_key": args.key_path}
								)
				else:
								__A   			=rh.cluster(
								    name="rh-cluster", instance_type=args.instance, provider=args.provider, use_spot=args.use_spot
								)
				__A   			=args.example.rsplit("/", 1)[0]

				# Set up remote environment
				cluster.install_packages(["pip:./"])  # Installs transformers from local source
				# Note transformers is copied into the home directory on the remote machine, so we can install from there
				cluster.run([f'''pip install -r transformers/examples/{example_dir}/requirements.txt'''])
				cluster.run(["pip install torch --upgrade --extra-index-url https://download.pytorch.org/whl/cu117"])

				# Run example. You can bypass the CLI wrapper and paste your own code here.
				cluster.run([f'''python transformers/examples/{args.example} {" ".join(shlex.quote(arg) for arg in unknown)}'''])

				# Alternatively, we can just import and run a training function (especially if there's no wrapper CLI):
				# from my_script... import train
				# reqs = ['pip:./', 'torch', 'datasets', 'accelerate', 'evaluate', 'tqdm', 'scipy', 'scikit-learn', 'tensorboard']
				# launch_train_gpu = rh.function(fn=train,
				#                                system=gpu,
				#                                reqs=reqs,
				#                                name='train_bert_glue')
				#
				# We can pass in arguments just like we would to a function:
				# launch_train_gpu(num_epochs = 3, lr = 2e-5, seed = 42, batch_size = 16
				#                  stream_logs=True)





