usrbinenv python codingutf8 2020 the huggingface team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license finetuning the library models for sequence classification import logging import os from dataclasses import dataclass field from typing import dict optional import datasets import numpy as np import tensorflow as tf from transformers import autoconfig autotokenizer evalprediction hfargumentparser pretrainedtokenizer tfautomodelforsequenceclassification tftrainer tftrainingarguments from transformers utils import logging as hflogging hflogging setverbosityinfo hflogging enabledefaulthandler hflogging enableexplicitformat def gettfds trainfile str evalfile str testfile str tokenizer pretrainedtokenizer labelcolumnid int maxseqlength optionalint none files if trainfile is not none filesdatasets split train trainfile if evalfile is not none filesdatasets split validation evalfile if testfile is not none filesdatasets split test testfile ds datasets loaddatasetcsv datafilesfiles featuresname listdslistfiles keys0 features keys labelname featuresname poplabelcolumnid labellist listsetdslistfiles keys0labelname label2id label i for i label in enumeratelabellist inputnames tokenizer modelinputnames transformedds if lenfeaturesname 1 for k in files keys transformeddsk dsk map lambda example tokenizer batchencodeplus examplefeaturesname0 truncationtrue maxlengthmaxseqlength paddingmaxlength batchedtrue elif lenfeaturesname 2 for k in files keys transformeddsk dsk map lambda example tokenizer batchencodeplus examplefeaturesname0 examplefeaturesname1 truncationtrue maxlengthmaxseqlength paddingmaxlength batchedtrue def gentrain for ex in transformeddsdatasets split train d k v for k v in ex items if k in inputnames label label2idexlabelname yield d label def genval for ex in transformeddsdatasets split validation d k v for k v in ex items if k in inputnames label label2idexlabelname yield d label def gentest for ex in transformeddsdatasets split test d k v for k v in ex items if k in inputnames label label2idexlabelname yield d label trainds tf data dataset fromgenerator gentrain k tf int32 for k in inputnames tf int64 k tf tensorshapenone for k in inputnames tf tensorshape if datasets split train in transformedds else none if trainds is not none trainds trainds applytf data experimental assertcardinalitylendsdatasets split train valds tf data dataset fromgenerator genval k tf int32 for k in inputnames tf int64 k tf tensorshapenone for k in inputnames tf tensorshape if datasets split validation in transformedds else none if valds is not none valds valds applytf data experimental assertcardinalitylendsdatasets split validation testds tf data dataset fromgenerator gentest k tf int32 for k in inputnames tf int64 k tf tensorshapenone for k in inputnames tf tensorshape if datasets split test in transformedds else none if testds is not none testds testds applytf data experimental assertcardinalitylendsdatasets split test return trainds valds testds label2id logger logging getloggername dataclass class datatrainingarguments labelcolumnid int fieldmetadatahelp which column contains the label trainfile str fielddefaultnone metadatahelp the path of the training file devfile optionalstr fielddefaultnone metadatahelp the path of the development file testfile optionalstr fielddefaultnone metadatahelp the path of the test file maxseqlength int field default128 metadata help the maximum total input sequence length after tokenization sequences longer than this will be truncated sequences shorter will be padded overwritecache bool field defaultfalse metadatahelp overwrite the cached training and evaluation sets dataclass class modelarguments modelnameorpath str field metadatahelp path to pretrained model or model identifier from huggingface comodels configname optionalstr field defaultnone metadatahelp pretrained config name or path if not the same as modelname tokenizername optionalstr field defaultnone metadatahelp pretrained tokenizer name or path if not the same as modelname usefast bool fielddefaultfalse metadatahelp set this flag to use fast tokenization if you want to tweak more attributes on your tokenizer you should do it in a distinct script or just modify its tokenizerconfig json cachedir optionalstr field defaultnone metadatahelp where do you want to store the pretrained models downloaded from huggingface co def main see all possible arguments in srctransformerstrainingargs py or by passing the help flag to this script we now keep distinct sets of args for a cleaner separation of concerns parser hfargumentparsermodelarguments datatrainingarguments tftrainingarguments modelargs dataargs trainingargs parser parseargsintodataclasses if os path existstrainingargs outputdir and os listdirtrainingargs outputdir and trainingargs dotrain and not trainingargs overwriteoutputdir raise valueerror foutput directory trainingargs outputdir already exists and is not empty use overwriteoutputdir to overcome setup logging logging basicconfig formatasctimes levelnames names messages datefmtmdy h m s levellogging info logger info fnreplicas trainingargs nreplicas distributed training booltrainingargs nreplicas 1 f16bits training trainingargs fp16 logger infoftrainingevaluation parameters trainingargs load pretrained model and tokenizer distributed training the frompretrained methods guarantee that only one local process can concurrently download model vocab tokenizer autotokenizer frompretrained modelargs tokenizername if modelargs tokenizername else modelargs modelnameorpath cachedirmodelargs cachedir traindataset evaldataset testds label2id gettfds trainfiledataargs trainfile evalfiledataargs devfile testfiledataargs testfile tokenizertokenizer labelcolumniddataargs labelcolumnid maxseqlengthdataargs maxseqlength config autoconfig frompretrained modelargs configname if modelargs configname else modelargs modelnameorpath numlabelslenlabel2id label2idlabel2id id2labelid label for label id in label2id items finetuningtasktextclassification cachedirmodelargs cachedir with trainingargs strategy scope model tfautomodelforsequenceclassification frompretrained modelargs modelnameorpath fromptbool bin in modelargs modelnameorpath configconfig cachedirmodelargs cachedir def computemetricsp evalprediction dict preds np argmaxp predictions axis1 return acc preds p labelids mean initialize our trainer trainer tftrainer modelmodel argstrainingargs traindatasettraindataset evaldatasetevaldataset computemetricscomputemetrics training if trainingargs dotrain trainer train trainer savemodel tokenizer savepretrainedtrainingargs outputdir evaluation results if trainingargs doeval logger info evaluate result trainer evaluate outputevalfile os path jointrainingargs outputdir evalresults txt with openoutputevalfile w as writer logger info eval results for key value in result items logger infof key value writer writefkey valuen results updateresult return results if name main main usr bin env python coding utf 8 2020 the huggingface team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license fine tuning the library models for sequence classification arguments pertaining to what data we are going to input our model for training and eval using hfargumentparser we can turn this class into argparse arguments to be able to specify them on the command line arguments pertaining to which model config tokenizer we are going to fine tune from if you want to tweak more attributes on your tokenizer you should do it in a distinct script or just modify its tokenizer_config json see all possible arguments in src transformers training_args py or by passing the help flag to this script we now keep distinct sets of args for a cleaner separation of concerns setup logging load pretrained model and tokenizer distributed training the from_pretrained methods guarantee that only one local process can concurrently download model vocab initialize our trainer training evaluation
import logging
import os
from dataclasses import dataclass, field
from typing import Dict, Optional

import datasets
import numpy as np
import tensorflow as tf

from transformers import (
    AutoConfig,
    AutoTokenizer,
    EvalPrediction,
    HfArgumentParser,
    PreTrainedTokenizer,
    TFAutoModelForSequenceClassification,
    TFTrainer,
    TFTrainingArguments,
)
from transformers.utils import logging as hf_logging


hf_logging.set_verbosity_info()
hf_logging.enable_default_handler()
hf_logging.enable_explicit_format()


def get_tfds(
    train_file: str,
    eval_file: str,
    test_file: str,
    tokenizer: PreTrainedTokenizer,
    label_column_id: int,
    max_seq_length: Optional[int] = None,
):
    files = {}

    if train_file is not None:
        files[datasets.Split.TRAIN] = [train_file]
    if eval_file is not None:
        files[datasets.Split.VALIDATION] = [eval_file]
    if test_file is not None:
        files[datasets.Split.TEST] = [test_file]

    ds = datasets.load_dataset("csv", data_files=files)
    features_name = list(ds[list(files.keys())[0]].features.keys())
    label_name = features_name.pop(label_column_id)
    label_list = list(set(ds[list(files.keys())[0]][label_name]))
    label2id = {label: i for i, label in enumerate(label_list)}
    input_names = tokenizer.model_input_names
    transformed_ds = {}

    if len(features_name) == 1:
        for k in files.keys():
            transformed_ds[k] = ds[k].map(
                lambda example: tokenizer.batch_encode_plus(
                    example[features_name[0]], truncation=True, max_length=max_seq_length, padding="max_length"
                ),
                batched=True,
            )
    elif len(features_name) == 2:
        for k in files.keys():
            transformed_ds[k] = ds[k].map(
                lambda example: tokenizer.batch_encode_plus(
                    (example[features_name[0]], example[features_name[1]]),
                    truncation=True,
                    max_length=max_seq_length,
                    padding="max_length",
                ),
                batched=True,
            )

    def gen_train():
        for ex in transformed_ds[datasets.Split.TRAIN]:
            d = {k: v for k, v in ex.items() if k in input_names}
            label = label2id[ex[label_name]]
            yield (d, label)

    def gen_val():
        for ex in transformed_ds[datasets.Split.VALIDATION]:
            d = {k: v for k, v in ex.items() if k in input_names}
            label = label2id[ex[label_name]]
            yield (d, label)

    def gen_test():
        for ex in transformed_ds[datasets.Split.TEST]:
            d = {k: v for k, v in ex.items() if k in input_names}
            label = label2id[ex[label_name]]
            yield (d, label)

    train_ds = (
        tf.data.Dataset.from_generator(
            gen_train,
            ({k: tf.int32 for k in input_names}, tf.int64),
            ({k: tf.TensorShape([None]) for k in input_names}, tf.TensorShape([])),
        )
        if datasets.Split.TRAIN in transformed_ds
        else None
    )

    if train_ds is not None:
        train_ds = train_ds.apply(tf.data.experimental.assert_cardinality(len(ds[datasets.Split.TRAIN])))

    val_ds = (
        tf.data.Dataset.from_generator(
            gen_val,
            ({k: tf.int32 for k in input_names}, tf.int64),
            ({k: tf.TensorShape([None]) for k in input_names}, tf.TensorShape([])),
        )
        if datasets.Split.VALIDATION in transformed_ds
        else None
    )

    if val_ds is not None:
        val_ds = val_ds.apply(tf.data.experimental.assert_cardinality(len(ds[datasets.Split.VALIDATION])))

    test_ds = (
        tf.data.Dataset.from_generator(
            gen_test,
            ({k: tf.int32 for k in input_names}, tf.int64),
            ({k: tf.TensorShape([None]) for k in input_names}, tf.TensorShape([])),
        )
        if datasets.Split.TEST in transformed_ds
        else None
    )

    if test_ds is not None:
        test_ds = test_ds.apply(tf.data.experimental.assert_cardinality(len(ds[datasets.Split.TEST])))

    return train_ds, val_ds, test_ds, label2id


logger = logging.getLogger(__name__)


@dataclass
class DataTrainingArguments:
    

    label_column_id: int = field(metadata={"help": "Which column contains the label"})
    train_file: str = field(default=None, metadata={"help": "The path of the training file"})
    dev_file: Optional[str] = field(default=None, metadata={"help": "The path of the development file"})
    test_file: Optional[str] = field(default=None, metadata={"help": "The path of the test file"})
    max_seq_length: int = field(
        default=128,
        metadata={
            "help": (
                "The maximum total input sequence length after tokenization. Sequences longer "
                "than this will be truncated, sequences shorter will be padded."
            )
        },
    )
    overwrite_cache: bool = field(
        default=False, metadata={"help": "Overwrite the cached training and evaluation sets"}
    )


@dataclass
class ModelArguments:
    

    model_name_or_path: str = field(
        metadata={"help": "Path to pretrained model or model identifier from huggingface.co/models"}
    )
    config_name: Optional[str] = field(
        default=None, metadata={"help": "Pretrained config name or path if not the same as model_name"}
    )
    tokenizer_name: Optional[str] = field(
        default=None, metadata={"help": "Pretrained tokenizer name or path if not the same as model_name"}
    )
    use_fast: bool = field(default=False, metadata={"help": "Set this flag to use fast tokenization."})
    
    
    cache_dir: Optional[str] = field(
        default=None,
        metadata={"help": "Where do you want to store the pretrained models downloaded from huggingface.co"},
    )


def main():
    
    
    
    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TFTrainingArguments))
    model_args, data_args, training_args = parser.parse_args_into_dataclasses()

    if (
        os.path.exists(training_args.output_dir)
        and os.listdir(training_args.output_dir)
        and training_args.do_train
        and not training_args.overwrite_output_dir
    ):
        raise ValueError(
            f"Output directory ({training_args.output_dir}) already exists and is not empty. Use"
            " --overwrite_output_dir to overcome."
        )

    
    logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
        datefmt="%m/%d/%Y %H:%M:%S",
        level=logging.INFO,
    )
    logger.info(
        f"n_replicas: {training_args.n_replicas}, distributed training: {bool(training_args.n_replicas > 1)}, "
        f"16-bits training: {training_args.fp16}"
    )
    logger.info(f"Training/evaluation parameters {training_args}")

    
    
    
    
    

    tokenizer = AutoTokenizer.from_pretrained(
        model_args.tokenizer_name if model_args.tokenizer_name else model_args.model_name_or_path,
        cache_dir=model_args.cache_dir,
    )

    train_dataset, eval_dataset, test_ds, label2id = get_tfds(
        train_file=data_args.train_file,
        eval_file=data_args.dev_file,
        test_file=data_args.test_file,
        tokenizer=tokenizer,
        label_column_id=data_args.label_column_id,
        max_seq_length=data_args.max_seq_length,
    )

    config = AutoConfig.from_pretrained(
        model_args.config_name if model_args.config_name else model_args.model_name_or_path,
        num_labels=len(label2id),
        label2id=label2id,
        id2label={id: label for label, id in label2id.items()},
        finetuning_task="text-classification",
        cache_dir=model_args.cache_dir,
    )

    with training_args.strategy.scope():
        model = TFAutoModelForSequenceClassification.from_pretrained(
            model_args.model_name_or_path,
            from_pt=bool(".bin" in model_args.model_name_or_path),
            config=config,
            cache_dir=model_args.cache_dir,
        )

    def compute_metrics(p: EvalPrediction) -> Dict:
        preds = np.argmax(p.predictions, axis=1)

        return {"acc": (preds == p.label_ids).mean()}

    
    trainer = TFTrainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        compute_metrics=compute_metrics,
    )

    
    if training_args.do_train:
        trainer.train()
        trainer.save_model()
        tokenizer.save_pretrained(training_args.output_dir)

    
    results = {}
    if training_args.do_eval:
        logger.info("*** Evaluate ***")
        result = trainer.evaluate()
        output_eval_file = os.path.join(training_args.output_dir, "eval_results.txt")

        with open(output_eval_file, "w") as writer:
            logger.info("***** Eval results *****")

            for key, value in result.items():
                logger.info(f"  {key} = {value}")
                writer.write(f"{key} = {value}\n")

            results.update(result)

    return results


if __name__ == "__main__":
    main()