| 33 | |
| 34 | |
| 35 | class LoadBestPeftModelCallback(TrainerCallback): |
| 36 | def on_train_end( |
| 37 | self, |
| 38 | args: TrainingArguments, |
| 39 | state: TrainerState, |
| 40 | control: TrainerControl, |
| 41 | **kwargs, |
| 42 | ): |
| 43 | print(f"Loading best peft model from {state.best_model_checkpoint} (score: {state.best_metric}).") |
| 44 | best_model_path = os.path.join(state.best_model_checkpoint, "adapter_model.bin") |
| 45 | adapters_weights = torch.load(best_model_path) |
| 46 | model = kwargs["model"] |
| 47 | set_peft_model_state_dict(model, adapters_weights) |
| 48 | return control |
| 49 | |
| 50 | |
| 51 | def get_args(): |
nothing calls this directly
no outgoing calls
no test coverage detected