support streaming data, fix #284 #274 #268

Former-commit-id: 819cc1353599e5fa45658bc56dd0dbe4b258b197

src/llmtuner/tuner/rm/collator.py

@@ -15,5 +15,8 @@ class PairwiseDataCollatorWithPadding(DataCollatorWithPadding):
         We generate 2 * n examples where the first n examples represent chosen examples and
         the last n examples represent rejected examples.
         """
-        features = [{"input_ids": feature[key]} for key in ("accept_ids", "reject_ids") for feature in features]
+        features = [
+            {"input_ids": feature[key], "attention_mask": [1] * len(feature[key])}
+            for key in ("accept_ids", "reject_ids") for feature in features
+        ]
         return super().__call__(features)

src/llmtuner/tuner/rm/trainer.py

@@ -1,13 +1,15 @@
 import os
 import json
 import torch
-from typing import Dict, List, Optional, Tuple, Union
-from transformers.trainer import PredictionOutput
-from transformers.modeling_utils import PreTrainedModel
+from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union
 
 from llmtuner.extras.logging import get_logger
 from llmtuner.tuner.core.trainer import PeftTrainer
 
+if TYPE_CHECKING:
+    from transformers.trainer import PredictionOutput
+    from transformers.modeling_utils import PreTrainedModel
+
 
 logger = get_logger(__name__)
 
@@ -23,7 +25,7 @@ class PairwisePeftTrainer(PeftTrainer):
 
     def compute_loss(
         self,
-        model: PreTrainedModel,
+        model: "PreTrainedModel",
         inputs: Dict[str, torch.Tensor],
         return_outputs: Optional[bool] = False
     ) -> Union[torch.Tensor, Tuple[torch.Tensor, List[torch.Tensor]]]:
@@ -46,7 +48,7 @@ class PairwisePeftTrainer(PeftTrainer):
 
     def save_predictions(
         self,
-        predict_results: PredictionOutput
+        predict_results: "PredictionOutput"
     ) -> None:
         r"""
         Saves model predictions to `output_dir`.

src/llmtuner/tuner/rm/workflow.py

@@ -2,25 +2,27 @@
 # https://github.com/lvwerra/trl/blob/main/examples/summarization/scripts/reward_summarization.py
 # https://github.com/CarperAI/trlx/blob/main/examples/summarize_rlhf/reward_model/train_reward_model_gptj.py
 
-from typing import Optional, List
-from transformers import Seq2SeqTrainingArguments, TrainerCallback
+from typing import TYPE_CHECKING, Optional, List
 
 from llmtuner.dsets import get_dataset, preprocess_dataset, split_dataset
 from llmtuner.extras.callbacks import LogCallback
 from llmtuner.extras.ploting import plot_loss
-from llmtuner.hparams import ModelArguments, DataArguments, FinetuningArguments
 from llmtuner.tuner.core import load_model_and_tokenizer
 from llmtuner.tuner.rm.metric import compute_accuracy
 from llmtuner.tuner.rm.collator import PairwiseDataCollatorWithPadding
 from llmtuner.tuner.rm.trainer import PairwisePeftTrainer
 
+if TYPE_CHECKING:
+    from transformers import Seq2SeqTrainingArguments, TrainerCallback
+    from llmtuner.hparams import ModelArguments, DataArguments, FinetuningArguments
+
 
 def run_rm(
-    model_args: ModelArguments,
-    data_args: DataArguments,
-    training_args: Seq2SeqTrainingArguments,
-    finetuning_args: FinetuningArguments,
-    callbacks: Optional[List[TrainerCallback]] = [LogCallback()]
+    model_args: "ModelArguments",
+    data_args: "DataArguments",
+    training_args: "Seq2SeqTrainingArguments",
+    finetuning_args: "FinetuningArguments",
+    callbacks: Optional[List["TrainerCallback"]] = [LogCallback()]
 ):
     dataset = get_dataset(model_args, data_args)
     model, tokenizer = load_model_and_tokenizer(model_args, finetuning_args, training_args.do_train, stage="rm")