codingutf8 2020present the huggingface inc 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 callbacks to use with the trainer class and customize the training loop a class containing the trainer inner state that will be saved along the model and optimizer when checkpointing and passed to the trainercallback tip in all this class one step is to be understood as one update step when using gradient accumulation one update step may require several forward and backward passes if you use gradientaccumulationstepsn then one update step requires going through n batches tip args epoch float optional only set during training will represent the epoch the training is at the decimal part being the percentage of the current epoch completed globalstep int optional defaults to 0 during training represents the number of update steps completed maxsteps int optional defaults to 0 the number of update steps to do during the current training loggingsteps int optional defaults to 500 log every x updates steps evalsteps int optional run an evaluation every x steps savesteps int optional defaults to 500 save checkpoint every x updates steps numinputtokensseen int optional defaults to 0 the number of tokens seen during training number of input tokens not the number of prediction tokens totalflos float optional defaults to 0 the total number of floating operations done by the model since the beginning of training stored as floats to avoid overflow loghistory listdictstr float optional the list of logs done since the beginning of training bestmetric float optional when tracking the best model the value of the best metric encountered so far bestmodelcheckpoint str optional when tracking the best model the value of the name of the checkpoint for the best model encountered so far islocalprocesszero bool optional defaults to true whether or not this process is the local e g on one machine if training in a distributed fashion on several machines main process isworldprocesszero bool optional defaults to true whether or not this process is the global main process when training in a distributed fashion on several machines this is only going to be true for one process ishyperparamsearch bool optional defaults to false whether we are in the process of a hyper parameter search using trainer hyperparametersearch this will impact the way data will be logged in tensorboard save the content of this instance in json format inside jsonpath jsonstring json dumpsdataclasses asdictself indent2 sortkeystrue n with openjsonpath w encodingutf8 as f f writejsonstring classmethod def loadfromjsoncls jsonpath str a class that handles the trainer control flow this class is used by the trainercallback to activate some switches in the training loop args shouldtrainingstop bool optional defaults to false whether or not the training should be interrupted if true this variable will not be set back to false the training will just stop shouldepochstop bool optional defaults to false whether or not the current epoch should be interrupted if true this variable will be set back to false at the beginning of the next epoch shouldsave bool optional defaults to false whether or not the model should be saved at this step if true this variable will be set back to false at the beginning of the next step shouldevaluate bool optional defaults to false whether or not the model should be evaluated at this step if true this variable will be set back to false at the beginning of the next step shouldlog bool optional defaults to false whether or not the logs should be reported at this step if true this variable will be set back to false at the beginning of the next step internal method that resets the variable for a new training self shouldtrainingstop false def newepochself internal method that resets the variable for a new step self shouldsave false self shouldevaluate false self shouldlog false class trainercallback noformat event called at the end of the initialization of the trainer event called at the beginning of training event called at the end of training event called at the beginning of an epoch event called at the end of an epoch event called at the beginning of a training step if using gradient accumulation one training step might take several inputs event called at the end of an substep during gradient accumulation event called at the end of a training step if using gradient accumulation one training step might take several inputs event called after an evaluation phase event called after a successful prediction event called after a checkpoint save event called after logging the last logs event called after a prediction step internal class that just calls the list of callbacks in order def initself callbacks model tokenizer optimizer lrscheduler self callbacks for cb in callbacks self addcallbackcb self model model self tokenizer tokenizer self optimizer optimizer self lrscheduler lrscheduler self traindataloader none self evaldataloader none if not anyisinstancecb defaultflowcallback for cb in self callbacks logger warning the trainer will not work properly if you don t have a defaultflowcallback in its callbacks youn should add one before training with trainer addcallbackdefaultflowcallback the current list of callbacks isn self callbacklist def addcallbackself callback cb callback if isinstancecallback type else callback cbclass callback if isinstancecallback type else callback class if cbclass in c class for c in self callbacks logger warning fyou are adding a cbclass to the callbacks of this trainer but there is already one the current list of callbacks isn self callbacklist self callbacks appendcb def popcallbackself callback if isinstancecallback type for cb in self callbacks if isinstancecb callback self callbacks removecb return cb else for cb in self callbacks if cb callback self callbacks removecb return cb def removecallbackself callback if isinstancecallback type for cb in self callbacks if isinstancecb callback self callbacks removecb return else self callbacks removecallback property def callbacklistself return n joincb class name for cb in self callbacks def oninitendself args trainingarguments state trainerstate control trainercontrol return self calleventoninitend args state control def ontrainbeginself args trainingarguments state trainerstate control trainercontrol control shouldtrainingstop false return self calleventontrainbegin args state control def ontrainendself args trainingarguments state trainerstate control trainercontrol return self calleventontrainend args state control def onepochbeginself args trainingarguments state trainerstate control trainercontrol control shouldepochstop false return self calleventonepochbegin args state control def onepochendself args trainingarguments state trainerstate control trainercontrol return self calleventonepochend args state control def onstepbeginself args trainingarguments state trainerstate control trainercontrol control shouldlog false control shouldevaluate false control shouldsave false return self calleventonstepbegin args state control def onsubstependself args trainingarguments state trainerstate control trainercontrol return self calleventonsubstepend args state control def onstependself args trainingarguments state trainerstate control trainercontrol return self calleventonstepend args state control def onevaluateself args trainingarguments state trainerstate control trainercontrol metrics control shouldevaluate false return self calleventonevaluate args state control metricsmetrics def onpredictself args trainingarguments state trainerstate control trainercontrol metrics return self calleventonpredict args state control metricsmetrics def onsaveself args trainingarguments state trainerstate control trainercontrol control shouldsave false return self calleventonsave args state control def onlogself args trainingarguments state trainerstate control trainercontrol logs control shouldlog false return self calleventonlog args state control logslogs def onpredictionstepself args trainingarguments state trainerstate control trainercontrol return self calleventonpredictionstep args state control def calleventself event args state control kwargs for callback in self callbacks result getattrcallback event args state control modelself model tokenizerself tokenizer optimizerself optimizer lrschedulerself lrscheduler traindataloaderself traindataloader evaldataloaderself evaldataloader kwargs a callback can skip the return of control if it doesn t change it if result is not none control result return control class defaultflowcallbacktrainercallback def onstependself args trainingarguments state trainerstate control trainercontrol kwargs log if state globalstep 1 and args loggingfirststep control shouldlog true if args loggingstrategy intervalstrategy steps and state globalstep state loggingsteps 0 control shouldlog true evaluate if args evaluationstrategy intervalstrategy steps and state globalstep state evalsteps 0 and args evaldelay state globalstep control shouldevaluate true save if args savestrategy intervalstrategy steps and state savesteps 0 and state globalstep state savesteps 0 control shouldsave true end training if state globalstep state maxsteps control shouldtrainingstop true return control def onepochendself args trainingarguments state trainerstate control trainercontrol kwargs log if args loggingstrategy intervalstrategy epoch control shouldlog true evaluate if args evaluationstrategy intervalstrategy epoch and args evaldelay state epoch control shouldevaluate true save if args savestrategy intervalstrategy epoch control shouldsave true return control class progresscallbacktrainercallback def initself self trainingbar none self predictionbar none def ontrainbeginself args state control kwargs if state islocalprocesszero self trainingbar tqdmtotalstate maxsteps dynamicncolstrue self currentstep 0 def onstependself args state control kwargs if state islocalprocesszero self trainingbar updatestate globalstep self currentstep self currentstep state globalstep def onpredictionstepself args state control evaldataloadernone kwargs if state islocalprocesszero and haslengthevaldataloader if self predictionbar is none self predictionbar tqdm totallenevaldataloader leaveself trainingbar is none dynamicncolstrue self predictionbar update1 def onevaluateself args state control kwargs if state islocalprocesszero if self predictionbar is not none self predictionbar close self predictionbar none def onpredictself args state control kwargs if state islocalprocesszero if self predictionbar is not none self predictionbar close self predictionbar none def onlogself args state control logsnone kwargs if state islocalprocesszero and self trainingbar is not none logs poptotalflos none self trainingbar writestrlogs def ontrainendself args state control kwargs if state islocalprocesszero self trainingbar close self trainingbar none class printercallbacktrainercallback def onlogself args state control logsnone kwargs logs poptotalflos none if state islocalprocesszero printlogs class earlystoppingcallbacktrainercallback def initself earlystoppingpatience int 1 earlystoppingthreshold optionalfloat 0 0 self earlystoppingpatience earlystoppingpatience self earlystoppingthreshold earlystoppingthreshold earlystoppingpatiencecounter denotes the number of times validation metrics failed to improve self earlystoppingpatiencecounter 0 def checkmetricvalueself args state control metricvalue bestmetric is set by code for loadbestmodel operator np greater if args greaterisbetter else np less if state bestmetric is none or operatormetricvalue state bestmetric and absmetricvalue state bestmetric self earlystoppingthreshold self earlystoppingpatiencecounter 0 else self earlystoppingpatiencecounter 1 def ontrainbeginself args state control kwargs assert args loadbestmodelatend earlystoppingcallback requires loadbestmodelatend true assert args metricforbestmodel is not none earlystoppingcallback requires metricforbestmodel is defined assert args evaluationstrategy intervalstrategy no earlystoppingcallback requires intervalstrategy of steps or epoch def onevaluateself args state control metrics kwargs metrictocheck args metricforbestmodel if not metrictocheck startswitheval metrictocheck fevalmetrictocheck metricvalue metrics getmetrictocheck if metricvalue is none logger warning fearly stopping required metricforbestmodel but did not find metrictocheck so early stopping is disabled return self checkmetricvalueargs state control metricvalue if self earlystoppingpatiencecounter self earlystoppingpatience control shouldtrainingstop true coding utf 8 2020 present the huggingface inc 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 callbacks to use with the trainer class and customize the training loop a class containing the trainer inner state that will be saved along the model and optimizer when checkpointing and passed to the trainercallback tip in all this class one step is to be understood as one update step when using gradient accumulation one update step may require several forward and backward passes if you use gradient_accumulation_steps n then one update step requires going through n batches tip args epoch float optional only set during training will represent the epoch the training is at the decimal part being the percentage of the current epoch completed global_step int optional defaults to 0 during training represents the number of update steps completed max_steps int optional defaults to 0 the number of update steps to do during the current training logging_steps int optional defaults to 500 log every x updates steps eval_steps int optional run an evaluation every x steps save_steps int optional defaults to 500 save checkpoint every x updates steps num_input_tokens_seen int optional defaults to 0 the number of tokens seen during training number of input tokens not the number of prediction tokens total_flos float optional defaults to 0 the total number of floating operations done by the model since the beginning of training stored as floats to avoid overflow log_history list dict str float optional the list of logs done since the beginning of training best_metric float optional when tracking the best model the value of the best metric encountered so far best_model_checkpoint str optional when tracking the best model the value of the name of the checkpoint for the best model encountered so far is_local_process_zero bool optional defaults to true whether or not this process is the local e g on one machine if training in a distributed fashion on several machines main process is_world_process_zero bool optional defaults to true whether or not this process is the global main process when training in a distributed fashion on several machines this is only going to be true for one process is_hyper_param_search bool optional defaults to false whether we are in the process of a hyper parameter search using trainer hyperparameter_search this will impact the way data will be logged in tensorboard save the content of this instance in json format inside json_path create an instance from the content of json_path a class that handles the trainer control flow this class is used by the trainercallback to activate some switches in the training loop args should_training_stop bool optional defaults to false whether or not the training should be interrupted if true this variable will not be set back to false the training will just stop should_epoch_stop bool optional defaults to false whether or not the current epoch should be interrupted if true this variable will be set back to false at the beginning of the next epoch should_save bool optional defaults to false whether or not the model should be saved at this step if true this variable will be set back to false at the beginning of the next step should_evaluate bool optional defaults to false whether or not the model should be evaluated at this step if true this variable will be set back to false at the beginning of the next step should_log bool optional defaults to false whether or not the logs should be reported at this step if true this variable will be set back to false at the beginning of the next step internal method that resets the variable for a new training internal method that resets the variable for a new epoch internal method that resets the variable for a new step no format a class for objects that will inspect the state of the training loop at some events and take some decisions at each of those events the following arguments are available args args trainingarguments the training arguments used to instantiate the trainer state trainerstate the current state of the trainer control trainercontrol the object that is returned to the trainer and can be used to make some decisions model pretrainedmodel or torch nn module the model being trained tokenizer pretrainedtokenizer the tokenizer used for encoding the data optimizer torch optim optimizer the optimizer used for the training steps lr_scheduler torch optim lr_scheduler lambdalr the scheduler used for setting the learning rate train_dataloader torch utils data dataloader optional the current dataloader used for training eval_dataloader torch utils data dataloader optional the current dataloader used for training metrics dict str float the metrics computed by the last evaluation phase those are only accessible in the event on_evaluate logs dict str float the values to log those are only accessible in the event on_log the control object is the only one that can be changed by the callback in which case the event that changes it should return the modified version the argument args state and control are positionals for all events all the others are grouped in kwargs you can unpack the ones you need in the signature of the event using them as an example see the code of the simple transformers printercallback example python class printercallback trainercallback def on_log self args state control logs none kwargs _ logs pop total_flos none if state is_local_process_zero print logs event called at the end of the initialization of the trainer event called at the beginning of training event called at the end of training event called at the beginning of an epoch event called at the end of an epoch event called at the beginning of a training step if using gradient accumulation one training step might take several inputs event called at the end of an substep during gradient accumulation event called at the end of a training step if using gradient accumulation one training step might take several inputs event called after an evaluation phase event called after a successful prediction event called after a checkpoint save event called after logging the last logs event called after a prediction step internal class that just calls the list of callbacks in order a callback can skip the return of control if it doesn t change it a trainercallback that handles the default flow of the training loop for logs evaluation and checkpoints log evaluate save end training log evaluate save a trainercallback that displays the progress of training or evaluation a bare trainercallback that just prints the logs a trainercallback that handles early stopping args early_stopping_patience int use with metric_for_best_model to stop training when the specified metric worsens for early_stopping_patience evaluation calls early_stopping_threshold float optional use with trainingarguments metric_for_best_model and early_stopping_patience to denote how much the specified metric must improve to satisfy early stopping conditions this callback depends on trainingarguments argument load_best_model_at_end functionality to set best_metric in trainerstate note that if the trainingarguments argument save_steps differs from eval_steps the early stopping will not occur until the next save step early_stopping_patience_counter denotes the number of times validation metrics failed to improve best_metric is set by code for load_best_model
import dataclasses
import json
from dataclasses import dataclass
from typing import Dict, List, Optional, Union

import numpy as np
from tqdm.auto import tqdm

from .trainer_utils import IntervalStrategy, has_length
from .training_args import TrainingArguments
from .utils import logging


logger = logging.get_logger(__name__)


@dataclass
class TrainerState:
    

    epoch: Optional[float] = None
    global_step: int = 0
    max_steps: int = 0
    logging_steps: int = 500
    eval_steps: int = 500
    save_steps: int = 500
    num_train_epochs: int = 0
    num_input_tokens_seen: int = 0
    total_flos: float = 0
    log_history: List[Dict[str, float]] = None
    best_metric: Optional[float] = None
    best_model_checkpoint: Optional[str] = None
    is_local_process_zero: bool = True
    is_world_process_zero: bool = True
    is_hyper_param_search: bool = False
    trial_name: str = None
    trial_params: Dict[str, Union[str, float, int, bool]] = None

    def __post_init__(self):
        if self.log_history is None:
            self.log_history = []

    def save_to_json(self, json_path: str):
        
        json_string = json.dumps(dataclasses.asdict(self), indent=2, sort_keys=True) + "\n"
        with open(json_path, "w", encoding="utf-8") as f:
            f.write(json_string)

    @classmethod
    def load_from_json(cls, json_path: str):
        
        with open(json_path, "r", encoding="utf-8") as f:
            text = f.read()
        return cls(**json.loads(text))


@dataclass
class TrainerControl:
    

    should_training_stop: bool = False
    should_epoch_stop: bool = False
    should_save: bool = False
    should_evaluate: bool = False
    should_log: bool = False

    def _new_training(self):
        
        self.should_training_stop = False

    def _new_epoch(self):
        
        self.should_epoch_stop = False

    def _new_step(self):
        
        self.should_save = False
        self.should_evaluate = False
        self.should_log = False


class TrainerCallback:
    
    

    def on_init_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
        
        pass

    def on_train_begin(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
        
        pass

    def on_train_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
        
        pass

    def on_epoch_begin(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
        
        pass

    def on_epoch_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
        
        pass

    def on_step_begin(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
        
        pass

    def on_substep_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
        
        pass

    def on_step_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
        
        pass

    def on_evaluate(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
        
        pass

    def on_predict(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, metrics, **kwargs):
        
        pass

    def on_save(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
        
        pass

    def on_log(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
        
        pass

    def on_prediction_step(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
        
        pass


class CallbackHandler(TrainerCallback):
    

    def __init__(self, callbacks, model, tokenizer, optimizer, lr_scheduler):
        self.callbacks = []
        for cb in callbacks:
            self.add_callback(cb)
        self.model = model
        self.tokenizer = tokenizer
        self.optimizer = optimizer
        self.lr_scheduler = lr_scheduler
        self.train_dataloader = None
        self.eval_dataloader = None

        if not any(isinstance(cb, DefaultFlowCallback) for cb in self.callbacks):
            logger.warning(
                "The Trainer will not work properly if you don't have a `DefaultFlowCallback` in its callbacks. You\n"
                + "should add one before training with `trainer.add_callback(DefaultFlowCallback). The current list of"
                + "callbacks is\n:"
                + self.callback_list
            )

    def add_callback(self, callback):
        cb = callback() if isinstance(callback, type) else callback
        cb_class = callback if isinstance(callback, type) else callback.__class__
        if cb_class in [c.__class__ for c in self.callbacks]:
            logger.warning(
                f"You are adding a {cb_class} to the callbacks of this Trainer, but there is already one. The current"
                + "list of callbacks is\n:"
                + self.callback_list
            )
        self.callbacks.append(cb)

    def pop_callback(self, callback):
        if isinstance(callback, type):
            for cb in self.callbacks:
                if isinstance(cb, callback):
                    self.callbacks.remove(cb)
                    return cb
        else:
            for cb in self.callbacks:
                if cb == callback:
                    self.callbacks.remove(cb)
                    return cb

    def remove_callback(self, callback):
        if isinstance(callback, type):
            for cb in self.callbacks:
                if isinstance(cb, callback):
                    self.callbacks.remove(cb)
                    return
        else:
            self.callbacks.remove(callback)

    @property
    def callback_list(self):
        return "\n".join(cb.__class__.__name__ for cb in self.callbacks)

    def on_init_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl):
        return self.call_event("on_init_end", args, state, control)

    def on_train_begin(self, args: TrainingArguments, state: TrainerState, control: TrainerControl):
        control.should_training_stop = False
        return self.call_event("on_train_begin", args, state, control)

    def on_train_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl):
        return self.call_event("on_train_end", args, state, control)

    def on_epoch_begin(self, args: TrainingArguments, state: TrainerState, control: TrainerControl):
        control.should_epoch_stop = False
        return self.call_event("on_epoch_begin", args, state, control)

    def on_epoch_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl):
        return self.call_event("on_epoch_end", args, state, control)

    def on_step_begin(self, args: TrainingArguments, state: TrainerState, control: TrainerControl):
        control.should_log = False
        control.should_evaluate = False
        control.should_save = False
        return self.call_event("on_step_begin", args, state, control)

    def on_substep_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl):
        return self.call_event("on_substep_end", args, state, control)

    def on_step_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl):
        return self.call_event("on_step_end", args, state, control)

    def on_evaluate(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, metrics):
        control.should_evaluate = False
        return self.call_event("on_evaluate", args, state, control, metrics=metrics)

    def on_predict(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, metrics):
        return self.call_event("on_predict", args, state, control, metrics=metrics)

    def on_save(self, args: TrainingArguments, state: TrainerState, control: TrainerControl):
        control.should_save = False
        return self.call_event("on_save", args, state, control)

    def on_log(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, logs):
        control.should_log = False
        return self.call_event("on_log", args, state, control, logs=logs)

    def on_prediction_step(self, args: TrainingArguments, state: TrainerState, control: TrainerControl):
        return self.call_event("on_prediction_step", args, state, control)

    def call_event(self, event, args, state, control, **kwargs):
        for callback in self.callbacks:
            result = getattr(callback, event)(
                args,
                state,
                control,
                model=self.model,
                tokenizer=self.tokenizer,
                optimizer=self.optimizer,
                lr_scheduler=self.lr_scheduler,
                train_dataloader=self.train_dataloader,
                eval_dataloader=self.eval_dataloader,
                **kwargs,
            )
            
            if result is not None:
                control = result
        return control


class DefaultFlowCallback(TrainerCallback):
    

    def on_step_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
        
        if state.global_step == 1 and args.logging_first_step:
            control.should_log = True
        if args.logging_strategy == IntervalStrategy.STEPS and state.global_step % state.logging_steps == 0:
            control.should_log = True

        
        if (
            args.evaluation_strategy == IntervalStrategy.STEPS
            and state.global_step % state.eval_steps == 0
            and args.eval_delay <= state.global_step
        ):
            control.should_evaluate = True

        
        if (
            args.save_strategy == IntervalStrategy.STEPS
            and state.save_steps > 0
            and state.global_step % state.save_steps == 0
        ):
            control.should_save = True

        
        if state.global_step >= state.max_steps:
            control.should_training_stop = True

        return control

    def on_epoch_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
        
        if args.logging_strategy == IntervalStrategy.EPOCH:
            control.should_log = True

        
        if args.evaluation_strategy == IntervalStrategy.EPOCH and args.eval_delay <= state.epoch:
            control.should_evaluate = True

        
        if args.save_strategy == IntervalStrategy.EPOCH:
            control.should_save = True

        return control


class ProgressCallback(TrainerCallback):
    

    def __init__(self):
        self.training_bar = None
        self.prediction_bar = None

    def on_train_begin(self, args, state, control, **kwargs):
        if state.is_local_process_zero:
            self.training_bar = tqdm(total=state.max_steps, dynamic_ncols=True)
        self.current_step = 0

    def on_step_end(self, args, state, control, **kwargs):
        if state.is_local_process_zero:
            self.training_bar.update(state.global_step - self.current_step)
            self.current_step = state.global_step

    def on_prediction_step(self, args, state, control, eval_dataloader=None, **kwargs):
        if state.is_local_process_zero and has_length(eval_dataloader):
            if self.prediction_bar is None:
                self.prediction_bar = tqdm(
                    total=len(eval_dataloader), leave=self.training_bar is None, dynamic_ncols=True
                )
            self.prediction_bar.update(1)

    def on_evaluate(self, args, state, control, **kwargs):
        if state.is_local_process_zero:
            if self.prediction_bar is not None:
                self.prediction_bar.close()
            self.prediction_bar = None

    def on_predict(self, args, state, control, **kwargs):
        if state.is_local_process_zero:
            if self.prediction_bar is not None:
                self.prediction_bar.close()
            self.prediction_bar = None

    def on_log(self, args, state, control, logs=None, **kwargs):
        if state.is_local_process_zero and self.training_bar is not None:
            _ = logs.pop("total_flos", None)
            self.training_bar.write(str(logs))

    def on_train_end(self, args, state, control, **kwargs):
        if state.is_local_process_zero:
            self.training_bar.close()
            self.training_bar = None


class PrinterCallback(TrainerCallback):
    

    def on_log(self, args, state, control, logs=None, **kwargs):
        _ = logs.pop("total_flos", None)
        if state.is_local_process_zero:
            print(logs)


class EarlyStoppingCallback(TrainerCallback):
    

    def __init__(self, early_stopping_patience: int = 1, early_stopping_threshold: Optional[float] = 0.0):
        self.early_stopping_patience = early_stopping_patience
        self.early_stopping_threshold = early_stopping_threshold
        
        self.early_stopping_patience_counter = 0

    def check_metric_value(self, args, state, control, metric_value):
        
        operator = np.greater if args.greater_is_better else np.less
        if state.best_metric is None or (
            operator(metric_value, state.best_metric)
            and abs(metric_value - state.best_metric) > self.early_stopping_threshold
        ):
            self.early_stopping_patience_counter = 0
        else:
            self.early_stopping_patience_counter += 1

    def on_train_begin(self, args, state, control, **kwargs):
        assert args.load_best_model_at_end, "EarlyStoppingCallback requires load_best_model_at_end = True"
        assert (
            args.metric_for_best_model is not None
        ), "EarlyStoppingCallback requires metric_for_best_model is defined"
        assert (
            args.evaluation_strategy != IntervalStrategy.NO
        ), "EarlyStoppingCallback requires IntervalStrategy of steps or epoch"

    def on_evaluate(self, args, state, control, metrics, **kwargs):
        metric_to_check = args.metric_for_best_model
        if not metric_to_check.startswith("eval_"):
            metric_to_check = f"eval_{metric_to_check}"
        metric_value = metrics.get(metric_to_check)

        if metric_value is None:
            logger.warning(
                f"early stopping required metric_for_best_model, but did not find {metric_to_check} so early stopping"
                " is disabled"
            )
            return

        self.check_metric_value(args, state, control, metric_value)
        if self.early_stopping_patience_counter >= self.early_stopping_patience:
            control.should_training_stop = True