




import unittest

import numpy as np

from transformers import RobertaConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask


if is_flax_available():
				from transformers.models.roberta.modeling_flax_roberta import (
				    FlaxRobertaForCausalLM,
				    FlaxRobertaForMaskedLM,
				    FlaxRobertaForMultipleChoice,
				    FlaxRobertaForQuestionAnswering,
				    FlaxRobertaForSequenceClassification,
				    FlaxRobertaForTokenClassification,
				    FlaxRobertaModel,
				)



class       a_       (     unittest.TestCase     ):





	def __init__(							self      ,     SCREAMING_SNAKE_CASE      ,     SCREAMING_SNAKE_CASE=13      ,     SCREAMING_SNAKE_CASE=7      ,     SCREAMING_SNAKE_CASE=True      ,     SCREAMING_SNAKE_CASE=True      ,     SCREAMING_SNAKE_CASE=True      ,     SCREAMING_SNAKE_CASE=True      ,     SCREAMING_SNAKE_CASE=99      ,     SCREAMING_SNAKE_CASE=32      ,     SCREAMING_SNAKE_CASE=5      ,     SCREAMING_SNAKE_CASE=4      ,     SCREAMING_SNAKE_CASE=37      ,     SCREAMING_SNAKE_CASE="gelu"      ,     SCREAMING_SNAKE_CASE=0.1      ,     SCREAMING_SNAKE_CASE=0.1      ,     SCREAMING_SNAKE_CASE=512      ,     SCREAMING_SNAKE_CASE=16      ,     SCREAMING_SNAKE_CASE=2      ,     SCREAMING_SNAKE_CASE=0.0_2      ,     SCREAMING_SNAKE_CASE=4      ,     )  ->					List[str]:




				"""simple docstring"""



				SCREAMING_SNAKE_CASE_								=						parent
				SCREAMING_SNAKE_CASE_								=						batch_size
				SCREAMING_SNAKE_CASE_								=						seq_length
				SCREAMING_SNAKE_CASE_								=						is_training
				SCREAMING_SNAKE_CASE_								=						use_attention_mask
				SCREAMING_SNAKE_CASE_								=						use_token_type_ids
				SCREAMING_SNAKE_CASE_								=						use_labels
				SCREAMING_SNAKE_CASE_								=						vocab_size
				SCREAMING_SNAKE_CASE_								=						hidden_size
				SCREAMING_SNAKE_CASE_								=						num_hidden_layers
				SCREAMING_SNAKE_CASE_								=						num_attention_heads
				SCREAMING_SNAKE_CASE_								=						intermediate_size
				SCREAMING_SNAKE_CASE_								=						hidden_act
				SCREAMING_SNAKE_CASE_								=						hidden_dropout_prob
				SCREAMING_SNAKE_CASE_								=						attention_probs_dropout_prob
				SCREAMING_SNAKE_CASE_								=						max_position_embeddings
				SCREAMING_SNAKE_CASE_								=						type_vocab_size
				SCREAMING_SNAKE_CASE_								=						type_sequence_label_size
				SCREAMING_SNAKE_CASE_								=						initializer_range
				SCREAMING_SNAKE_CASE_								=						num_choices





	def 						A_(							self      )  ->					Tuple:




				"""simple docstring"""



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

				SCREAMING_SNAKE_CASE_								=						None
				if self.use_attention_mask:
							SCREAMING_SNAKE_CASE_								=						random_attention_mask([self.batch_size, self.seq_length]      )

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

				SCREAMING_SNAKE_CASE_								=						RobertaConfig(
				    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=SCREAMING_SNAKE_CASE      ,     initializer_range=self.initializer_range      ,     )

				return config, input_ids, token_type_ids, attention_mask





	def 						A_(							self      )  ->					Tuple:




				"""simple docstring"""



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





	def 						A_(							self      )  ->					Dict:




				"""simple docstring"""



				SCREAMING_SNAKE_CASE_								=						self.prepare_config_and_inputs()
				SCREAMING_SNAKE_CASE_						,		SCREAMING_SNAKE_CASE_						,		SCREAMING_SNAKE_CASE_						,		SCREAMING_SNAKE_CASE_								=						config_and_inputs

				SCREAMING_SNAKE_CASE_								=						True
				SCREAMING_SNAKE_CASE_								=						floats_tensor([self.batch_size, self.seq_length, self.hidden_size]      )
				SCREAMING_SNAKE_CASE_								=						ids_tensor([self.batch_size, self.seq_length]      ,     vocab_size=2      )

				return (
				    config,
				    input_ids,
				    token_type_ids,
				    encoder_hidden_states,
				    encoder_attention_mask,
				)




@require_flax
class       a_       (     SCREAMING_SNAKE_CASE__  ,    unittest.TestCase     ):
	A    							= True

	A    							= (
	    (
	        FlaxRobertaModel,
	        FlaxRobertaForCausalLM,
	        FlaxRobertaForMaskedLM,
	        FlaxRobertaForSequenceClassification,
	        FlaxRobertaForTokenClassification,
	        FlaxRobertaForMultipleChoice,
	        FlaxRobertaForQuestionAnswering,
	    )
	    if is_flax_available()
	    else ()
	)





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




				"""simple docstring"""



				SCREAMING_SNAKE_CASE_								=						FlaxRobertaModelTester(self      )





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




				"""simple docstring"""



				for model_class_name in self.all_model_classes:
							SCREAMING_SNAKE_CASE_								=						model_class_name.from_pretrained('roberta-base'      ,     from_pt=SCREAMING_SNAKE_CASE      )
							SCREAMING_SNAKE_CASE_								=						model(np.ones((1, 1)      )      )
							self.assertIsNotNone(SCREAMING_SNAKE_CASE      )






import os
from argparse import ArgumentParser, Namespace

from ..data import SingleSentenceClassificationProcessor as Processor
from ..pipelines import TextClassificationPipeline
from ..utils import is_tf_available, is_torch_available, logging
from . import BaseTransformersCLICommand


if not is_tf_available() and not is_torch_available():
				raise RuntimeError("At least one of PyTorch or TensorFlow 2.0+ should be installed to use CLI training")

# TF training parameters
SCREAMING_SNAKE_CASE__     :     Union[str, Any]		    =       False
SCREAMING_SNAKE_CASE__     :     Optional[int]		    =       False

def lowercase    (  SCREAMING_SNAKE_CASE )			->							List[Any]:


			'''simple docstring'''







			return TrainCommand(SCREAMING_SNAKE_CASE )



class       a_       (     SCREAMING_SNAKE_CASE__     ):





	@staticmethod
	def 						A_(							SCREAMING_SNAKE_CASE      )  ->					Optional[int]:




				"""simple docstring"""



				SCREAMING_SNAKE_CASE_								=						parser.add_parser('train'      ,     help='CLI tool to train a model on a task.'      )

				train_parser.add_argument(
				    '--train_data'      ,     type=SCREAMING_SNAKE_CASE      ,     required=SCREAMING_SNAKE_CASE      ,     help='path to train (and optionally evaluation) dataset as a csv with tab separated labels and sentences.'      ,     )
				train_parser.add_argument(
				    '--column_label'      ,     type=SCREAMING_SNAKE_CASE      ,     default=0      ,     help='Column of the dataset csv file with example labels.'      )
				train_parser.add_argument(
				    '--column_text'      ,     type=SCREAMING_SNAKE_CASE      ,     default=1      ,     help='Column of the dataset csv file with example texts.'      )
				train_parser.add_argument(
				    '--column_id'      ,     type=SCREAMING_SNAKE_CASE      ,     default=2      ,     help='Column of the dataset csv file with example ids.'      )
				train_parser.add_argument(
				    '--skip_first_row'      ,     action='store_true'      ,     help='Skip the first row of the csv file (headers).'      )

				train_parser.add_argument('--validation_data'      ,     type=SCREAMING_SNAKE_CASE      ,     default=''      ,     help='path to validation dataset.'      )
				train_parser.add_argument(
				    '--validation_split'      ,     type=SCREAMING_SNAKE_CASE      ,     default=0.1      ,     help='if validation dataset is not provided, fraction of train dataset to use as validation dataset.'      ,     )

				train_parser.add_argument('--output'      ,     type=SCREAMING_SNAKE_CASE      ,     default='./'      ,     help='path to saved the trained model.'      )

				train_parser.add_argument(
				    '--task'      ,     type=SCREAMING_SNAKE_CASE      ,     default='text_classification'      ,     help='Task to train the model on.'      )
				train_parser.add_argument(
				    '--model'      ,     type=SCREAMING_SNAKE_CASE      ,     default='bert-base-uncased'      ,     help='Model\'s name or path to stored model.'      )
				train_parser.add_argument('--train_batch_size'      ,     type=SCREAMING_SNAKE_CASE      ,     default=32      ,     help='Batch size for training.'      )
				train_parser.add_argument('--valid_batch_size'      ,     type=SCREAMING_SNAKE_CASE      ,     default=64      ,     help='Batch size for validation.'      )
				train_parser.add_argument('--learning_rate'      ,     type=SCREAMING_SNAKE_CASE      ,     default=3e-5      ,     help='Learning rate.'      )
				train_parser.add_argument('--adam_epsilon'      ,     type=SCREAMING_SNAKE_CASE      ,     default=1e-08      ,     help='Epsilon for Adam optimizer.'      )
				train_parser.set_defaults(func=SCREAMING_SNAKE_CASE      )





	def __init__(							self      ,     SCREAMING_SNAKE_CASE      )  ->					List[Any]:




				"""simple docstring"""



				SCREAMING_SNAKE_CASE_								=						logging.get_logger('transformers-cli/training'      )

				SCREAMING_SNAKE_CASE_								=						'tf' if is_tf_available() else 'torch'

				os.makedirs(args.output      ,     exist_ok=SCREAMING_SNAKE_CASE      )
				SCREAMING_SNAKE_CASE_								=						args.output

				SCREAMING_SNAKE_CASE_								=						args.column_label
				SCREAMING_SNAKE_CASE_								=						args.column_text
				SCREAMING_SNAKE_CASE_								=						args.column_id

				self.logger.info(f'Loading {args.task} pipeline for {args.model}'      )
				if args.task == "text_classification":
							SCREAMING_SNAKE_CASE_								=						TextClassificationPipeline.from_pretrained(args.model      )
				elif args.task == "token_classification":
							raise NotImplementedError
				elif args.task == "question_answering":
							raise NotImplementedError

				self.logger.info(f'Loading dataset from {args.train_data}'      )
				SCREAMING_SNAKE_CASE_								=						Processor.create_from_csv(
				    args.train_data      ,     column_label=args.column_label      ,     column_text=args.column_text      ,     column_id=args.column_id      ,     skip_first_row=args.skip_first_row      ,     )
				SCREAMING_SNAKE_CASE_								=						None
				if args.validation_data:
							self.logger.info(f'Loading validation dataset from {args.validation_data}'      )
							SCREAMING_SNAKE_CASE_								=						Processor.create_from_csv(
							    args.validation_data      ,     column_label=args.column_label      ,     column_text=args.column_text      ,     column_id=args.column_id      ,     skip_first_row=args.skip_first_row      ,     )

				SCREAMING_SNAKE_CASE_								=						args.validation_split
				SCREAMING_SNAKE_CASE_								=						args.train_batch_size
				SCREAMING_SNAKE_CASE_								=						args.valid_batch_size
				SCREAMING_SNAKE_CASE_								=						args.learning_rate
				SCREAMING_SNAKE_CASE_								=						args.adam_epsilon





	def 						A_(							self      )  ->					Dict:




				"""simple docstring"""



				if self.framework == "tf":
							return self.run_tf()
				return self.run_torch()





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




				"""simple docstring"""



				raise NotImplementedError





	def 						A_(							self      )  ->					Dict:




				"""simple docstring"""



				self.pipeline.fit(
				    self.train_dataset      ,     validation_data=self.valid_dataset      ,     validation_split=self.validation_split      ,     learning_rate=self.learning_rate      ,     adam_epsilon=self.adam_epsilon      ,     train_batch_size=self.train_batch_size      ,     valid_batch_size=self.valid_batch_size      ,     )

				# Save trained pipeline
				self.pipeline.save_pretrained(self.output      )
