src.core.callbacks.CustomWandbCallback

class CustomWandbCallback(project: str, json_file: str, group: Optional[str] = None, resume: bool = False, resume_run_id: Optional[str] = None, wandb_dir: Optional[str] = None, api_key_path: Optional[str] = None)

Bases: WandbCallback

Custom Weights and Biases Callback used by Mistral for logging information from the Huggingface Trainer.

Methods

on_epoch_begin	Event called at the beginning of an epoch.
on_epoch_end	Event called at the end of an epoch.
on_evaluate	Event called after an evaluation phase.
on_init_end	Event called at the end of the initialization of the [Trainer].
on_log	Event called after logging the last logs.
on_predict	Event called after a successful prediction.
on_prediction_step	Event called after a prediction step.
on_save	Event called after a checkpoint save.
on_step_begin	Event called at the beginning of a training step.
on_step_end	Event called at the end of a training step.
on_substep_end	Event called at the end of an substep during gradient accumulation.
on_train_begin	Calls wandb.init, we add additional arguments to that call using this method.
on_train_end	Event called at the end of training.
setup	Note: have to override this method in order to inject additional arguments into the wandb.init call.

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

Event called at the beginning of an epoch.

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

Event called at the end of an epoch.

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

Event called after an evaluation phase.

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

Event called at the end of the initialization of the [Trainer].

on_log(args, state, control, model: Optional[PreTrainedModel] = None, tokenizer=None, optimizer=None, lr_scheduler=None, train_dataloader=None, eval_dataloader=None, logs=None, **kwargs)

Event called after logging the last logs.

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

Event called after a successful prediction.

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

Event called after a prediction step.

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

Event called after a checkpoint save.

on_step_begin(args: TrainingArguments, state: TrainerState, control: TrainerControl, model: Optional[PreTrainedModel] = None, tokenizer=None, optimizer=None, lr_scheduler=None, train_dataloader=None, eval_dataloader=None, **kwargs)

Event called at the beginning of a training step. If using gradient accumulation, one training step might take several inputs.

on_step_end(args: TrainingArguments, state: TrainerState, control: TrainerControl, model: Optional[PreTrainedModel] = None, tokenizer=None, optimizer=None, lr_scheduler=None, train_dataloader=None, eval_dataloader=None, **kwargs)

Event called at the end of a training step. If using gradient accumulation, one training step might take several inputs.

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

Event called at the end of an substep during gradient accumulation.

on_train_begin(args, state, control, model: Optional[PreTrainedModel] = None, tokenizer=None, optimizer=None, lr_scheduler=None, train_dataloader=None, eval_dataloader=None, **kwargs)

Calls wandb.init, we add additional arguments to that call using this method.

on_train_end(args, state, control, model=None, tokenizer=None, **kwargs)

Event called at the end of training.

setup(args, state, model, **kwargs)

Note: have to override this method in order to inject additional arguments into the wandb.init call. Currently, HF provides no way to pass kwargs to that.

Setup the optional Weights & Biases (wandb) integration.

One can subclass and override this method to customize the setup if needed. Find more information here. You can also override the following environment variables:

Environment:

WANDB_LOG_MODEL (bool, optional, defaults to False):

Whether or not to log model as artifact at the end of training.

WANDB_WATCH (str, optional defaults to "gradients"):

Can be "gradients", "all" or "false". Set to "false" to disable gradient logging or "all" to log gradients and parameters.

WANDB_PROJECT (str, optional, defaults to "huggingface"):

Set this to a custom string to store results in a different project.

WANDB_DISABLED (bool, optional, defaults to False):

Whether or not to disable wandb entirely. Set WANDB_DISABLED=true to disable.