From b2ef4695a0e0b72ff9e3d4c14ae85b9c35ec24da Mon Sep 17 00:00:00 2001
From: theblackcat102 <theblackcat102@github.com>
Date: Sat, 31 Dec 2022 03:47:54 +0000
Subject: [PATCH] [fix] Fix missing accuracy and eval loss

---
 model/reward/instructor/trainer.py | 43 +++++++++++++++++++++++-------
 1 file changed, 34 insertions(+), 9 deletions(-)

diff --git a/model/reward/instructor/trainer.py b/model/reward/instructor/trainer.py
index 43a5f8ef..45ee76c6 100644
--- a/model/reward/instructor/trainer.py
+++ b/model/reward/instructor/trainer.py
@@ -1,4 +1,6 @@
-from typing import Callable, List, Optional, Tuple, Union, Dict
+import os
+os.environ['WANDB_PROJECT'] = 'reward-model'
+from typing import Any, Callable, List, Optional, Tuple, Union, Dict
 import torch
 from torch import nn
 import numpy as np
@@ -60,6 +62,29 @@ class RankTrainer(Trainer):
 
         return (loss, outputs) if return_outputs else loss
 
+    def _compute_loss(self, model, inputs):
+        inputs = self._prepare_inputs(inputs)
+        outputs = model(**inputs)
+        logits = outputs.get("logits").view(-1, 2)
+        if self.loss_function == 'rank':
+            loss = self.loss_fct(logits[:, 0], logits[:, 1])
+        else:
+            loss = self.loss_fct(logits, torch.zeros(logits.shape[0], device=logits.device, dtype=torch.long))
+
+        return loss, logits
+
+    def prediction_step(self, model: nn.Module, inputs: Dict[str, Union[torch.Tensor, Any]], prediction_loss_only: bool, ignore_keys: Optional[List[str]] = None) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]:
+
+        with torch.no_grad():
+            # compute loss on predict data
+            loss, logits = self._compute_loss(model, inputs)
+
+        loss = loss.mean().detach()
+        labels = torch.zeros(logits.shape[0], device=logits.device, dtype=torch.long)
+        if self.args.prediction_loss_only:
+            return (loss, None, None)
+
+        return (loss, logits, labels)
 
 if __name__ == "__main__":
     model_name = 'bigscience/bloomz-560m'
@@ -67,26 +92,25 @@ if __name__ == "__main__":
     model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=1, problem_type='regression')
     tokenizer = get_tokenizer(model_name)
     args = CustomTrainingArguments(
-        output_dir=f"outputs/{model_name}-finetuned",
-        fp16=True,
+        output_dir=f"{model_name}-finetuned",
         num_train_epochs=4,
         warmup_steps=500,
+        loss_function='rank',
         learning_rate=3e-5,
         # half_precision_backend="apex",
+        fp16=True,
         gradient_checkpointing=False,
-        gradient_accumulation_steps=6,
-        per_device_train_batch_size=12,
+        gradient_accumulation_steps=5,
+        per_device_train_batch_size=16,
         per_device_eval_batch_size=5,
         weight_decay=0.01,
         max_grad_norm=2.0,
         logging_steps=10,
         save_total_limit=4,
         evaluation_strategy='steps',
-        loss_function='rank',
         eval_steps=500,
         save_steps=1000,
-        report_to="wandb",
-        run_name='reward-model'
+        report_to='wandb'
     )
     dataset = WebGPT()
     train, eval = train_val_dataset(dataset)
@@ -97,6 +121,7 @@ if __name__ == "__main__":
         train_dataset=train,
         eval_dataset=eval,
         data_collator=collate_fn,
-        tokenizer=tokenizer
+        tokenizer=tokenizer,
+        compute_metrics=compute_metrics
     )
     trainer.train()