| 16 | """ |
| 17 | |
| 18 | class SavePeftModelCallback(TrainerCallback): |
| 19 | def on_save( |
| 20 | self, |
| 21 | args: TrainingArguments, |
| 22 | state: TrainerState, |
| 23 | control: TrainerControl, |
| 24 | **kwargs, |
| 25 | ): |
| 26 | checkpoint_folder = os.path.join(args.output_dir, f"{PREFIX_CHECKPOINT_DIR}-{state.global_step}") |
| 27 | |
| 28 | kwargs["model"].save_pretrained(checkpoint_folder) |
| 29 | |
| 30 | pytorch_model_path = os.path.join(checkpoint_folder, "pytorch_model.bin") |
| 31 | torch.save({}, pytorch_model_path) |
| 32 | return control |
| 33 | |
| 34 | |
| 35 | class LoadBestPeftModelCallback(TrainerCallback): |
nothing calls this directly
no outgoing calls
no test coverage detected