format style

Former-commit-id: 53b683531b83cd1d19de97c6565f16c1eca6f5e1

src/llmtuner/train/rm/collator.py

@@ -1,6 +1,7 @@
-import torch
 from dataclasses import dataclass
 from typing import Any, Dict, Sequence
+
+import torch
 from transformers import DataCollatorWithPadding
 
 
@@ -20,8 +21,9 @@ class PairwiseDataCollatorWithPadding(DataCollatorWithPadding):
         features = [
             {
                 "input_ids": feature["prompt_ids"] + feature[key],
-                "attention_mask": [1] * (len(feature["prompt_ids"]) + len(feature[key]))
+                "attention_mask": [1] * (len(feature["prompt_ids"]) + len(feature[key])),
             }
-            for key in ("chosen_ids", "rejected_ids") for feature in features
+            for key in ("chosen_ids", "rejected_ids")
+            for feature in features
         ]
         return super().__call__(features)

src/llmtuner/train/rm/metric.py

@@ -1,6 +1,7 @@
-import numpy as np
 from typing import Dict, Sequence, Tuple, Union
 
+import numpy as np
+
 
 def compute_accuracy(eval_preds: Sequence[Union[np.ndarray, Tuple[np.ndarray]]]) -> Dict[str, float]:
     preds, _ = eval_preds

src/llmtuner/train/rm/trainer.py

@@ -1,14 +1,16 @@
-import os
 import json
-import torch
+import os
 from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union
+
+import torch
 from transformers import Trainer
 
 from ...extras.logging import get_logger
 
+
 if TYPE_CHECKING:
-    from transformers.trainer import PredictionOutput
     from transformers.modeling_utils import PreTrainedModel
+    from transformers.trainer import PredictionOutput
 
 
 logger = get_logger(__name__)
@@ -21,13 +23,10 @@ class PairwiseTrainer(Trainer):
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
-        self.can_return_loss = True # override property to return eval_loss
+        self.can_return_loss = True  # override property to return eval_loss
 
     def compute_loss(
-        self,
-        model: "PreTrainedModel",
-        inputs: Dict[str, torch.Tensor],
-        return_outputs: Optional[bool] = False
+        self, model: "PreTrainedModel", inputs: Dict[str, torch.Tensor], return_outputs: Optional[bool] = False
     ) -> Union[torch.Tensor, Tuple[torch.Tensor, List[torch.Tensor]]]:
         r"""
         Computes pairwise loss. The first n examples are chosen and the last n examples are rejected.
@@ -68,9 +67,9 @@ class PairwiseTrainer(Trainer):
             assert div_index > 0
             chosen_trunc_rewards = chosen_rewards[i, div_index:end_index]
             rejected_trunc_rewards = rejected_rewards[i, div_index:end_index]
-            if return_outputs: # use the score on the last token except pad token for inference
-                chosen_scores.append(chosen_rewards[i, chosen_length-1])
-                rejected_scores.append(rejected_rewards[i, rejected_length-1])
+            if return_outputs:  # use the score on the last token except pad token for inference
+                chosen_scores.append(chosen_rewards[i, chosen_length - 1])
+                rejected_scores.append(rejected_rewards[i, rejected_length - 1])
             loss += -torch.nn.functional.logsigmoid(chosen_trunc_rewards - rejected_trunc_rewards).mean()
 
         loss = loss / batch_size
@@ -80,10 +79,7 @@ class PairwiseTrainer(Trainer):
 
         return loss
 
-    def save_predictions(
-        self,
-        predict_results: "PredictionOutput"
-    ) -> None:
+    def save_predictions(self, predict_results: "PredictionOutput") -> None:
         r"""
         Saves model predictions to `output_dir`.
 

src/llmtuner/train/rm/workflow.py

@@ -1,6 +1,7 @@
 # Inspired by: https://github.com/CarperAI/trlx/blob/main/examples/summarize_rlhf/reward_model/train_reward_model_gptj.py
 
-from typing import TYPE_CHECKING, Optional, List
+from typing import TYPE_CHECKING, List, Optional
+
 from transformers import Seq2SeqTrainingArguments
 
 from ...data import get_dataset, split_dataset
@@ -13,9 +14,11 @@ from ...train.rm.metric import compute_accuracy
 from ...train.rm.trainer import PairwiseTrainer
 from ...train.utils import create_modelcard_and_push
 
+
 if TYPE_CHECKING:
     from transformers import TrainerCallback
-    from ...hparams import ModelArguments, DataArguments, FinetuningArguments
+
+    from ...hparams import DataArguments, FinetuningArguments, ModelArguments
 
 
 def run_rm(
@@ -23,15 +26,17 @@ def run_rm(
     data_args: "DataArguments",
     training_args: "Seq2SeqTrainingArguments",
     finetuning_args: "FinetuningArguments",
-    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="rm")
     data_collator = PairwiseDataCollatorWithPadding(tokenizer, pad_to_multiple_of=8)
 
     # Update arguments
     training_args_dict = training_args.to_dict()
-    training_args_dict.update(dict(remove_unused_columns=False)) # important for pairwise dataset
+    training_args_dict.update(dict(remove_unused_columns=False))  # important for pairwise dataset
     training_args = Seq2SeqTrainingArguments(**training_args_dict)
 
     # Initialize our Trainer
@@ -42,7 +47,7 @@ def run_rm(
         data_collator=data_collator,
         callbacks=callbacks + [FixValueHeadModelCallback()],
         compute_metrics=compute_accuracy,
-        **split_dataset(dataset, data_args, training_args)
+        **split_dataset(dataset, data_args, training_args),
     )
 
     # Training