
"""simple docstring"""




import os
import tempfile
import unittest
from pathlib import Path

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


if is_tf_available():
    import tensorflow as tf

    from transformers import TensorFlowBenchmark, TensorFlowBenchmarkArguments







@require_tf
class     lowercase(unittest.TestCase       ):



     def 			lowercase__  (							self     ,				__SCREAMING_SNAKE_CASE   )		->     str:







         """simple docstring"""



         for model_result in results.values():
             for batch_size, sequence_length in zip(model_result['bs']     ,				model_result['ss']   ):
                 a__		=			model_result['result'][batch_size][sequence_length]
                 self.assertIsNotNone(__SCREAMING_SNAKE_CASE   )



     def 			lowercase__  (							self   )		->     Tuple:







         """simple docstring"""



         a__		=			'sshleifer/tiny-gpt2'
         a__		=			TensorFlowBenchmarkArguments(
             models=[MODEL_ID]     ,				training=__SCREAMING_SNAKE_CASE     ,				inference=__SCREAMING_SNAKE_CASE     ,				sequence_lengths=[8]     ,				batch_sizes=[1]     ,				eager_mode=__SCREAMING_SNAKE_CASE     ,				multi_process=__SCREAMING_SNAKE_CASE     ,				)
         a__		=			TensorFlowBenchmark(__SCREAMING_SNAKE_CASE   )
         a__		=			benchmark.run()
         self.check_results_dict_not_empty(results.time_inference_result   )
         self.check_results_dict_not_empty(results.memory_inference_result   )



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







         """simple docstring"""



         a__		=			'sgugger/tiny-distilbert-classification'
         a__		=			TensorFlowBenchmarkArguments(
             models=[MODEL_ID]     ,				training=__SCREAMING_SNAKE_CASE     ,				inference=__SCREAMING_SNAKE_CASE     ,				sequence_lengths=[8]     ,				batch_sizes=[1]     ,				multi_process=__SCREAMING_SNAKE_CASE     ,				only_pretrain_model=__SCREAMING_SNAKE_CASE     ,				)
         a__		=			TensorFlowBenchmark(__SCREAMING_SNAKE_CASE   )
         a__		=			benchmark.run()
         self.check_results_dict_not_empty(results.time_inference_result   )
         self.check_results_dict_not_empty(results.memory_inference_result   )



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







         """simple docstring"""



         a__		=			'sshleifer/tiny-gpt2'
         a__		=			TensorFlowBenchmarkArguments(
             models=[MODEL_ID]     ,				training=__SCREAMING_SNAKE_CASE     ,				inference=__SCREAMING_SNAKE_CASE     ,				sequence_lengths=[8]     ,				batch_sizes=[1]     ,				multi_process=__SCREAMING_SNAKE_CASE     ,				)
         a__		=			TensorFlowBenchmark(__SCREAMING_SNAKE_CASE   )
         a__		=			benchmark.run()
         self.check_results_dict_not_empty(results.time_inference_result   )
         self.check_results_dict_not_empty(results.memory_inference_result   )



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







         """simple docstring"""



         a__		=			'sshleifer/tiny-gpt2'
         a__		=			AutoConfig.from_pretrained(__SCREAMING_SNAKE_CASE   )
         a__		=			TensorFlowBenchmarkArguments(
             models=[MODEL_ID]     ,				training=__SCREAMING_SNAKE_CASE     ,				inference=__SCREAMING_SNAKE_CASE     ,				sequence_lengths=[8]     ,				batch_sizes=[1]     ,				eager_mode=__SCREAMING_SNAKE_CASE     ,				multi_process=__SCREAMING_SNAKE_CASE     ,				)
         a__		=			TensorFlowBenchmark(__SCREAMING_SNAKE_CASE     ,				[config]   )
         a__		=			benchmark.run()
         self.check_results_dict_not_empty(results.time_inference_result   )
         self.check_results_dict_not_empty(results.memory_inference_result   )



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







         """simple docstring"""



         a__		=			'sshleifer/tiny-gpt2'
         a__		=			AutoConfig.from_pretrained(__SCREAMING_SNAKE_CASE   )
         a__		=			TensorFlowBenchmarkArguments(
             models=[MODEL_ID]     ,				training=__SCREAMING_SNAKE_CASE     ,				inference=__SCREAMING_SNAKE_CASE     ,				sequence_lengths=[8]     ,				batch_sizes=[1]     ,				multi_process=__SCREAMING_SNAKE_CASE     ,				)
         a__		=			TensorFlowBenchmark(__SCREAMING_SNAKE_CASE     ,				[config]   )
         a__		=			benchmark.run()
         self.check_results_dict_not_empty(results.time_inference_result   )
         self.check_results_dict_not_empty(results.memory_inference_result   )



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







         """simple docstring"""



         a__		=			'sshleifer/tiny-gpt2'
         a__		=			TensorFlowBenchmarkArguments(
             models=[MODEL_ID]     ,				training=__SCREAMING_SNAKE_CASE     ,				inference=__SCREAMING_SNAKE_CASE     ,				sequence_lengths=[8]     ,				batch_sizes=[1]     ,				multi_process=__SCREAMING_SNAKE_CASE     ,				)
         a__		=			TensorFlowBenchmark(__SCREAMING_SNAKE_CASE   )
         a__		=			benchmark.run()
         self.check_results_dict_not_empty(results.time_train_result   )
         self.check_results_dict_not_empty(results.memory_train_result   )



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







         """simple docstring"""



         a__		=			'sshleifer/tiny-gpt2'
         a__		=			AutoConfig.from_pretrained(__SCREAMING_SNAKE_CASE   )
         a__		=			TensorFlowBenchmarkArguments(
             models=[MODEL_ID]     ,				training=__SCREAMING_SNAKE_CASE     ,				inference=__SCREAMING_SNAKE_CASE     ,				sequence_lengths=[8]     ,				batch_sizes=[1]     ,				multi_process=__SCREAMING_SNAKE_CASE     ,				)
         a__		=			TensorFlowBenchmark(__SCREAMING_SNAKE_CASE     ,				[config]   )
         a__		=			benchmark.run()
         self.check_results_dict_not_empty(results.time_train_result   )
         self.check_results_dict_not_empty(results.memory_train_result   )



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







         """simple docstring"""



         a__		=			'patrickvonplaten/t5-tiny-random'
         a__		=			AutoConfig.from_pretrained(__SCREAMING_SNAKE_CASE   )
         a__		=			TensorFlowBenchmarkArguments(
             models=[MODEL_ID]     ,				training=__SCREAMING_SNAKE_CASE     ,				inference=__SCREAMING_SNAKE_CASE     ,				sequence_lengths=[8]     ,				batch_sizes=[1]     ,				multi_process=__SCREAMING_SNAKE_CASE     ,				)
         a__		=			TensorFlowBenchmark(__SCREAMING_SNAKE_CASE     ,				configs=[config]   )
         a__		=			benchmark.run()
         self.check_results_dict_not_empty(results.time_inference_result   )
         self.check_results_dict_not_empty(results.memory_inference_result   )



     @unittest.skipIf(is_tf_available() and len(tf.config.list_physical_devices('GPU'   )   ) == 0     ,				'Cannot do xla on CPU.'   )
     def 			lowercase__  (							self   )		->     Optional[Any]:







         """simple docstring"""



         a__		=			'sshleifer/tiny-gpt2'
         a__		=			TensorFlowBenchmarkArguments(
             models=[MODEL_ID]     ,				training=__SCREAMING_SNAKE_CASE     ,				inference=__SCREAMING_SNAKE_CASE     ,				sequence_lengths=[8]     ,				batch_sizes=[1]     ,				use_xla=__SCREAMING_SNAKE_CASE     ,				multi_process=__SCREAMING_SNAKE_CASE     ,				)
         a__		=			TensorFlowBenchmark(__SCREAMING_SNAKE_CASE   )
         a__		=			benchmark.run()
         self.check_results_dict_not_empty(results.time_inference_result   )
         self.check_results_dict_not_empty(results.memory_inference_result   )



     def 			lowercase__  (							self   )		->     Tuple:







         """simple docstring"""



         a__		=			'sshleifer/tiny-gpt2'
         with tempfile.TemporaryDirectory() as tmp_dir:
             a__		=			TensorFlowBenchmarkArguments(
                 models=[MODEL_ID]     ,				inference=__SCREAMING_SNAKE_CASE     ,				save_to_csv=__SCREAMING_SNAKE_CASE     ,				sequence_lengths=[8]     ,				batch_sizes=[1]     ,				inference_time_csv_file=os.path.join(__SCREAMING_SNAKE_CASE     ,				'inf_time.csv'   )     ,				inference_memory_csv_file=os.path.join(__SCREAMING_SNAKE_CASE     ,				'inf_mem.csv'   )     ,				env_info_csv_file=os.path.join(__SCREAMING_SNAKE_CASE     ,				'env.csv'   )     ,				multi_process=__SCREAMING_SNAKE_CASE     ,				)
             a__		=			TensorFlowBenchmark(__SCREAMING_SNAKE_CASE   )
             benchmark.run()
             self.assertTrue(Path(os.path.join(__SCREAMING_SNAKE_CASE     ,				'inf_time.csv'   )   ).exists()   )
             self.assertTrue(Path(os.path.join(__SCREAMING_SNAKE_CASE     ,				'inf_mem.csv'   )   ).exists()   )
             self.assertTrue(Path(os.path.join(__SCREAMING_SNAKE_CASE     ,				'env.csv'   )   ).exists()   )




     def 			lowercase__  (							self   )		->     str:







         """simple docstring"""



         a__		=			'sshleifer/tiny-gpt2'

         def _check_summary_is_not_empty(__SCREAMING_SNAKE_CASE   ):
             self.assertTrue(hasattr(__SCREAMING_SNAKE_CASE     ,				'sequential'   )   )
             self.assertTrue(hasattr(__SCREAMING_SNAKE_CASE     ,				'cumulative'   )   )
             self.assertTrue(hasattr(__SCREAMING_SNAKE_CASE     ,				'current'   )   )
             self.assertTrue(hasattr(__SCREAMING_SNAKE_CASE     ,				'total'   )   )

         with tempfile.TemporaryDirectory() as tmp_dir:
             a__		=			TensorFlowBenchmarkArguments(
                 models=[MODEL_ID]     ,				inference=__SCREAMING_SNAKE_CASE     ,				sequence_lengths=[8]     ,				batch_sizes=[1]     ,				log_filename=os.path.join(__SCREAMING_SNAKE_CASE     ,				'log.txt'   )     ,				log_print=__SCREAMING_SNAKE_CASE     ,				trace_memory_line_by_line=__SCREAMING_SNAKE_CASE     ,				eager_mode=__SCREAMING_SNAKE_CASE     ,				multi_process=__SCREAMING_SNAKE_CASE     ,				)
             a__		=			TensorFlowBenchmark(__SCREAMING_SNAKE_CASE   )
             a__		=			benchmark.run()
             _check_summary_is_not_empty(result.inference_summary   )
             self.assertTrue(Path(os.path.join(__SCREAMING_SNAKE_CASE     ,				'log.txt'   )   ).exists()   )




"""simple docstring"""




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
a		:       Optional[Any]     						=       False
a		:       str     						=       False



def     __magic_name__		(					UpperCamelCase     :		Namespace     )     ->						Optional[int]:
    return TrainCommand(UpperCamelCase     )







class     lowercase(_lowercase       ):



     @staticmethod
     def 			lowercase__  (							__SCREAMING_SNAKE_CASE   )		->     List[Any]:







         """simple docstring"""



         a__		=			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=3_2     ,				help='Batch size for training.'   )
         train_parser.add_argument('--valid_batch_size'     ,				type=__SCREAMING_SNAKE_CASE     ,				default=6_4     ,				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   )		->     Union[str, Any]:







         """simple docstring"""



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

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

         os.makedirs(args.output     ,				exist_ok=__SCREAMING_SNAKE_CASE   )
         a__		=			args.output

         a__		=			args.column_label
         a__		=			args.column_text
         a__		=			args.column_id

         self.logger.info(f'Loading {args.task} pipeline for {args.model}'   )
         if args.task == "text_classification":
             a__		=			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}'   )
         a__		=			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     ,				)
         a__		=			None
         if args.validation_data:
             self.logger.info(f'Loading validation dataset from {args.validation_data}'   )
             a__		=			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     ,				)

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



     def 			lowercase__  (							self   )		->     Tuple:







         """simple docstring"""



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



     def 			lowercase__  (							self   )		->     Any:







         """simple docstring"""



         raise NotImplementedError




     def 			lowercase__  (							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   )


