From 3a10f1024ab16a00acb42b400ac5195a0aec07b5 Mon Sep 17 00:00:00 2001
From: theblackcat102 <theblackcat102@github.com>
Date: Sat, 31 Dec 2022 09:27:09 +0000
Subject: [PATCH] [fix] Fix truncation in collate fn

---
 model/reward/instructor/rank_datasets.py | 11 +++++++----
 model/reward/instructor/trainer.py       | 15 ++++++++-------
 2 files changed, 15 insertions(+), 11 deletions(-)

diff --git a/model/reward/instructor/rank_datasets.py b/model/reward/instructor/rank_datasets.py
index e407b30f..128baafe 100644
--- a/model/reward/instructor/rank_datasets.py
+++ b/model/reward/instructor/rank_datasets.py
@@ -38,8 +38,10 @@ class DataCollatorForPairRank:
         batch_size = 0
         for question, pairs in features:
             for (pos, neg) in pairs:
-                flatten_features.append(self.tokenizer(question, pos, truncation=True))
-                flatten_features.append(self.tokenizer(question, neg, truncation=True))
+                flatten_features.append(self.tokenizer(question, pos,
+                    truncation=True, max_length=self.max_length))
+                flatten_features.append(self.tokenizer(question, neg,
+                    truncation=True, max_length=self.max_length))
                 batch_size += 1
         
         batch = self.tokenizer.pad(
@@ -147,6 +149,8 @@ class HFSummary(Dataset):
 
         self.summaries = summaries
 
+        self.postfix_prompt = ' TLDR;'
+
     def __len__(self):
         return len(self.index2summary)
 
@@ -159,6 +163,5 @@ class HFSummary(Dataset):
         # not optimal but good for now
         valid_idx = np.random.choice(len(rows), self.max_comparison_per_sample)
         # optimize the format later
-        return context, [ r for idx, r in enumerate(rows) if idx in valid_idx ]
-
+        return context+self.postfix_prompt, [ r for idx, r in enumerate(rows) if idx in valid_idx ]
 
diff --git a/model/reward/instructor/trainer.py b/model/reward/instructor/trainer.py
index 45ee76c6..586c8d47 100644
--- a/model/reward/instructor/trainer.py
+++ b/model/reward/instructor/trainer.py
@@ -6,10 +6,10 @@ from torch import nn
 import numpy as np
 import evaluate
 from dataclasses import dataclass
-from torch.utils.data import Dataset
+from torch.utils.data import Dataset, ConcatDataset
 from transformers import AutoModelForSequenceClassification, AutoModelForMultipleChoice
 from transformers import Trainer, PreTrainedModel, TrainingArguments, DataCollator, EvalPrediction, TrainerCallback, PreTrainedTokenizerBase
-from rank_datasets import DataCollatorForPairRank, WebGPT
+from rank_datasets import DataCollatorForPairRank, WebGPT, HFSummary
 from utils import get_tokenizer, train_val_dataset
 
 accuracy = evaluate.load("accuracy")
@@ -88,7 +88,7 @@ class RankTrainer(Trainer):
 
 if __name__ == "__main__":
     model_name = 'bigscience/bloomz-560m'
-    model_name = 'google/electra-base-discriminator'
+    model_name = 'google/electra-large-discriminator'
     model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=1, problem_type='regression')
     tokenizer = get_tokenizer(model_name)
     args = CustomTrainingArguments(
@@ -99,9 +99,9 @@ if __name__ == "__main__":
         learning_rate=3e-5,
         # half_precision_backend="apex",
         fp16=True,
-        gradient_checkpointing=False,
-        gradient_accumulation_steps=5,
-        per_device_train_batch_size=16,
+        gradient_checkpointing=True,
+        gradient_accumulation_steps=8,
+        per_device_train_batch_size=8,
         per_device_eval_batch_size=5,
         weight_decay=0.01,
         max_grad_norm=2.0,
@@ -114,7 +114,8 @@ if __name__ == "__main__":
     )
     dataset = WebGPT()
     train, eval = train_val_dataset(dataset)
-    collate_fn = DataCollatorForPairRank(tokenizer, max_length=400)
+    train = ConcatDataset([train, HFSummary()])
+    collate_fn = DataCollatorForPairRank(tokenizer, max_length=440)
     trainer = RankTrainer(
         model,
         args,