@@ -96,38 +96,45 @@ def fix_valuehead_checkpoint(
9696
9797
9898class FixValueHeadModelCallback (TrainerCallback ):
99+ r"""
100+ A callback for fixing the checkpoint for valuehead models.
101+ """
102+
99103 @override
100104 def on_save (self , args : "TrainingArguments" , state : "TrainerState" , control : "TrainerControl" , ** kwargs ):
101105 r"""
102106 Event called after a checkpoint save.
103107 """
104108 if args .should_save :
109+ output_dir = os .path .join (args .output_dir , "{}-{}" .format (PREFIX_CHECKPOINT_DIR , state .global_step ))
105110 fix_valuehead_checkpoint (
106- model = kwargs .pop ("model" ),
107- output_dir = os .path .join (args .output_dir , "{}-{}" .format (PREFIX_CHECKPOINT_DIR , state .global_step )),
108- safe_serialization = args .save_safetensors ,
111+ model = kwargs .pop ("model" ), output_dir = output_dir , safe_serialization = args .save_safetensors
109112 )
110113
111114
112115class SaveProcessorCallback (TrainerCallback ):
116+ r"""
117+ A callback for saving the processor.
118+ """
119+
113120 def __init__ (self , processor : "ProcessorMixin" ) -> None :
114- r"""
115- Initializes a callback for saving the processor.
116- """
117121 self .processor = processor
118122
123+ @override
124+ def on_save (self , args : "TrainingArguments" , state : "TrainerState" , control : "TrainerControl" , ** kwargs ):
125+ if args .should_save :
126+ output_dir = os .path .join (args .output_dir , "{}-{}" .format (PREFIX_CHECKPOINT_DIR , state .global_step ))
127+ getattr (self .processor , "image_processor" ).save_pretrained (output_dir )
128+
119129 @override
120130 def on_train_end (self , args : "TrainingArguments" , state : "TrainerState" , control : "TrainerControl" , ** kwargs ):
121- r"""
122- Event called at the end of training.
123- """
124131 if args .should_save :
125132 getattr (self .processor , "image_processor" ).save_pretrained (args .output_dir )
126133
127134
128135class PissaConvertCallback (TrainerCallback ):
129136 r"""
130- Initializes a callback for converting the PiSSA adapter to a normal one.
137+ A callback for converting the PiSSA adapter to a normal one.
131138 """
132139
133140 @override
@@ -147,9 +154,6 @@ def on_train_begin(self, args: "TrainingArguments", state: "TrainerState", contr
147154
148155 @override
149156 def on_train_end (self , args : "TrainingArguments" , state : "TrainerState" , control : "TrainerControl" , ** kwargs ):
150- r"""
151- Event called at the end of training.
152- """
153157 if args .should_save :
154158 model = kwargs .pop ("model" )
155159 pissa_init_dir = os .path .join (args .output_dir , "pissa_init" )
@@ -177,21 +181,22 @@ def on_train_end(self, args: "TrainingArguments", state: "TrainerState", control
177181
178182
179183class LogCallback (TrainerCallback ):
184+ r"""
185+ A callback for logging training and evaluation status.
186+ """
187+
180188 def __init__ (self ) -> None :
181- r"""
182- Initializes a callback for logging training and evaluation status.
183- """
184- """ Progress """
189+ # Progress
185190 self .start_time = 0
186191 self .cur_steps = 0
187192 self .max_steps = 0
188193 self .elapsed_time = ""
189194 self .remaining_time = ""
190195 self .thread_pool : Optional ["ThreadPoolExecutor" ] = None
191- """ Status """
196+ # Status
192197 self .aborted = False
193198 self .do_train = False
194- """ Web UI """
199+ # Web UI
195200 self .webui_mode = os .environ .get ("LLAMABOARD_ENABLED" , "0" ).lower () in ["true" , "1" ]
196201 if self .webui_mode :
197202 signal .signal (signal .SIGABRT , self ._set_abort )
@@ -233,9 +238,6 @@ def _close_thread_pool(self) -> None:
233238
234239 @override
235240 def on_init_end (self , args : "TrainingArguments" , state : "TrainerState" , control : "TrainerControl" , ** kwargs ):
236- r"""
237- Event called at the end of the initialization of the `Trainer`.
238- """
239241 if (
240242 args .should_save
241243 and os .path .exists (os .path .join (args .output_dir , TRAINER_LOG ))
@@ -246,60 +248,39 @@ def on_init_end(self, args: "TrainingArguments", state: "TrainerState", control:
246248
247249 @override
248250 def on_train_begin (self , args : "TrainingArguments" , state : "TrainerState" , control : "TrainerControl" , ** kwargs ):
249- r"""
250- Event called at the beginning of training.
251- """
252251 if args .should_save :
253252 self .do_train = True
254253 self ._reset (max_steps = state .max_steps )
255254 self ._create_thread_pool (output_dir = args .output_dir )
256255
257256 @override
258257 def on_train_end (self , args : "TrainingArguments" , state : "TrainerState" , control : "TrainerControl" , ** kwargs ):
259- r"""
260- Event called at the end of training.
261- """
262258 self ._close_thread_pool ()
263259
264260 @override
265261 def on_substep_end (self , args : "TrainingArguments" , state : "TrainerState" , control : "TrainerControl" , ** kwargs ):
266- r"""
267- Event called at the end of an substep during gradient accumulation.
268- """
269262 if self .aborted :
270263 control .should_epoch_stop = True
271264 control .should_training_stop = True
272265
273266 @override
274267 def on_step_end (self , args : "TrainingArguments" , state : "TrainerState" , control : "TrainerControl" , ** kwargs ):
275- r"""
276- Event called at the end of a training step.
277- """
278268 if self .aborted :
279269 control .should_epoch_stop = True
280270 control .should_training_stop = True
281271
282272 @override
283273 def on_evaluate (self , args : "TrainingArguments" , state : "TrainerState" , control : "TrainerControl" , ** kwargs ):
284- r"""
285- Event called after an evaluation phase.
286- """
287274 if not self .do_train :
288275 self ._close_thread_pool ()
289276
290277 @override
291278 def on_predict (self , args : "TrainingArguments" , state : "TrainerState" , control : "TrainerControl" , ** kwargs ):
292- r"""
293- Event called after a successful prediction.
294- """
295279 if not self .do_train :
296280 self ._close_thread_pool ()
297281
298282 @override
299283 def on_log (self , args : "TrainingArguments" , state : "TrainerState" , control : "TrainerControl" , ** kwargs ):
300- r"""
301- Event called after logging the last logs.
302- """
303284 if not args .should_save :
304285 return
305286
@@ -342,9 +323,6 @@ def on_log(self, args: "TrainingArguments", state: "TrainerState", control: "Tra
342323 def on_prediction_step (
343324 self , args : "TrainingArguments" , state : "TrainerState" , control : "TrainerControl" , ** kwargs
344325 ):
345- r"""
346- Event called after a prediction step.
347- """
348326 if self .do_train :
349327 return
350328
0 commit comments