U &c@sddlZddlmZmZddlmZddlmZddlm Z m Z m Z ddl m Z e sde sdeddZdZed d d ZGd d d e ZdS)N)ArgumentParser Namespace) getLogger)%SingleSentenceClassificationProcessor)TextClassificationPipelineis_tf_availableis_torch_available)BaseTransformersCLICommandzRAt least one of PyTorch or TensorFlow 2.0+ should be installed to use CLI trainingFargscCst|S)z} Factory function used to instantiate serving server from provided command line arguments. :return: ServeCommand ) TrainCommandr r ?/tmp/pip-unpacked-wheel-ymerj3tt/transformers/commands/train.pytrain_command_factorysrc@sDeZdZeedddZedddZddZd d Z d d Z d S)r )parsercCs|jddd}|jdtddd|jdtd d d |jd td dd |jdtddd |jdddd|jdtddd |jdtddd |jdtddd |jdtd d!d |jd"td#d$d |jd%td&d'd |jd(td)d*d |jd+td,d-d |jd.td/d0d |jtd1d2S)3z Register this command to argparse so it's available for the transformer-cli :param parser: Root parser to register command-specific arguments :return: Ztrainz$CLI tool to train a model on a task.)helpz --train_dataTzcpath to train (and optionally evaluation) dataset as a csv with tab separated labels and sentences.)typerequiredrz--column_labelrz3Column of the dataset csv file with example labels.)rdefaultrz --column_textz2Column of the dataset csv file with example texts.z --column_idz0Column of the dataset csv file with example ids.z--skip_first_row store_truez-Skip the first row of the csv file (headers).)actionrz--validation_datazpath to validation dataset.z--validation_splitg?z^if validation dataset is not provided, fraction of train dataset to use as validation dataset.z--outputz./z path to saved the trained model.z--tasktext_classificationzTask to train the model on.z--modelzbert-base-uncasedz%Model's name or path to stored model.z--train_batch_size zBatch size for training.z--valid_batch_size@zBatch size for validation.z--learning_rategiUMu>zLearning rate.z--adam_epsilong:0yE>zEpsilon for Adam optimizer.)funcN) add_parser add_argumentstrintfloat set_defaultsr)rZ train_parserr r rregister_subcommandsnz TrainCommand.register_subcommandr cCsNtd|_trdnd|_tj|jddtj|js) r6Zfitr9r:r<r?r@r=r>Zsave_pretrainedr/rEr r rrCs zTrainCommand.run_tfN) __name__ __module__ __qualname__ staticmethodrr$rrBrFrDrCr r r rr s 2.r )r-argparserrloggingrZ transformersrr8rrrZtransformers.commandsr RuntimeErrorZUSE_XLAZUSE_AMPrr r r r rs