update web UI, support rm predict #210

Former-commit-id: 92cc6b655dc91b94d5bf9d8618c3b57d5cf94333

src/llmtuner/tuner/rm/trainer.py

@@ -1,10 +1,17 @@
+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 llmtuner.extras.logging import get_logger
 from llmtuner.tuner.core.trainer import PeftTrainer
 
 
+logger = get_logger(__name__)
+
+
 class PairwisePeftTrainer(PeftTrainer):
     r"""
     Inherits PeftTrainer to compute pairwise loss.
@@ -36,3 +43,26 @@ class PairwisePeftTrainer(PeftTrainer):
         r_accept, r_reject = values[:, -1].split(batch_size, dim=0)
         loss = -torch.log(torch.sigmoid(r_accept - r_reject)).mean()
         return (loss, [loss, r_accept, r_reject]) if return_outputs else loss
+
+    def save_predictions(
+        self,
+        predict_results: PredictionOutput
+    ) -> None:
+        r"""
+        Saves model predictions to `output_dir`.
+
+        A custom behavior that not contained in Seq2SeqTrainer.
+        """
+        if not self.is_world_process_zero():
+            return
+
+        output_prediction_file = os.path.join(self.args.output_dir, "generated_predictions.jsonl")
+        logger.info(f"Saving prediction results to {output_prediction_file}")
+
+        acc_scores, rej_scores = predict_results.predictions
+
+        with open(output_prediction_file, "w", encoding="utf-8") as writer:
+            res: List[str] = []
+            for acc_score, rej_score in zip(acc_scores, rej_scores):
+                res.append(json.dumps({"accept": round(float(acc_score), 2), "reject": round(float(rej_score), 2)}))
+            writer.write("\n".join(res))

src/llmtuner/tuner/rm/workflow.py

@@ -56,3 +56,10 @@ def run_rm(
         metrics = trainer.evaluate(metric_key_prefix="eval")
         trainer.log_metrics("eval", metrics)
         trainer.save_metrics("eval", metrics)
+
+    # Predict
+    if training_args.do_predict:
+        predict_results = trainer.predict(dataset, metric_key_prefix="predict")
+        trainer.log_metrics("predict", predict_results.metrics)
+        trainer.save_metrics("predict", predict_results.metrics)
+        trainer.save_predictions(predict_results)