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 tensorflow trainer class import datetime import math import os import warnings from typing import callable dict optional tuple from utils import envvarstruevalues integrations must be imported before ml frameworks isort off from integrations import iscometavailable iswandbavailable isort on import numpy as np import tensorflow as tf from tensorflow python distribute values import perreplica from modelingtfutils import tfpretrainedmodel from optimizationtf import gradientaccumulator createoptimizer from trainerutils import prefixcheckpointdir evalprediction intervalstrategy predictionoutput enablefulldeterminism setseed from trainingargstf import tftrainingarguments from utils import logging if iswandbavailable import wandb if iscometavailable import cometml logger logging getloggername class tftrainer def init self model tfpretrainedmodel args tftrainingarguments traindataset optionaltf data dataset none evaldataset optionaltf data dataset none computemetrics optionalcallableevalprediction dict none tbwriter optionaltf summary summarywriter none optimizers tupletf keras optimizers optimizer tf keras optimizers schedules learningrateschedule none none self model model self args args self traindataset traindataset self evaldataset evaldataset self computemetrics computemetrics self optimizer self lrscheduler optimizers self gradientaccumulator gradientaccumulator self globalstep 0 self epochlogging 0 self evalloss tf keras metrics sum warnings warn the class tftrainer is deprecated and will be removed in version 5 of transformers we recommend using native keras instead by calling methods like fit and predict directly on the model object detailed examples of the keras style can be found in our examples at https github comhuggingfacetransformerstreemainexamplestensorflow futurewarning if tbwriter is not none self tbwriter tbwriter else self tbwriter tf summary createfilewriterself args loggingdir if iswandbavailable self setupwandb elif os getenvwandbdisabled upper not in envvarstruevalues logger info you are instantiating a trainer but wb is not installed to use wandb logging run pip install wandb wandb login see https docs wandb comhuggingface if iscometavailable self setupcomet elif os environ getcometmode disabled logger info to use cometml logging run pipconda install cometml see https www comet mldocspythonsdkhuggingface enablefulldeterminismself args seed if self args fulldeterminism else setseedself args seed def gettraintfdatasetself tf data dataset if self traindataset is none raise valueerrortrainer training requires a traindataset self totaltrainbatchsize self args trainbatchsize self args gradientaccumulationsteps self numtrainexamples self traindataset cardinality numpy if self numtrainexamples 0 raise valueerrorthe training dataset must have an asserted cardinality ds self traindataset repeat shuffleself numtrainexamples seedself args seed batchself totaltrainbatchsize dropremainderself args dataloaderdroplast prefetchtf data experimental autotune return self args strategy experimentaldistributedatasetds def getevaltfdatasetself evaldataset optionaltf data dataset none tf data dataset if evaldataset is none and self evaldataset is none raise valueerrortrainer evaluation requires an evaldataset evaldataset evaldataset if evaldataset is not none else self evaldataset numexamples evaldataset cardinality numpy if numexamples 0 raise valueerrorthe training dataset must have an asserted cardinality approx math floor if self args dataloaderdroplast else math ceil steps approxnumexamples self args evalbatchsize ds evaldataset repeat batchself args evalbatchsize dropremainderself args dataloaderdroplast prefetchtf data experimental autotune return self args strategy experimentaldistributedatasetds steps numexamples def gettesttfdatasetself testdataset tf data dataset tf data dataset numexamples testdataset cardinality numpy if numexamples 0 raise valueerrorthe training dataset must have an asserted cardinality steps math ceilnumexamples self args evalbatchsize ds testdataset batchself args evalbatchsize prefetchtf data experimental autotune return self args strategy experimentaldistributedatasetds steps numexamples def createoptimizerandschedulerself numtrainingsteps int if not self optimizer and not self lrscheduler warmupsteps self args warmupsteps if self args warmupsteps 0 else math ceilnumtrainingsteps self args warmupratio self optimizer self lrscheduler createoptimizer self args learningrate numtrainingsteps warmupsteps adambeta1self args adambeta1 adambeta2self args adambeta2 adamepsilonself args adamepsilon weightdecayrateself args weightdecay powerself args polypower def setupwandbself logger info automatic weights biases logging enabled to disable set os environwandbdisabled true combineddict self model config todict self args tosanitizeddict wandb initprojectos getenvwandbproject huggingface configcombineddict nameself args runname def setupcometself cometmode os getenvcometmode online upper args projectname os getenvcometprojectname huggingface experiment none if cometmode online experiment cometml experimentargs logger infoautomatic comet ml online logging enabled elif cometmode offline argsofflinedirectory os getenvcometofflinedirectory experiment cometml offlineexperimentargs logger infoautomatic comet ml offline logging enabled use comet upload when finished if experiment is not none experiment setmodelgraphself model frameworktransformers experiment logparametersself args prefixargs frameworktransformers experiment logparametersself model config prefixconfig frameworktransformers def predictionloop self dataset tf data dataset steps int numexamples int description str predictionlossonly optionalbool none predictionoutput predictionlossonly predictionlossonly if predictionlossonly is not none else self args predictionlossonly logger infof running description logger infof num examples in dataset numexamples if description evaluation logger infof num examples in used in evaluation self args evalbatchsize steps logger infof batch size self args evalbatchsize labelids np ndarray none preds np ndarray none self evalloss resetstates reset the past mems state at the beginning of the evaluation if necessary if self args pastindex 0 self past none for step batch in enumeratedataset logits self distributedpredictionstepsbatch labels batch if not predictionlossonly if isinstancelogits tuple logits logits0 if isinstancelabels tuple labels labels0 if self args nreplicas 1 for val in logits values if preds is none preds val numpy else preds np appendpreds val numpy axis0 for val in labels values if labelids is none labelids val numpy else labelids np appendlabelids val numpy axis0 else if preds is none preds logits numpy else preds np appendpreds logits numpy axis0 if labelids is none labelids labels numpy else labelids np appendlabelids labels numpy axis0 if step steps 1 break if self computemetrics is not none and preds is not none and labelids is not none metrics self computemetricsevalpredictionpredictionspreds labelidslabelids else metrics metricsevalloss self evalloss result numpy steps for key in listmetrics keys if not key startswitheval metricsfevalkey metrics popkey if self args pastindex and hasattrself past clean the state at the end of training delattrself past return predictionoutputpredictionspreds labelidslabelids metricsmetrics def logself logs dictstr float none logsepoch self epochlogging if self tbwriter with self tbwriter asdefault for k v in logs items tf summary scalark v stepself globalstep self tbwriter flush if iswandbavailable wandb loglogs stepself globalstep if iscometavailable experiment cometml config getglobalexperiment if experiment is not none experiment logmetrics logs stepself globalstep epochself epochlogging frameworktransformers output logs step self globalstep logger infooutput def evaluateself evaldataset optionaltf data dataset none dictstr float evalds steps numexamples self getevaltfdatasetevaldataset output self predictionloopevalds steps numexamples descriptionevaluation logs output metrics logsepoch self epochlogging self loglogs return output metrics def predictionstep self features tf tensor labels tf tensor nbinstancesinglobalbatch tf tensor tf tensor perexampleloss logits self runmodelfeatures labels false scaledloss perexampleloss tf castnbinstancesinglobalbatch dtypeperexampleloss dtype self evalloss updatestatescaledloss return logits tf function def distributedpredictionstepsself batch nbinstancesinbatch self computenbinstancesbatch inputs self getstepinputsbatch nbinstancesinbatch logits self args strategy runself predictionstep inputs return logits def trainself none trainds self gettraintfdataset if self args debug tf summary traceongraphtrue profilertrue self gradientaccumulator reset numupdatestepsperepoch self numtrainexamples self totaltrainbatchsize in fact self args dataloaderdroplast has no effect in trainertf py because the dataset is repeated before being batched it has the effect only when tpu is used which requires explicit tensor shape in order to make the gradient accumulation implementation work approx math floor if self args dataloaderdroplast else math ceil numupdatestepsperepoch approxnumupdatestepsperepoch at least one update for each epoch numupdatestepsperepoch maxnumupdatestepsperepoch 1 self stepsperepoch numupdatestepsperepoch if self args maxsteps 0 ttotal self args maxsteps epochs self args maxsteps self stepsperepoch int self args maxsteps self stepsperepoch 0 else ttotal self stepsperepoch self args numtrainepochs epochs self args numtrainepochs since self args numtrainepochs can be float we make epochs be a float always epochs floatepochs with self args strategy scope self createoptimizerandschedulernumtrainingstepsttotal folder os path joinself args outputdir prefixcheckpointdir ckpt tf train checkpointoptimizerself optimizer modelself model self model ckptmanager tf train checkpointmanagerckpt folder maxtokeepself args savetotallimit iterations self optimizer iterations epochstrained 0 stepstrainedincurrentepoch 0 if self model ckptmanager latestcheckpoint logger info fcheckpoint file self model ckptmanager latestcheckpoint found and restoring from checkpoint ckpt restoreself model ckptmanager latestcheckpoint expectpartial self globalstep iterations numpy epochstrained self globalstep self stepsperepoch stepstrainedincurrentepoch self globalstep self stepsperepoch logger info continuing training from checkpoint will skip to saved globalstep logger infof continuing training from epoch epochstrained logger infof continuing training from global step self globalstep logger infof will skip the first stepstrainedincurrentepoch steps in the first epoch tf summary experimental setstepself globalstep with self tbwriter asdefault tf summary textargs self args tojsonstring self tbwriter flush logger info running training logger infof num examples self numtrainexamples todo we might want to print a more precise epochs if self args maxsteps 0 logger infof num epochs epochs logger infof instantaneous batch size per device self args perdevicetrainbatchsize logger info f total train batch size w parallel distributed accumulation self totaltrainbatchsize logger infof gradient accumulation steps self args gradientaccumulationsteps logger infof steps per epoch self stepsperepoch logger infof total optimization steps ttotal self trainloss tf keras metrics sum starttime datetime datetime now for epochiter in rangeepochstrained intepochs reset the past mems state at the beginning of each epoch if necessary if self args pastindex 0 self past none for step batch in enumeratetrainds skip past any already trained steps if resuming training if stepstrainedincurrentepoch 0 stepstrainedincurrentepoch 1 continue self distributedtrainingstepsbatch self globalstep iterations numpy self epochlogging epochiter step 1 self stepsperepoch trainingloss self trainloss result step 1 if self args debug logs logsloss trainingloss numpy logsepoch self epochlogging self loglogs if self globalstep 1 and self args debug with self tbwriter asdefault tf summary traceexport nametraining stepself globalstep profileroutdirself args loggingdir if self args evalsteps 0 and self args evaluationstrategy intervalstrategy steps and self globalstep self args evalsteps 0 self evaluate if self args loggingsteps 0 and self globalstep self args loggingsteps 0 or self globalstep 1 and self args loggingfirststep logs logsloss trainingloss numpy logslearningrate self lrschedulerself globalstep numpy logsepoch self epochlogging self loglogs if self args savesteps 0 and self globalstep self args savesteps 0 ckptsavepath self model ckptmanager save logger infofsaving checkpoint for step self globalstep at ckptsavepath if self args maxsteps 0 and self globalstep ttotal break if self globalstep self stepsperepoch 0 break self trainloss resetstates if self args maxsteps 0 and self globalstep self args maxsteps break endtime datetime datetime now logger infoftraining took strendtime starttime if self args pastindex and hasattrself past clean the state at the end of training delattrself past def trainingstepself features labels nbinstancesinglobalbatch perexampleloss self runmodelfeatures labels true scaledloss perexampleloss tf castnbinstancesinglobalbatch dtypeperexampleloss dtype gradients tf gradientsscaledloss self model trainablevariables gradients g if g is not none else tf zeroslikev for g v in zipgradients self model trainablevariables if self args gradientaccumulationsteps 1 self gradientaccumulatorgradients self trainloss updatestatescaledloss if self args gradientaccumulationsteps 1 return gradients def applygradientsself features labels nbinstancesinglobalbatch if self args gradientaccumulationsteps 1 gradients self trainingstepfeatures labels nbinstancesinglobalbatch self optimizer applygradientslistzipgradients self model trainablevariables else for in tf rangeself args gradientaccumulationsteps reducedfeatures k ft self args trainbatchsize self args nreplicas for k ft in features items if tf istensorlabels reducedlabels labels self args trainbatchsize self args nreplicas elif isinstancelabels dict reducedlabels k lbl self args trainbatchsize self args nreplicas for k lbl in labels items else raise valueerrorthe labels must be either a tf tensor or a dict self trainingstepreducedfeatures reducedlabels nbinstancesinglobalbatch features k tf concat ftself args trainbatchsize self args nreplicas reducedfeaturesk axis0 for k ft in features items if tf istensorlabels labels tf concat labelsself args trainbatchsize self args nreplicas reducedlabels axis0 elif isinstancelabels dict labels k tf concat lblself args trainbatchsize self args nreplicas reducedlabelsk axis0 for k lbl in labels items else raise valueerrorthe labels must be either a tf tensor or a dict gradients self gradientaccumulator gradients gradients tf clipbyvaluegrad self args maxgradnorm self args maxgradnorm for grad in gradients self optimizer applygradientslistzipgradients self model trainablevariables self gradientaccumulator reset tf function def distributedtrainingstepsself batch with self args strategy scope nbinstancesinbatch self computenbinstancesbatch inputs self getstepinputsbatch nbinstancesinbatch self args strategy runself applygradients inputs staticmethod def computenbinstancesbatch labels batch1 if isinstancelabels perreplica labels tf concatlabels values axis0 nbinstances tf reducesumtf castlabels 100 dtypetf int32 return nbinstances staticmethod def getstepinputsbatch nbinstances features labels batch if isinstancelabels perreplica need to make a perreplica objects for nbinstances nbinstances perreplicanbinstances lenlabels values stepinputs features labels nbinstances return stepinputs def runmodelself features labels training if self args pastindex 0 and getattrself past none is not none featuresmems self past if isinstancelabels dict outputs self modelfeatures trainingtraining labels 2 else outputs self modelfeatures labelslabels trainingtraining 2 loss logits outputs 2 if self args pastindex 0 self past outputsself args pastindex return loss logits def predictself testdataset tf data dataset predictionoutput testds steps numexamples self gettesttfdatasettestdataset return self predictionlooptestds steps numexamples descriptionprediction def savemodelself outputdir optionalstr none outputdir outputdir if outputdir is not none else self args outputdir logger infofsaving model in outputdir if not isinstanceself model tfpretrainedmodel raise valueerrortrainer model appears to not be a pretrainedmodel self model savepretrainedoutputdir 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 tensorflow trainer class integrations must be imported before ml frameworks isort off isort on tftrainer is a simple but feature complete training and eval loop for tensorflow optimized for transformers args model tfpretrainedmodel the model to train evaluate or use for predictions args tftrainingarguments the arguments to tweak training train_dataset tf data dataset optional the dataset to use for training the dataset should yield tuples of features labels where features is a dict of input features and labels is the labels if labels is a tensor the loss is calculated by the model by calling model features labels labels if labels is a dict such as when using a questionanswering head model with multiple targets the loss is instead calculated by calling model features labels eval_dataset tf data dataset optional the dataset to use for evaluation the dataset should yield tuples of features labels where features is a dict of input features and labels is the labels if labels is a tensor the loss is calculated by the model by calling model features labels labels if labels is a dict such as when using a questionanswering head model with multiple targets the loss is instead calculated by calling model features labels compute_metrics callable evalprediction dict optional the function that will be used to compute metrics at evaluation must take a evalprediction and return a dictionary string to metric values tb_writer tf summary summarywriter optional object to write to tensorboard optimizers tuple tf keras optimizers optimizer tf keras optimizers schedules learningrateschedule optional a tuple containing the optimizer and the scheduler to use the optimizer default to an instance of tf keras optimizers adam if args weight_decay_rate is 0 else an instance of adamweightdecay the scheduler will default to an instance of tf keras optimizers schedules polynomialdecay if args num_warmup_steps is 0 else an instance of warmup returns the training tf data dataset subclass and override this method if you want to inject some custom behavior returns the evaluation tf data dataset args eval_dataset tf data dataset optional if provided will override self eval_dataset the dataset should yield tuples of features labels where features is a dict of input features and labels is the labels if labels is a tensor the loss is calculated by the model by calling model features labels labels if labels is a dict such as when using a questionanswering head model with multiple targets the loss is instead calculated by calling model features labels subclass and override this method if you want to inject some custom behavior returns a test tf data dataset args test_dataset tf data dataset the dataset to use the dataset should yield tuples of features labels where features is a dict of input features and labels is the labels if labels is a tensor the loss is calculated by the model by calling model features labels labels if labels is a dict such as when using a questionanswering head model with multiple targets the loss is instead calculated by calling model features labels subclass and override this method if you want to inject some custom behavior setup the optimizer and the learning rate scheduler we provide a reasonable default that works well if you want to use something else you can pass a tuple in the tftrainer s init through optimizers or subclass and override this method setup the optional weights biases wandb integration one can subclass and override this method to customize the setup if needed find more information here https docs wandb com huggingface __ you can also override the following environment variables environment wandb_project optional str huggingface by default set this to a custom string to store results in a different project wandb_disabled optional boolean defaults to false set to true to disable wandb entirely setup the optional comet ml integration environment comet_mode optional str offline online or disabled comet_project_name optional str comet ml project name for experiments comet_offline_directory optional str folder to use for saving offline experiments when comet_mode is offline for a number of configurable items in the environment see here https www comet ml docs python sdk advanced comet configuration variables __ prediction evaluation loop shared by tftrainer evaluate and tftrainer predict works both with or without labels reset the past mems state at the beginning of the evaluation if necessary clean the state at the end of training log logs on the various objects watching training subclass and override this method to inject custom behavior args logs dict str float the values to log run evaluation and returns metrics the calling script will be responsible for providing a method to compute metrics as they are task dependent pass it to the init compute_metrics argument args eval_dataset tf data dataset optional pass a dataset if you wish to override self eval_dataset the dataset should yield tuples of features labels where features is a dict of input features and labels is the labels if labels is a tensor the loss is calculated by the model by calling model features labels labels if labels is a dict such as when using a questionanswering head model with multiple targets the loss is instead calculated by calling model features labels returns a dictionary containing the evaluation loss and the potential metrics computed from the predictions compute the prediction on features and update the loss with labels subclass and override to inject some custom behavior train method to train the model in fact self args dataloader_drop_last has no effect in trainer_tf py because the dataset is repeated before being batched it has the effect only when tpu is used which requires explicit tensor shape in order to make the gradient accumulation implementation work at least one update for each epoch since self args num_train_epochs can be float we make epochs be a float always todo we might want to print a more precise epochs if self args max_steps 0 reset the past mems state at the beginning of each epoch if necessary skip past any already trained steps if resuming training clean the state at the end of training perform a training step on features and labels subclass and override to inject some custom behavior need to make a perreplica objects for nb_instances computes the loss of the given features and labels pair subclass and override this method if you want to inject some custom behavior args features tf tensor a batch of input features labels tf tensor a batch of labels training bool whether or not to run the model in training mode returns a tuple of two tf tensor the loss and logits run prediction and returns predictions and potential metrics depending on the dataset and your use case your test dataset may contain labels in that case this method will also return metrics like in evaluate args test_dataset tf data dataset dataset to run the predictions on the dataset should yield tuples of features labels where features is a dict of input features and labels is the labels if labels is a tensor the loss is calculated by the model by calling model features labels labels if labels is a dict such as when using a questionanswering head model with multiple targets the loss is instead calculated by calling model features labels returns namedtuple a namedtuple with the following keys predictions np ndarray the predictions on test_dataset label_ids np ndarray optional the labels if the dataset contained some metrics dict str float optional the potential dictionary of metrics if the dataset contained labels will save the model so you can reload it using from_pretrained
import datetime
import math
import os
import warnings
from typing import Callable, Dict, Optional, Tuple

from .utils import ENV_VARS_TRUE_VALUES




from .integrations import (
    is_comet_available,
    is_wandb_available,
)



import numpy as np
import tensorflow as tf
from tensorflow.python.distribute.values import PerReplica

from .modeling_tf_utils import TFPreTrainedModel
from .optimization_tf import GradientAccumulator, create_optimizer
from .trainer_utils import (
    PREFIX_CHECKPOINT_DIR,
    EvalPrediction,
    IntervalStrategy,
    PredictionOutput,
    enable_full_determinism,
    set_seed,
)
from .training_args_tf import TFTrainingArguments
from .utils import logging


if is_wandb_available():
    import wandb

if is_comet_available():
    import comet_ml

logger = logging.get_logger(__name__)


class TFTrainer:
    

    def __init__(
        self,
        model: TFPreTrainedModel,
        args: TFTrainingArguments,
        train_dataset: Optional[tf.data.Dataset] = None,
        eval_dataset: Optional[tf.data.Dataset] = None,
        compute_metrics: Optional[Callable[[EvalPrediction], Dict]] = None,
        tb_writer: Optional[tf.summary.SummaryWriter] = None,
        optimizers: Tuple[tf.keras.optimizers.Optimizer, tf.keras.optimizers.schedules.LearningRateSchedule] = (
            None,
            None,
        ),
    ):
        self.model = model
        self.args = args
        self.train_dataset = train_dataset
        self.eval_dataset = eval_dataset
        self.compute_metrics = compute_metrics
        self.optimizer, self.lr_scheduler = optimizers
        self.gradient_accumulator = GradientAccumulator()
        self.global_step = 0
        self.epoch_logging = 0
        self.eval_loss = tf.keras.metrics.Sum()

        warnings.warn(
            "The class `TFTrainer` is deprecated and will be removed in version 5 of Transformers. "
            "We recommend using native Keras instead, by calling methods like `fit()` and `predict()` "
            "directly on the model object. Detailed examples of the Keras style can be found in our "
            "examples at https://github.com/huggingface/transformers/tree/main/examples/tensorflow",
            FutureWarning,
        )

        if tb_writer is not None:
            self.tb_writer = tb_writer
        else:
            self.tb_writer = tf.summary.create_file_writer(self.args.logging_dir)

        if is_wandb_available():
            self.setup_wandb()
        elif os.getenv("WANDB_DISABLED", "").upper() not in ENV_VARS_TRUE_VALUES:
            logger.info(
                "You are instantiating a Trainer but W&B is not installed. To use wandb logging, "
                "run `pip install wandb && wandb login` see https://docs.wandb.com/huggingface."
            )

        if is_comet_available():
            self.setup_comet()
        elif os.environ.get("COMET_MODE") != "DISABLED":
            logger.info(
                "To use comet_ml logging, run `pip/conda install comet_ml` "
                "see https://www.comet.ml/docs/python-sdk/huggingface/"
            )

        enable_full_determinism(self.args.seed) if self.args.full_determinism else set_seed(self.args.seed)

    def get_train_tfdataset(self) -> tf.data.Dataset:
        
        if self.train_dataset is None:
            raise ValueError("Trainer: training requires a train_dataset.")

        self.total_train_batch_size = self.args.train_batch_size * self.args.gradient_accumulation_steps
        self.num_train_examples = self.train_dataset.cardinality().numpy()

        if self.num_train_examples < 0:
            raise ValueError("The training dataset must have an asserted cardinality")

        ds = (
            self.train_dataset.repeat()
            .shuffle(self.num_train_examples, seed=self.args.seed)
            .batch(self.total_train_batch_size, drop_remainder=self.args.dataloader_drop_last)
            .prefetch(tf.data.experimental.AUTOTUNE)
        )

        return self.args.strategy.experimental_distribute_dataset(ds)

    def get_eval_tfdataset(self, eval_dataset: Optional[tf.data.Dataset] = None) -> tf.data.Dataset:
        
        if eval_dataset is None and self.eval_dataset is None:
            raise ValueError("Trainer: evaluation requires an eval_dataset.")

        eval_dataset = eval_dataset if eval_dataset is not None else self.eval_dataset
        num_examples = eval_dataset.cardinality().numpy()

        if num_examples < 0:
            raise ValueError("The training dataset must have an asserted cardinality")

        approx = math.floor if self.args.dataloader_drop_last else math.ceil
        steps = approx(num_examples / self.args.eval_batch_size)
        ds = (
            eval_dataset.repeat()
            .batch(self.args.eval_batch_size, drop_remainder=self.args.dataloader_drop_last)
            .prefetch(tf.data.experimental.AUTOTUNE)
        )

        return self.args.strategy.experimental_distribute_dataset(ds), steps, num_examples

    def get_test_tfdataset(self, test_dataset: tf.data.Dataset) -> tf.data.Dataset:
        

        num_examples = test_dataset.cardinality().numpy()

        if num_examples < 0:
            raise ValueError("The training dataset must have an asserted cardinality")

        steps = math.ceil(num_examples / self.args.eval_batch_size)
        ds = test_dataset.batch(self.args.eval_batch_size).prefetch(tf.data.experimental.AUTOTUNE)

        return self.args.strategy.experimental_distribute_dataset(ds), steps, num_examples

    def create_optimizer_and_scheduler(self, num_training_steps: int):
        
        if not self.optimizer and not self.lr_scheduler:
            warmup_steps = (
                self.args.warmup_steps
                if self.args.warmup_steps > 0
                else math.ceil(num_training_steps * self.args.warmup_ratio)
            )

            self.optimizer, self.lr_scheduler = create_optimizer(
                self.args.learning_rate,
                num_training_steps,
                warmup_steps,
                adam_beta1=self.args.adam_beta1,
                adam_beta2=self.args.adam_beta2,
                adam_epsilon=self.args.adam_epsilon,
                weight_decay_rate=self.args.weight_decay,
                power=self.args.poly_power,
            )

    def setup_wandb(self):
        

        logger.info('Automatic Weights & Biases logging enabled, to disable set os.environ["WANDB_DISABLED"] = "true"')
        combined_dict = {**self.model.config.to_dict(), **self.args.to_sanitized_dict()}
        wandb.init(project=os.getenv("WANDB_PROJECT", "huggingface"), config=combined_dict, name=self.args.run_name)

    def setup_comet(self):
        
        comet_mode = os.getenv("COMET_MODE", "ONLINE").upper()
        args = {"project_name": os.getenv("COMET_PROJECT_NAME", "huggingface")}
        experiment = None
        if comet_mode == "ONLINE":
            experiment = comet_ml.Experiment(**args)
            logger.info("Automatic Comet.ml online logging enabled")
        elif comet_mode == "OFFLINE":
            args["offline_directory"] = os.getenv("COMET_OFFLINE_DIRECTORY", "./")
            experiment = comet_ml.OfflineExperiment(**args)
            logger.info("Automatic Comet.ml offline logging enabled; use `comet upload` when finished")
        if experiment is not None:
            experiment._set_model_graph(self.model, framework="transformers")
            experiment._log_parameters(self.args, prefix="args/", framework="transformers")
            experiment._log_parameters(self.model.config, prefix="config/", framework="transformers")

    def prediction_loop(
        self,
        dataset: tf.data.Dataset,
        steps: int,
        num_examples: int,
        description: str,
        prediction_loss_only: Optional[bool] = None,
    ) -> PredictionOutput:
        

        prediction_loss_only = (
            prediction_loss_only if prediction_loss_only is not None else self.args.prediction_loss_only
        )

        logger.info(f"***** Running {description} *****")
        logger.info(f"  Num examples in dataset = {num_examples}")
        if description == "Evaluation":
            logger.info(f"  Num examples in used in evaluation = {self.args.eval_batch_size * steps}")
        logger.info(f"  Batch size = {self.args.eval_batch_size}")

        label_ids: np.ndarray = None
        preds: np.ndarray = None
        self.eval_loss.reset_states()

        
        if self.args.past_index >= 0:
            self._past = None

        for step, batch in enumerate(dataset):
            logits = self.distributed_prediction_steps(batch)
            _, labels = batch

            if not prediction_loss_only:
                if isinstance(logits, tuple):
                    logits = logits[0]

                if isinstance(labels, tuple):
                    labels = labels[0]

                if self.args.n_replicas > 1:
                    for val in logits.values:
                        if preds is None:
                            preds = val.numpy()
                        else:
                            preds = np.append(preds, val.numpy(), axis=0)

                    for val in labels.values:
                        if label_ids is None:
                            label_ids = val.numpy()
                        else:
                            label_ids = np.append(label_ids, val.numpy(), axis=0)
                else:
                    if preds is None:
                        preds = logits.numpy()
                    else:
                        preds = np.append(preds, logits.numpy(), axis=0)

                    if label_ids is None:
                        label_ids = labels.numpy()
                    else:
                        label_ids = np.append(label_ids, labels.numpy(), axis=0)

                if step == steps - 1:
                    break

        if self.compute_metrics is not None and preds is not None and label_ids is not None:
            metrics = self.compute_metrics(EvalPrediction(predictions=preds, label_ids=label_ids))
        else:
            metrics = {}

        metrics["eval_loss"] = self.eval_loss.result().numpy() / steps

        for key in list(metrics.keys()):
            if not key.startswith("eval_"):
                metrics[f"eval_{key}"] = metrics.pop(key)

        if self.args.past_index and hasattr(self, "_past"):
            
            delattr(self, "_past")

        return PredictionOutput(predictions=preds, label_ids=label_ids, metrics=metrics)

    def log(self, logs: Dict[str, float]) -> None:
        
        logs["epoch"] = self.epoch_logging

        if self.tb_writer:
            with self.tb_writer.as_default():
                for k, v in logs.items():
                    tf.summary.scalar(k, v, step=self.global_step)
            self.tb_writer.flush()

        if is_wandb_available():
            wandb.log(logs, step=self.global_step)

        if is_comet_available():
            experiment = comet_ml.config.get_global_experiment()
            if experiment is not None:
                experiment._log_metrics(
                    logs, step=self.global_step, epoch=self.epoch_logging, framework="transformers"
                )

        output = {**logs, **{"step": self.global_step}}

        logger.info(output)

    def evaluate(self, eval_dataset: Optional[tf.data.Dataset] = None) -> Dict[str, float]:
        
        eval_ds, steps, num_examples = self.get_eval_tfdataset(eval_dataset)

        output = self.prediction_loop(eval_ds, steps, num_examples, description="Evaluation")
        logs = {**output.metrics}
        logs["epoch"] = self.epoch_logging

        self.log(logs)

        return output.metrics

    def prediction_step(
        self, features: tf.Tensor, labels: tf.Tensor, nb_instances_in_global_batch: tf.Tensor
    ) -> tf.Tensor:
        
        per_example_loss, logits = self.run_model(features, labels, False)
        scaled_loss = per_example_loss / tf.cast(nb_instances_in_global_batch, dtype=per_example_loss.dtype)

        self.eval_loss.update_state(scaled_loss)

        return logits

    @tf.function
    def distributed_prediction_steps(self, batch):
        nb_instances_in_batch = self._compute_nb_instances(batch)
        inputs = self._get_step_inputs(batch, nb_instances_in_batch)

        logits = self.args.strategy.run(self.prediction_step, inputs)

        return logits

    def train(self) -> None:
        
        train_ds = self.get_train_tfdataset()

        if self.args.debug:
            tf.summary.trace_on(graph=True, profiler=True)

        self.gradient_accumulator.reset()

        num_update_steps_per_epoch = self.num_train_examples / self.total_train_batch_size

        
        
        
        
        approx = math.floor if self.args.dataloader_drop_last else math.ceil
        num_update_steps_per_epoch = approx(num_update_steps_per_epoch)

        
        num_update_steps_per_epoch = max(num_update_steps_per_epoch, 1)
        self.steps_per_epoch = num_update_steps_per_epoch

        if self.args.max_steps > 0:
            t_total = self.args.max_steps
            epochs = (self.args.max_steps // self.steps_per_epoch) + int(
                self.args.max_steps % self.steps_per_epoch > 0
            )
        else:
            t_total = self.steps_per_epoch * self.args.num_train_epochs
            epochs = self.args.num_train_epochs

        
        epochs = float(epochs)

        with self.args.strategy.scope():
            self.create_optimizer_and_scheduler(num_training_steps=t_total)
            folder = os.path.join(self.args.output_dir, PREFIX_CHECKPOINT_DIR)
            ckpt = tf.train.Checkpoint(optimizer=self.optimizer, model=self.model)
            self.model.ckpt_manager = tf.train.CheckpointManager(ckpt, folder, max_to_keep=self.args.save_total_limit)

            iterations = self.optimizer.iterations
            epochs_trained = 0
            steps_trained_in_current_epoch = 0
            if self.model.ckpt_manager.latest_checkpoint:
                logger.info(
                    f"Checkpoint file {self.model.ckpt_manager.latest_checkpoint} found and restoring from checkpoint"
                )
                ckpt.restore(self.model.ckpt_manager.latest_checkpoint).expect_partial()

                self.global_step = iterations.numpy()

                epochs_trained = self.global_step // self.steps_per_epoch
                steps_trained_in_current_epoch = self.global_step % self.steps_per_epoch

                logger.info("  Continuing training from checkpoint, will skip to saved global_step")
                logger.info(f"  Continuing training from epoch {epochs_trained}")
                logger.info(f"  Continuing training from global step {self.global_step}")
                logger.info(f"  Will skip the first {steps_trained_in_current_epoch} steps in the first epoch")

            tf.summary.experimental.set_step(self.global_step)

            with self.tb_writer.as_default():
                tf.summary.text("args", self.args.to_json_string())

            self.tb_writer.flush()

            logger.info("***** Running training *****")
            logger.info(f"  Num examples = {self.num_train_examples}")
            
            logger.info(f"  Num Epochs = {epochs}")
            logger.info(f"  Instantaneous batch size per device = {self.args.per_device_train_batch_size}")
            logger.info(
                f"  Total train batch size (w. parallel, distributed & accumulation) = {self.total_train_batch_size}"
            )
            logger.info(f"  Gradient Accumulation steps = {self.args.gradient_accumulation_steps}")
            logger.info(f"  Steps per epoch = {self.steps_per_epoch}")
            logger.info(f"  Total optimization steps = {t_total}")

            self.train_loss = tf.keras.metrics.Sum()
            start_time = datetime.datetime.now()

            for epoch_iter in range(epochs_trained, int(epochs)):
                
                if self.args.past_index >= 0:
                    self._past = None

                for step, batch in enumerate(train_ds):
                    
                    if steps_trained_in_current_epoch > 0:
                        steps_trained_in_current_epoch -= 1
                        continue

                    self.distributed_training_steps(batch)

                    self.global_step = iterations.numpy()
                    self.epoch_logging = epoch_iter + (step + 1) / self.steps_per_epoch

                    training_loss = self.train_loss.result() / (step + 1)

                    if self.args.debug:
                        logs = {}
                        logs["loss"] = training_loss.numpy()
                        logs["epoch"] = self.epoch_logging

                        self.log(logs)

                    if self.global_step == 1 and self.args.debug:
                        with self.tb_writer.as_default():
                            tf.summary.trace_export(
                                name="training", step=self.global_step, profiler_outdir=self.args.logging_dir
                            )

                    if (
                        self.args.eval_steps > 0
                        and self.args.evaluation_strategy == IntervalStrategy.STEPS
                        and self.global_step % self.args.eval_steps == 0
                    ):
                        self.evaluate()

                    if (self.args.logging_steps > 0 and self.global_step % self.args.logging_steps == 0) or (
                        self.global_step == 1 and self.args.logging_first_step
                    ):
                        logs = {}
                        logs["loss"] = training_loss.numpy()
                        logs["learning_rate"] = self.lr_scheduler(self.global_step).numpy()
                        logs["epoch"] = self.epoch_logging

                        self.log(logs)

                    if self.args.save_steps > 0 and self.global_step % self.args.save_steps == 0:
                        ckpt_save_path = self.model.ckpt_manager.save()

                        logger.info(f"Saving checkpoint for step {self.global_step} at {ckpt_save_path}")

                    if self.args.max_steps > 0 and self.global_step >= t_total:
                        break

                    if self.global_step % self.steps_per_epoch == 0:
                        break

                self.train_loss.reset_states()

                if self.args.max_steps > 0 and self.global_step >= self.args.max_steps:
                    break

            end_time = datetime.datetime.now()

            logger.info(f"Training took: {str(end_time - start_time)}")

        if self.args.past_index and hasattr(self, "_past"):
            
            delattr(self, "_past")

    def training_step(self, features, labels, nb_instances_in_global_batch):
        
        per_example_loss, _ = self.run_model(features, labels, True)
        scaled_loss = per_example_loss / tf.cast(nb_instances_in_global_batch, dtype=per_example_loss.dtype)
        gradients = tf.gradients(scaled_loss, self.model.trainable_variables)
        gradients = [
            g if g is not None else tf.zeros_like(v) for g, v in zip(gradients, self.model.trainable_variables)
        ]

        if self.args.gradient_accumulation_steps > 1:
            self.gradient_accumulator(gradients)

        self.train_loss.update_state(scaled_loss)

        if self.args.gradient_accumulation_steps == 1:
            return gradients

    def apply_gradients(self, features, labels, nb_instances_in_global_batch):
        if self.args.gradient_accumulation_steps == 1:
            gradients = self.training_step(features, labels, nb_instances_in_global_batch)

            self.optimizer.apply_gradients(list(zip(gradients, self.model.trainable_variables)))
        else:
            for _ in tf.range(self.args.gradient_accumulation_steps):
                reduced_features = {
                    k: ft[: self.args.train_batch_size // self.args.n_replicas] for k, ft in features.items()
                }

                if tf.is_tensor(labels):
                    reduced_labels = labels[: self.args.train_batch_size // self.args.n_replicas]
                elif isinstance(labels, dict):
                    reduced_labels = {
                        k: lbl[: self.args.train_batch_size // self.args.n_replicas] for k, lbl in labels.items()
                    }
                else:
                    raise ValueError("The labels must be either a tf.Tensor or a dict.")

                self.training_step(reduced_features, reduced_labels, nb_instances_in_global_batch)

                features = {
                    k: tf.concat(
                        [ft[self.args.train_batch_size // self.args.n_replicas :], reduced_features[k]],
                        axis=0,
                    )
                    for k, ft in features.items()
                }

                if tf.is_tensor(labels):
                    labels = tf.concat(
                        [labels[self.args.train_batch_size // self.args.n_replicas :], reduced_labels], axis=0
                    )
                elif isinstance(labels, dict):
                    labels = {
                        k: tf.concat(
                            [lbl[self.args.train_batch_size // self.args.n_replicas :], reduced_labels[k]],
                            axis=0,
                        )
                        for k, lbl in labels.items()
                    }
                else:
                    raise ValueError("The labels must be either a tf.Tensor or a dict.")

            gradients = self.gradient_accumulator.gradients
            gradients = [
                (tf.clip_by_value(grad, -self.args.max_grad_norm, self.args.max_grad_norm)) for grad in gradients
            ]

            self.optimizer.apply_gradients(list(zip(gradients, self.model.trainable_variables)))
            self.gradient_accumulator.reset()

    @tf.function
    def distributed_training_steps(self, batch):
        with self.args.strategy.scope():
            nb_instances_in_batch = self._compute_nb_instances(batch)
            inputs = self._get_step_inputs(batch, nb_instances_in_batch)

            self.args.strategy.run(self.apply_gradients, inputs)

    @staticmethod
    def _compute_nb_instances(batch):
        labels = batch[-1]
        if isinstance(labels, PerReplica):
            labels = tf.concat(labels.values, axis=0)

        nb_instances = tf.reduce_sum(tf.cast(labels != -100, dtype=tf.int32))

        return nb_instances

    @staticmethod
    def _get_step_inputs(batch, nb_instances):
        features, labels = batch

        if isinstance(labels, PerReplica):
            
            nb_instances = PerReplica([nb_instances] * len(labels.values))

        step_inputs = (features, labels, nb_instances)

        return step_inputs

    def run_model(self, features, labels, training):
        

        if self.args.past_index >= 0 and getattr(self, "_past", None) is not None:
            features["mems"] = self._past

        if isinstance(labels, (dict)):
            outputs = self.model(features, training=training, **labels)[:2]
        else:
            outputs = self.model(features, labels=labels, training=training)[:2]

        loss, logits = outputs[:2]

        if self.args.past_index >= 0:
            self._past = outputs[self.args.past_index]

        return loss, logits

    def predict(self, test_dataset: tf.data.Dataset) -> PredictionOutput:
        
        test_ds, steps, num_examples = self.get_test_tfdataset(test_dataset)

        return self.prediction_loop(test_ds, steps, num_examples, description="Prediction")

    def save_model(self, output_dir: Optional[str] = None):
        
        output_dir = output_dir if output_dir is not None else self.args.output_dir

        logger.info(f"Saving model in {output_dir}")

        if not isinstance(self.model, TFPreTrainedModel):
            raise ValueError("Trainer.model appears to not be a PreTrainedModel")

        self.model.save_pretrained(output_dir)