format style

src/llmtuner/train/ppo/workflow.py

@@ -1,23 +1,26 @@
 # Inspired by: https://github.com/lvwerra/trl/blob/main/examples/research_projects/stack_llama/scripts/rl_training.py
 
 import math
-from trl import PPOConfig
+from typing import TYPE_CHECKING, List, Optional
+
 from torch.optim import AdamW
-from typing import TYPE_CHECKING, Optional, List
 from transformers import DataCollatorWithPadding
 from transformers.optimization import get_scheduler
+from trl import PPOConfig
 
 from ...data import get_dataset
 from ...extras.callbacks import FixValueHeadModelCallback
 from ...extras.misc import fix_valuehead_checkpoint
 from ...extras.ploting import plot_loss
 from ...model import load_model_and_tokenizer
-from ...train.utils import create_ref_model, create_reward_model
 from ...train.ppo.trainer import CustomPPOTrainer
+from ...train.utils import create_ref_model, create_reward_model
+
 
 if TYPE_CHECKING:
     from transformers import Seq2SeqTrainingArguments, TrainerCallback
-    from ...hparams import ModelArguments, DataArguments, FinetuningArguments, GeneratingArguments
+
+    from ...hparams import DataArguments, FinetuningArguments, GeneratingArguments, ModelArguments
 
 
 def run_ppo(
@@ -26,12 +29,14 @@ def run_ppo(
     training_args: "Seq2SeqTrainingArguments",
     finetuning_args: "FinetuningArguments",
     generating_args: "GeneratingArguments",
-    callbacks: Optional[List["TrainerCallback"]] = None
+    callbacks: Optional[List["TrainerCallback"]] = None,
 ):
-    model, tokenizer = load_model_and_tokenizer(model_args, finetuning_args, training_args.do_train, add_valuehead=True)
+    model, tokenizer = load_model_and_tokenizer(
+        model_args, finetuning_args, training_args.do_train, add_valuehead=True
+    )
     dataset = get_dataset(tokenizer, model_args, data_args, training_args, stage="ppo")
 
-    tokenizer.padding_side = "left" # use left-padding in generation while using right-padding in training
+    tokenizer.padding_side = "left"  # use left-padding in generation while using right-padding in training
     data_collator = DataCollatorWithPadding(tokenizer=tokenizer)
 
     # Create reference model and reward model
@@ -55,7 +60,7 @@ def run_ppo(
         use_score_scaling=finetuning_args.ppo_score_norm,
         use_score_norm=finetuning_args.ppo_score_norm,
         whiten_rewards=finetuning_args.ppo_whiten_rewards,
-        accelerator_kwargs={"step_scheduler_with_optimizer": False}
+        accelerator_kwargs={"step_scheduler_with_optimizer": False},
     )
 
     # Create optimizer and scheduler
@@ -70,7 +75,7 @@ def run_ppo(
         training_args.lr_scheduler_type,
         optimizer=optimizer,
         num_warmup_steps=training_args.get_warmup_steps(num_training_steps),
-        num_training_steps=num_training_steps
+        num_training_steps=num_training_steps,
     )
 
     # Initialize our Trainer
@@ -88,7 +93,7 @@ def run_ppo(
         dataset=dataset,
         data_collator=data_collator,
         optimizer=optimizer,
-        lr_scheduler=lr_scheduler
+        lr_scheduler=lr_scheduler,
     )
 
     # Training
@@ -97,6 +102,6 @@ def run_ppo(
         ppo_trainer.save_model()
         if training_args.should_save:
             fix_valuehead_checkpoint(model, training_args.output_dir, training_args.save_safetensors)
-        ppo_trainer.save_state() # must be called after save_model to have a folder
+        ppo_trainer.save_state()  # must be called after save_model to have a folder
         if ppo_trainer.is_world_process_zero() and finetuning_args.plot_loss:
             plot_loss(training_args.output_dir, keys=["loss", "reward"])