| """Save full model+tokenizer copies outside the rotating HF checkpoint window.""" |
|
|
| from __future__ import annotations |
|
|
| from pathlib import Path |
|
|
| from transformers import TrainerCallback, TrainerControl, TrainerState, TrainingArguments |
| from transformers.utils import logging |
|
|
| logger = logging.get_logger(__name__) |
|
|
|
|
| class PermanentCheckpointCallback(TrainerCallback): |
| """On each multiple of ``every_n_steps``, call ``trainer.save_model`` into a fixed subfolder.""" |
|
|
| def __init__(self, every_n_steps: int, subdir: str = "checkpoints/permanent") -> None: |
| self.every_n_steps = int(every_n_steps) |
| self.subdir = subdir |
| self._trainer = None |
|
|
| def bind_trainer(self, trainer) -> None: |
| self._trainer = trainer |
|
|
| def on_step_end( |
| self, |
| args: TrainingArguments, |
| state: TrainerState, |
| control: TrainerControl, |
| **kwargs, |
| ) -> TrainerControl: |
| if self.every_n_steps <= 0 or not state.is_world_process_zero: |
| return control |
| if state.global_step <= 0: |
| return control |
| if state.global_step % self.every_n_steps != 0: |
| return control |
| tr = self._trainer |
| if tr is None: |
| logger.warning("PermanentCheckpointCallback: trainer not bound; skipping save at step %s", state.global_step) |
| return control |
| out = Path(args.output_dir) / self.subdir / f"checkpoint-{state.global_step}" |
| if out.exists(): |
| return control |
| out.parent.mkdir(parents=True, exist_ok=True) |
| tr.save_model(str(out)) |
| logger.info("Permanent checkpoint: %s", out) |
| return control |
|
|
