usrbinenv python codingutf8 2018 the google ai language team s and the huggingface inc team c 2018 nvidia corporation 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 multilingual models on xnli e g bert distilbert xlm adapted from examplestextclassificationrunglue py import logging import os import random import sys import warnings from dataclasses import dataclass field from typing import optional import datasets import evaluate import numpy as np from datasets import loaddataset import transformers from transformers import autoconfig automodelforsequenceclassification autotokenizer datacollatorwithpadding evalprediction hfargumentparser trainer trainingarguments defaultdatacollator setseed from transformers trainerutils import getlastcheckpoint from transformers utils import checkminversion sendexampletelemetry from transformers utils versions import requireversion will error if the minimal version of transformers is not installed remove at your own risks checkminversion4 36 0 dev0 requireversiondatasets1 8 0 to fix pip install r examplespytorchtextclassificationrequirements txt logger logging getloggername dataclass class datatrainingarguments maxseqlength optionalint 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 preprocessed datasets or not padtomaxlength bool field defaulttrue metadata help whether to pad all samples to maxseqlength if false will pad the samples dynamically when batching to the maximum length in the batch maxtrainsamples optionalint field defaultnone metadata help for debugging purposes or quicker training truncate the number of training examples to this value if set maxevalsamples optionalint field defaultnone metadata help for debugging purposes or quicker training truncate the number of evaluation examples to this value if set maxpredictsamples optionalint field defaultnone metadata help for debugging purposes or quicker training truncate the number of prediction examples to this value if set dataclass class modelarguments modelnameorpath str field defaultnone metadatahelp path to pretrained model or model identifier from huggingface comodels language str field defaultnone metadatahelp evaluation language also train language if trainlanguage is set to none trainlanguage optionalstr field defaultnone metadatahelp train language if it is different from the evaluation language 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 cachedir optionalstr field defaultnone metadatahelp where do you want to store the pretrained models downloaded from huggingface co dolowercase optionalbool field defaultfalse metadatahelp arg to indicate if tokenizer should do lower case in autotokenizer frompretrained usefasttokenizer bool field defaulttrue metadatahelp whether to use one of the fast tokenizer backed by the tokenizers library or not modelrevision str field defaultmain metadatahelp the specific model version to use can be a branch name tag name or commit id token str field defaultnone metadata help the token to use as http bearer ization for remote files if not specified will use the token generated when running huggingfacecli login stored in huggingface useauthtoken bool field defaultnone metadata help the useauthtoken argument is deprecated and will be removed in v4 34 please use token instead trustremotecode bool field defaultfalse metadata help whether or not to allow for custom models defined on the hub in their own modeling files this option should only be set to true for repositories you trust and in which you have read the code as it will execute code present on the hub on your local machine ignoremismatchedsizes bool field defaultfalse metadatahelp will enable to load a pretrained model whose head dimensions are different 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 trainingarguments modelargs dataargs trainingargs parser parseargsintodataclasses if modelargs useauthtoken is not none warnings warn the useauthtoken argument is deprecated and will be removed in v4 34 please use token instead futurewarning if modelargs token is not none raise valueerrortoken and useauthtoken are both specified please set only the argument token modelargs token modelargs useauthtoken sending telemetry tracking the example usage helps us better allocate resources to maintain them the information sent is the one passed as arguments along with your pythonpytorch versions sendexampletelemetryrunxnli modelargs setup logging logging basicconfig formatasctimes levelnames names messages datefmtmdy h m s handlerslogging streamhandlersys stdout if trainingargs shouldlog the default of trainingargs loglevel is passive so we set log level at info here to have that default transformers utils logging setverbosityinfo loglevel trainingargs getprocessloglevel logger setlevelloglevel datasets utils logging setverbosityloglevel transformers utils logging setverbosityloglevel transformers utils logging enabledefaulthandler transformers utils logging enableexplicitformat log on each process the small summary logger warning fprocess rank trainingargs localrank device trainingargs device ngpu trainingargs ngpu fdistributed training trainingargs parallelmode value distributed 16bits training trainingargs fp16 logger infoftrainingevaluation parameters trainingargs detecting last checkpoint lastcheckpoint none if os path isdirtrainingargs outputdir and trainingargs dotrain and not trainingargs overwriteoutputdir lastcheckpoint getlastcheckpointtrainingargs outputdir if lastcheckpoint is none and lenos listdirtrainingargs outputdir 0 raise valueerror foutput directory trainingargs outputdir already exists and is not empty use overwriteoutputdir to overcome elif lastcheckpoint is not none logger info fcheckpoint detected resuming training at lastcheckpoint to avoid this behavior change the outputdir or add overwriteoutputdir to train from scratch set seed before initializing model setseedtrainingargs seed in distributed training the loaddataset function guarantees that only one local process can concurrently download the dataset downloading and loading xnli dataset from the hub if trainingargs dotrain if modelargs trainlanguage is none traindataset loaddataset xnli modelargs language splittrain cachedirmodelargs cachedir tokenmodelargs token else traindataset loaddataset xnli modelargs trainlanguage splittrain cachedirmodelargs cachedir tokenmodelargs token labellist traindataset featureslabel names if trainingargs doeval evaldataset loaddataset xnli modelargs language splitvalidation cachedirmodelargs cachedir tokenmodelargs token labellist evaldataset featureslabel names if trainingargs dopredict predictdataset loaddataset xnli modelargs language splittest cachedirmodelargs cachedir tokenmodelargs token labellist predictdataset featureslabel names labels numlabels lenlabellist load pretrained model and tokenizer in distributed training the frompretrained methods guarantee that only one local process can concurrently download model vocab config autoconfig frompretrained modelargs configname if modelargs configname else modelargs modelnameorpath numlabelsnumlabels id2labelstri label for i label in enumeratelabellist label2idlabel i for i label in enumeratelabellist finetuningtaskxnli cachedirmodelargs cachedir revisionmodelargs modelrevision tokenmodelargs token trustremotecodemodelargs trustremotecode tokenizer autotokenizer frompretrained modelargs tokenizername if modelargs tokenizername else modelargs modelnameorpath dolowercasemodelargs dolowercase cachedirmodelargs cachedir usefastmodelargs usefasttokenizer revisionmodelargs modelrevision tokenmodelargs token trustremotecodemodelargs trustremotecode model automodelforsequenceclassification frompretrained modelargs modelnameorpath fromtfbool ckpt in modelargs modelnameorpath configconfig cachedirmodelargs cachedir revisionmodelargs modelrevision tokenmodelargs token trustremotecodemodelargs trustremotecode ignoremismatchedsizesmodelargs ignoremismatchedsizes preprocessing the datasets padding strategy if dataargs padtomaxlength padding maxlength else we will pad later dynamically at batch creation to the max sequence length in each batch padding false def preprocessfunctionexamples tokenize the texts return tokenizer examplespremise exampleshypothesis paddingpadding maxlengthdataargs maxseqlength truncationtrue if trainingargs dotrain if dataargs maxtrainsamples is not none maxtrainsamples minlentraindataset dataargs maxtrainsamples traindataset traindataset selectrangemaxtrainsamples with trainingargs mainprocessfirstdesctrain dataset map preprocessing traindataset traindataset map preprocessfunction batchedtrue loadfromcachefilenot dataargs overwritecache descrunning tokenizer on train dataset log a few random samples from the training set for index in random samplerangelentraindataset 3 logger infofsample index of the training set traindatasetindex if trainingargs doeval if dataargs maxevalsamples is not none maxevalsamples minlenevaldataset dataargs maxevalsamples evaldataset evaldataset selectrangemaxevalsamples with trainingargs mainprocessfirstdescvalidation dataset map preprocessing evaldataset evaldataset map preprocessfunction batchedtrue loadfromcachefilenot dataargs overwritecache descrunning tokenizer on validation dataset if trainingargs dopredict if dataargs maxpredictsamples is not none maxpredictsamples minlenpredictdataset dataargs maxpredictsamples predictdataset predictdataset selectrangemaxpredictsamples with trainingargs mainprocessfirstdescprediction dataset map preprocessing predictdataset predictdataset map preprocessfunction batchedtrue loadfromcachefilenot dataargs overwritecache descrunning tokenizer on prediction dataset get the metric function metric evaluate loadxnli you can define your custom computemetrics function it takes an evalprediction object a namedtuple with a predictions and labelids field and has to return a dictionary string to float def computemetricsp evalprediction preds p predictions0 if isinstancep predictions tuple else p predictions preds np argmaxpreds axis1 return metric computepredictionspreds referencesp labelids data collator will default to datacollatorwithpadding so we change it if we already did the padding if dataargs padtomaxlength datacollator defaultdatacollator elif trainingargs fp16 datacollator datacollatorwithpaddingtokenizer padtomultipleof8 else datacollator none initialize our trainer trainer trainer modelmodel argstrainingargs traindatasettraindataset if trainingargs dotrain else none evaldatasetevaldataset if trainingargs doeval else none computemetricscomputemetrics tokenizertokenizer datacollatordatacollator training if trainingargs dotrain checkpoint none if trainingargs resumefromcheckpoint is not none checkpoint trainingargs resumefromcheckpoint elif lastcheckpoint is not none checkpoint lastcheckpoint trainresult trainer trainresumefromcheckpointcheckpoint metrics trainresult metrics maxtrainsamples dataargs maxtrainsamples if dataargs maxtrainsamples is not none else lentraindataset metricstrainsamples minmaxtrainsamples lentraindataset trainer savemodel saves the tokenizer too for easy upload trainer logmetricstrain metrics trainer savemetricstrain metrics trainer savestate evaluation if trainingargs doeval logger info evaluate metrics trainer evaluateevaldatasetevaldataset maxevalsamples dataargs maxevalsamples if dataargs maxevalsamples is not none else lenevaldataset metricsevalsamples minmaxevalsamples lenevaldataset trainer logmetricseval metrics trainer savemetricseval metrics prediction if trainingargs dopredict logger info predict predictions labels metrics trainer predictpredictdataset metrickeyprefixpredict maxpredictsamples dataargs maxpredictsamples if dataargs maxpredictsamples is not none else lenpredictdataset metricspredictsamples minmaxpredictsamples lenpredictdataset trainer logmetricspredict metrics trainer savemetricspredict metrics predictions np argmaxpredictions axis1 outputpredictfile os path jointrainingargs outputdir predictions txt if trainer isworldprocesszero with openoutputpredictfile w as writer writer writeindextpredictionn for index item in enumeratepredictions item labellistitem writer writefindextitemn if name main main usr bin env python coding utf 8 2018 the google ai language team s and the huggingface inc team c 2018 nvidia corporation 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 finetuning multi lingual models on xnli e g bert distilbert xlm adapted from examples text classification run_glue py will error if the minimal version of transformers is not installed remove at your own risks 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 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 sending telemetry tracking the example usage helps us better allocate resources to maintain them the information sent is the one passed as arguments along with your python pytorch versions setup logging the default of training_args log_level is passive so we set log level at info here to have that default log on each process the small summary detecting last checkpoint set seed before initializing model in distributed training the load_dataset function guarantees that only one local process can concurrently download the dataset downloading and loading xnli dataset from the hub labels load pretrained model and tokenizer in distributed training the from_pretrained methods guarantee that only one local process can concurrently download model vocab preprocessing the datasets padding strategy we will pad later dynamically at batch creation to the max sequence length in each batch tokenize the texts log a few random samples from the training set get the metric function you can define your custom compute_metrics function it takes an evalprediction object a namedtuple with a predictions and label_ids field and has to return a dictionary string to float data collator will default to datacollatorwithpadding so we change it if we already did the padding initialize our trainer training saves the tokenizer too for easy upload evaluation prediction
import logging
import os
import random
import sys
import warnings
from dataclasses import dataclass, field
from typing import Optional

import datasets
import evaluate
import numpy as np
from datasets import load_dataset

import transformers
from transformers import (
    AutoConfig,
    AutoModelForSequenceClassification,
    AutoTokenizer,
    DataCollatorWithPadding,
    EvalPrediction,
    HfArgumentParser,
    Trainer,
    TrainingArguments,
    default_data_collator,
    set_seed,
)
from transformers.trainer_utils import get_last_checkpoint
from transformers.utils import check_min_version, send_example_telemetry
from transformers.utils.versions import require_version



check_min_version("4.36.0.dev0")

require_version("datasets>=1.8.0", "To fix: pip install -r examples/pytorch/text-classification/requirements.txt")

logger = logging.getLogger(__name__)


@dataclass
class DataTrainingArguments:
    

    max_seq_length: Optional[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 preprocessed datasets or not."}
    )
    pad_to_max_length: bool = field(
        default=True,
        metadata={
            "help": (
                "Whether to pad all samples to `max_seq_length`. "
                "If False, will pad the samples dynamically when batching to the maximum length in the batch."
            )
        },
    )
    max_train_samples: Optional[int] = field(
        default=None,
        metadata={
            "help": (
                "For debugging purposes or quicker training, truncate the number of training examples to this "
                "value if set."
            )
        },
    )
    max_eval_samples: Optional[int] = field(
        default=None,
        metadata={
            "help": (
                "For debugging purposes or quicker training, truncate the number of evaluation examples to this "
                "value if set."
            )
        },
    )
    max_predict_samples: Optional[int] = field(
        default=None,
        metadata={
            "help": (
                "For debugging purposes or quicker training, truncate the number of prediction examples to this "
                "value if set."
            )
        },
    )


@dataclass
class ModelArguments:
    

    model_name_or_path: str = field(
        default=None, metadata={"help": "Path to pretrained model or model identifier from huggingface.co/models"}
    )
    language: str = field(
        default=None, metadata={"help": "Evaluation language. Also train language if `train_language` is set to None."}
    )
    train_language: Optional[str] = field(
        default=None, metadata={"help": "Train language if it is different from the evaluation language."}
    )
    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"}
    )
    cache_dir: Optional[str] = field(
        default=None,
        metadata={"help": "Where do you want to store the pretrained models downloaded from huggingface.co"},
    )
    do_lower_case: Optional[bool] = field(
        default=False,
        metadata={"help": "arg to indicate if tokenizer should do lower case in AutoTokenizer.from_pretrained()"},
    )
    use_fast_tokenizer: bool = field(
        default=True,
        metadata={"help": "Whether to use one of the fast tokenizer (backed by the tokenizers library) or not."},
    )
    model_revision: str = field(
        default="main",
        metadata={"help": "The specific model version to use (can be a branch name, tag name or commit id)."},
    )
    token: str = field(
        default=None,
        metadata={
            "help": (
                "The token to use as HTTP bearer authorization for remote files. If not specified, will use the token "
                "generated when running `huggingface-cli login` (stored in `~/.huggingface`)."
            )
        },
    )
    use_auth_token: bool = field(
        default=None,
        metadata={
            "help": "The `use_auth_token` argument is deprecated and will be removed in v4.34. Please use `token` instead."
        },
    )
    trust_remote_code: bool = field(
        default=False,
        metadata={
            "help": (
                "Whether or not to allow for custom models defined on the Hub in their own modeling files. This option"
                "should only be set to `True` for repositories you trust and in which you have read the code, as it will "
                "execute code present on the Hub on your local machine."
            )
        },
    )
    ignore_mismatched_sizes: bool = field(
        default=False,
        metadata={"help": "Will enable to load a pretrained model whose head dimensions are different."},
    )


def main():
    
    
    

    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments))
    model_args, data_args, training_args = parser.parse_args_into_dataclasses()

    if model_args.use_auth_token is not None:
        warnings.warn(
            "The `use_auth_token` argument is deprecated and will be removed in v4.34. Please use `token` instead.",
            FutureWarning,
        )
        if model_args.token is not None:
            raise ValueError("`token` and `use_auth_token` are both specified. Please set only the argument `token`.")
        model_args.token = model_args.use_auth_token

    
    
    send_example_telemetry("run_xnli", model_args)

    
    logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
        datefmt="%m/%d/%Y %H:%M:%S",
        handlers=[logging.StreamHandler(sys.stdout)],
    )

    if training_args.should_log:
        
        transformers.utils.logging.set_verbosity_info()

    log_level = training_args.get_process_log_level()
    logger.setLevel(log_level)
    datasets.utils.logging.set_verbosity(log_level)
    transformers.utils.logging.set_verbosity(log_level)
    transformers.utils.logging.enable_default_handler()
    transformers.utils.logging.enable_explicit_format()

    
    logger.warning(
        f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}, "
        + f"distributed training: {training_args.parallel_mode.value == 'distributed'}, 16-bits training: {training_args.fp16}"
    )
    logger.info(f"Training/evaluation parameters {training_args}")

    
    last_checkpoint = None
    if os.path.isdir(training_args.output_dir) and training_args.do_train and not training_args.overwrite_output_dir:
        last_checkpoint = get_last_checkpoint(training_args.output_dir)
        if last_checkpoint is None and len(os.listdir(training_args.output_dir)) > 0:
            raise ValueError(
                f"Output directory ({training_args.output_dir}) already exists and is not empty. "
                "Use --overwrite_output_dir to overcome."
            )
        elif last_checkpoint is not None:
            logger.info(
                f"Checkpoint detected, resuming training at {last_checkpoint}. To avoid this behavior, change "
                "the `--output_dir` or add `--overwrite_output_dir` to train from scratch."
            )

    
    set_seed(training_args.seed)

    
    
    
    if training_args.do_train:
        if model_args.train_language is None:
            train_dataset = load_dataset(
                "xnli",
                model_args.language,
                split="train",
                cache_dir=model_args.cache_dir,
                token=model_args.token,
            )
        else:
            train_dataset = load_dataset(
                "xnli",
                model_args.train_language,
                split="train",
                cache_dir=model_args.cache_dir,
                token=model_args.token,
            )
        label_list = train_dataset.features["label"].names

    if training_args.do_eval:
        eval_dataset = load_dataset(
            "xnli",
            model_args.language,
            split="validation",
            cache_dir=model_args.cache_dir,
            token=model_args.token,
        )
        label_list = eval_dataset.features["label"].names

    if training_args.do_predict:
        predict_dataset = load_dataset(
            "xnli",
            model_args.language,
            split="test",
            cache_dir=model_args.cache_dir,
            token=model_args.token,
        )
        label_list = predict_dataset.features["label"].names

    
    num_labels = len(label_list)

    
    
    
    config = AutoConfig.from_pretrained(
        model_args.config_name if model_args.config_name else model_args.model_name_or_path,
        num_labels=num_labels,
        id2label={str(i): label for i, label in enumerate(label_list)},
        label2id={label: i for i, label in enumerate(label_list)},
        finetuning_task="xnli",
        cache_dir=model_args.cache_dir,
        revision=model_args.model_revision,
        token=model_args.token,
        trust_remote_code=model_args.trust_remote_code,
    )
    tokenizer = AutoTokenizer.from_pretrained(
        model_args.tokenizer_name if model_args.tokenizer_name else model_args.model_name_or_path,
        do_lower_case=model_args.do_lower_case,
        cache_dir=model_args.cache_dir,
        use_fast=model_args.use_fast_tokenizer,
        revision=model_args.model_revision,
        token=model_args.token,
        trust_remote_code=model_args.trust_remote_code,
    )
    model = AutoModelForSequenceClassification.from_pretrained(
        model_args.model_name_or_path,
        from_tf=bool(".ckpt" in model_args.model_name_or_path),
        config=config,
        cache_dir=model_args.cache_dir,
        revision=model_args.model_revision,
        token=model_args.token,
        trust_remote_code=model_args.trust_remote_code,
        ignore_mismatched_sizes=model_args.ignore_mismatched_sizes,
    )

    
    
    if data_args.pad_to_max_length:
        padding = "max_length"
    else:
        
        padding = False

    def preprocess_function(examples):
        
        return tokenizer(
            examples["premise"],
            examples["hypothesis"],
            padding=padding,
            max_length=data_args.max_seq_length,
            truncation=True,
        )

    if training_args.do_train:
        if data_args.max_train_samples is not None:
            max_train_samples = min(len(train_dataset), data_args.max_train_samples)
            train_dataset = train_dataset.select(range(max_train_samples))
        with training_args.main_process_first(desc="train dataset map pre-processing"):
            train_dataset = train_dataset.map(
                preprocess_function,
                batched=True,
                load_from_cache_file=not data_args.overwrite_cache,
                desc="Running tokenizer on train dataset",
            )
        
        for index in random.sample(range(len(train_dataset)), 3):
            logger.info(f"Sample {index} of the training set: {train_dataset[index]}.")

    if training_args.do_eval:
        if data_args.max_eval_samples is not None:
            max_eval_samples = min(len(eval_dataset), data_args.max_eval_samples)
            eval_dataset = eval_dataset.select(range(max_eval_samples))
        with training_args.main_process_first(desc="validation dataset map pre-processing"):
            eval_dataset = eval_dataset.map(
                preprocess_function,
                batched=True,
                load_from_cache_file=not data_args.overwrite_cache,
                desc="Running tokenizer on validation dataset",
            )

    if training_args.do_predict:
        if data_args.max_predict_samples is not None:
            max_predict_samples = min(len(predict_dataset), data_args.max_predict_samples)
            predict_dataset = predict_dataset.select(range(max_predict_samples))
        with training_args.main_process_first(desc="prediction dataset map pre-processing"):
            predict_dataset = predict_dataset.map(
                preprocess_function,
                batched=True,
                load_from_cache_file=not data_args.overwrite_cache,
                desc="Running tokenizer on prediction dataset",
            )

    
    metric = evaluate.load("xnli")

    
    
    def compute_metrics(p: EvalPrediction):
        preds = p.predictions[0] if isinstance(p.predictions, tuple) else p.predictions
        preds = np.argmax(preds, axis=1)
        return metric.compute(predictions=preds, references=p.label_ids)

    
    if data_args.pad_to_max_length:
        data_collator = default_data_collator
    elif training_args.fp16:
        data_collator = DataCollatorWithPadding(tokenizer, pad_to_multiple_of=8)
    else:
        data_collator = None

    
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset if training_args.do_train else None,
        eval_dataset=eval_dataset if training_args.do_eval else None,
        compute_metrics=compute_metrics,
        tokenizer=tokenizer,
        data_collator=data_collator,
    )

    
    if training_args.do_train:
        checkpoint = None
        if training_args.resume_from_checkpoint is not None:
            checkpoint = training_args.resume_from_checkpoint
        elif last_checkpoint is not None:
            checkpoint = last_checkpoint
        train_result = trainer.train(resume_from_checkpoint=checkpoint)
        metrics = train_result.metrics
        max_train_samples = (
            data_args.max_train_samples if data_args.max_train_samples is not None else len(train_dataset)
        )
        metrics["train_samples"] = min(max_train_samples, len(train_dataset))

        trainer.save_model()  

        trainer.log_metrics("train", metrics)
        trainer.save_metrics("train", metrics)
        trainer.save_state()

    
    if training_args.do_eval:
        logger.info("*** Evaluate ***")
        metrics = trainer.evaluate(eval_dataset=eval_dataset)

        max_eval_samples = data_args.max_eval_samples if data_args.max_eval_samples is not None else len(eval_dataset)
        metrics["eval_samples"] = min(max_eval_samples, len(eval_dataset))

        trainer.log_metrics("eval", metrics)
        trainer.save_metrics("eval", metrics)

    
    if training_args.do_predict:
        logger.info("*** Predict ***")
        predictions, labels, metrics = trainer.predict(predict_dataset, metric_key_prefix="predict")

        max_predict_samples = (
            data_args.max_predict_samples if data_args.max_predict_samples is not None else len(predict_dataset)
        )
        metrics["predict_samples"] = min(max_predict_samples, len(predict_dataset))

        trainer.log_metrics("predict", metrics)
        trainer.save_metrics("predict", metrics)

        predictions = np.argmax(predictions, axis=1)
        output_predict_file = os.path.join(training_args.output_dir, "predictions.txt")
        if trainer.is_world_process_zero():
            with open(output_predict_file, "w") as writer:
                writer.write("index\tprediction\n")
                for index, item in enumerate(predictions):
                    item = label_list[item]
                    writer.write(f"{index}\t{item}\n")


if __name__ == "__main__":
    main()