| |
| """Hugging Face Trainer that saves TensorPEFT checkpoints as adapters (not the full model).""" |
| import os |
|
|
| from transformers import Trainer |
|
|
|
|
| class TensorPEFTTrainer(Trainer): |
| """Drop-in ``transformers.Trainer``; checkpoints and ``save_model`` write only the adapter.""" |
|
|
| def _save(self, output_dir=None, state_dict=None): |
| output_dir = output_dir if output_dir is not None else self.args.output_dir |
| os.makedirs(output_dir, exist_ok=True) |
| model = self.accelerator.unwrap_model(self.model) |
| if self.args.should_save: |
| model.save_pretrained(output_dir) |
| if self.processing_class is not None: |
| self.processing_class.save_pretrained(output_dir) |
|
|
