From ad98a282410664bf03cbbaf1b835b440a26d0409 Mon Sep 17 00:00:00 2001
From: theblackcat102 <theblackcat102@github.com>
Date: Fri, 30 Dec 2022 17:25:50 +0000
Subject: [PATCH 01/41] [feature] add rank dataset for webgpt and human
 feedback summary

---
 model/reward/instructor/README.md             |   7 +
 model/reward/instructor/TODO.md               |  12 ++
 model/reward/instructor/cls_dataset.py        |  73 +++++++++
 .../reward/instructor/experimental_dataset.py |  11 ++
 model/reward/instructor/rank_datasets.py      | 145 ++++++++++++++++++
 model/reward/instructor/tests/__init__.py     |   0
 model/reward/instructor/tests/test_dataset.py |  28 ++++
 model/reward/instructor/trainer.py            |   2 +
 model/reward/instructor/utils.py              |  18 +++
 model/utils.py                                |   4 +
 10 files changed, 300 insertions(+)
 create mode 100644 model/reward/instructor/README.md
 create mode 100644 model/reward/instructor/TODO.md
 create mode 100644 model/reward/instructor/cls_dataset.py
 create mode 100644 model/reward/instructor/experimental_dataset.py
 create mode 100644 model/reward/instructor/rank_datasets.py
 create mode 100644 model/reward/instructor/tests/__init__.py
 create mode 100644 model/reward/instructor/tests/test_dataset.py
 create mode 100644 model/reward/instructor/trainer.py
 create mode 100644 model/reward/instructor/utils.py
 create mode 100644 model/utils.py

diff --git a/model/reward/instructor/README.md b/model/reward/instructor/README.md
new file mode 100644
index 00000000..7dbfefbc
--- /dev/null
+++ b/model/reward/instructor/README.md
@@ -0,0 +1,7 @@
+
+
+
+```bash
+
+
+```
\ No newline at end of file
diff --git a/model/reward/instructor/TODO.md b/model/reward/instructor/TODO.md
new file mode 100644
index 00000000..33bc6595
--- /dev/null
+++ b/model/reward/instructor/TODO.md
@@ -0,0 +1,12 @@
+
+Some other reward features we can use
+
+
+Summaries from human feedback
+
+* use `confidence` score into the RM learning, ensure the output rank score correlates with confidence
+
+* each labeling has a labeling `note`, basically comments by labeler, not sure what else we can use
+
+
+
diff --git a/model/reward/instructor/cls_dataset.py b/model/reward/instructor/cls_dataset.py
new file mode 100644
index 00000000..54bbd19e
--- /dev/null
+++ b/model/reward/instructor/cls_dataset.py
@@ -0,0 +1,73 @@
+'''
+
+    classification based ranking
+
+'''
+import os
+import json
+import random
+import torch
+import numpy as np
+from dataset import load_dataset
+from torch.utils.data import Dataset
+from .utils import webgpt_return_format
+
+class WebGPTDataset(Dataset):
+    def __init__(self, mode='train', index_cache='dataset/webgpt_train_idx.pt', additional_dataset=None) -> None:
+        super().__init__()
+        '''
+            mode : train or val, used for validation purpose, has nothing to do with original split
+            additional_dataset : a list of jsonline format with idx, question and texts (generate candidates)
+                idx : must match the index you iterate from comparison enumerate order
+                question : for validation purpose
+                texts : list of K generate results from the question prompt
+        '''
+        os.makedirs('dataset', exist_ok=True)
+        dataset = load_dataset("openai/webgpt_comparisons")
+        if os.path.exists(index_cache):
+            train_idx = torch.load(index_cache)
+        else:
+            train_idx = np.random.choice(range(len(dataset['train'])), int(len(dataset['train'])*0.8), replace=False)
+            torch.save(set(train_idx.tolist()), index_cache)
+        self.dataset = []
+        self.dataset_index = []
+        for idx, row in enumerate(dataset['train']):
+            if mode == 'train' and idx in train_idx:
+                self.dataset.append(webgpt_return_format(row))
+                self.dataset_index.append(idx)
+            elif idx not in train_idx and mode != 'train':
+                self.dataset.append(webgpt_return_format(row))
+                self.dataset_index.append(idx)
+
+        # since this dataset was generated from 176B GPT-3
+        # we needed some more sample generated from the starting model
+        # since this model must rank model generated by GPT-3 being better than your starting model
+        self.sample_additional = False
+        if additional_dataset is not None:
+            self.sample_additional = True
+            self.additional = {}
+            with open(additional_dataset, 'r') as f:
+                for line in f:
+                    row = json.loads(line)
+                    if row['idx'] in self.dataset_index:
+                        self.additional[row['idx']] = row['negatives']
+            if len(self.additional) != len(self.dataset_index):
+                for match_idx in self.dataset_index:
+                    if match_idx in self.additional:
+                        continue
+
+                    idx = match_idx-900
+                    while idx not in self.additional:
+                        idx -= 1
+                    self.additional[match_idx] = self.additional[idx]
+
+    def __len__(self):
+        return len(self.dataset)
+
+    def __getitem__(self, index):
+        row = self.dataset[index]
+        if not self.sample_additional:
+            return row['question'], row['pos'], row['neg']
+
+        gen_neg = random.choice(self.additional[self.dataset_index[index]])
+        return row['question'], row['pos'], row['neg'], gen_neg
diff --git a/model/reward/instructor/experimental_dataset.py b/model/reward/instructor/experimental_dataset.py
new file mode 100644
index 00000000..145588c4
--- /dev/null
+++ b/model/reward/instructor/experimental_dataset.py
@@ -0,0 +1,11 @@
+'''
+    
+
+'''
+import os
+import json
+import random
+import torch
+import numpy as np
+from dataset import load_dataset
+from torch.utils.data import Dataset
diff --git a/model/reward/instructor/rank_datasets.py b/model/reward/instructor/rank_datasets.py
new file mode 100644
index 00000000..7fef5ab7
--- /dev/null
+++ b/model/reward/instructor/rank_datasets.py
@@ -0,0 +1,145 @@
+'''
+    author: theblackcat102
+
+    A list of rank based dataset for training using rank loss
+
+    Some nice features to have
+
+    [ ] 
+
+'''
+import os
+import glob
+import json
+import numpy as np
+from torch.utils.data import Dataset
+from datasets import load_dataset
+
+class CollateFN():
+    def __init__(self, tokenizer, max_length=400) -> None:
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+
+    def __call__(self, batch):
+        prompts = []
+        pos_sentences = []
+        neg_sentences = []
+        for prompt, pairs in batch:
+            for (pos, neg) in pairs:
+                prompts.append(prompt)
+                pos_sentences.append(pos)
+                neg_sentences.append(neg)
+
+        batch = [self.tokenizer(prompts, pos_sentences, return_tensors='pt', max_length=self.max_length, padding=True, truncation=True),\
+                self.tokenizer(prompts, neg_sentences, return_tensors='pt', max_length=self.max_length, padding=True, truncation=True)]
+        return batch
+
+class WebGPT(Dataset):
+
+    def __init__(self) -> None:
+        super().__init__()
+
+        dataset = load_dataset("openai/webgpt_comparisons")
+        questions = {}
+        # using prompt as our index will allows us
+        # to add additional generated prompt later
+        self.index2question = {}
+        for row in dataset['train']:
+            question = row['question']['full_text']
+            if question not in self.index2question:
+                self.index2question[len(self.index2question)] = question
+
+            if question not in questions:
+                questions[question] = []
+
+            if row['score_0'] > row['score_1']:
+                # not going to risk it
+                questions[question].append((
+                    row['answer_0'], row['answer_1']
+                ))
+            else:
+                questions[question].append((
+                    row['answer_1'], row['answer_0']
+                ))
+
+        self.questions = questions
+
+    def __len__(self):
+        return len(self.index2question)
+
+    def __getitem__(self, index):
+        question = self.index2question[index]
+        rows = self.questions[question]
+        # optimize the format later
+        return question, rows
+
+
+
+
+class HFSummary(Dataset):
+    '''
+        Human feedback data from OpenAI
+        https://github.com/openai/summarize-from-feedback
+
+            >> azcopy copy "https://openaipublic.blob.core.windows.net/summarize-from-feedback/dataset/*" . --recursive
+        
+        choice : 0 or 1
+
+    '''
+    def __init__(self, split='train',
+        path='summarize-from-feedback/comparisons/*.json',
+        conf_threshold=-1,
+        max_comparison_per_sample=5) -> None:
+        super().__init__()
+        assert split in ('train', 'valid1', 'valid2', 'test')
+        summaries = {}
+        # using prompt as our index will allows us
+        # to add additional generated prompt later
+        self.index2summary = {}
+        self.max_comparison_per_sample = max_comparison_per_sample
+        for jsonl_file in glob.glob(path):
+            with open(jsonl_file, 'r') as f:
+                for line in f:
+                    data = json.loads(line)
+                    if data['split'] != split:
+                        continue
+                    if 'extra' in data and \
+                        'confidence' in data['extra'] and \
+                        conf_threshold > data['extra']['confidence']:
+                        print('skipping {}'.format(data['info']['id']))
+                        continue
+
+                    if 'article' in data['info']:
+                        context = data['info']['article']
+                    elif 'post' in data['info']:
+                        context = data['info']['post']
+
+                    if context not in self.index2summary:
+                        self.index2summary[len(self.index2summary)] = context
+                    
+                    if context not in summaries:
+                        summaries[context] = []
+
+                    pos, neg = (0, 1) if data['choice'] == 0 else (1, 0)
+                    summaries[context].append((
+                        data['summaries'][pos]['text'],
+                        data['summaries'][neg]['text']
+                    ))
+
+        self.summaries = summaries
+
+    def __len__(self):
+        return len(self.index2summary)
+
+    def __getitem__(self, index):
+        context = self.index2summary[index]
+        # return pairs of comparison
+        rows = self.summaries[context]
+        # pair very big
+        # we are going to do some sampling
+        # 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 ]
+
+
diff --git a/model/reward/instructor/tests/__init__.py b/model/reward/instructor/tests/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/model/reward/instructor/tests/test_dataset.py b/model/reward/instructor/tests/test_dataset.py
new file mode 100644
index 00000000..4dd59c16
--- /dev/null
+++ b/model/reward/instructor/tests/test_dataset.py
@@ -0,0 +1,28 @@
+from transformers import AutoTokenizer
+from torch.utils.data import DataLoader
+from rank_datasets import WebGPT, HFSummary, CollateFN
+
+
+def test_hfsummary():
+    
+    tokenizer = AutoTokenizer.from_pretrained("bigscience/mt0-large")
+    collate_fn = CollateFN(tokenizer)
+    dataset = HFSummary()
+    dataloader = DataLoader(dataset, collate_fn=collate_fn, batch_size=8)
+    for batch in dataloader:
+        print(batch[0]['input_ids'].shape)
+ 
+
+def test_webgpt():
+    
+    tokenizer = AutoTokenizer.from_pretrained("bigscience/mt0-large")
+    collate_fn = CollateFN(tokenizer)
+    dataset = WebGPT()
+    dataloader = DataLoader(dataset, collate_fn=collate_fn, batch_size=32)
+    for batch in dataloader:
+        print(batch[0]['input_ids'].shape)
+
+
+if __name__ == "__main__":
+    test_hfsummary()
+    # test_webgpt()
\ No newline at end of file
diff --git a/model/reward/instructor/trainer.py b/model/reward/instructor/trainer.py
new file mode 100644
index 00000000..9ee5e043
--- /dev/null
+++ b/model/reward/instructor/trainer.py
@@ -0,0 +1,2 @@
+import wandb
+from accelerate import Accelerator
diff --git a/model/reward/instructor/utils.py b/model/reward/instructor/utils.py
new file mode 100644
index 00000000..1487947c
--- /dev/null
+++ b/model/reward/instructor/utils.py
@@ -0,0 +1,18 @@
+import re
+
+re_reference_remove = re.compile(r'\[([0-9])+\]|\[([0-9])+,([0-9])+\]')
+
+def webgpt_return_format(row):
+    if row['score_0'] >= row['score_1']:
+        # remove this to prevent information leak, since we are not using reference
+        return {
+                'question': row['question']['full_text'],
+                     'pos': re_reference_remove.sub('', row['answer_0']),
+                     'neg': re_reference_remove.sub('', row['answer_1'])
+                }
+
+    return {
+            'question': row['question']['full_text'],
+                 'pos': re_reference_remove.sub('', row['answer_1']),
+                 'neg': re_reference_remove.sub('', row['answer_0'])
+            }
diff --git a/model/utils.py b/model/utils.py
new file mode 100644
index 00000000..579b3f6e
--- /dev/null
+++ b/model/utils.py
@@ -0,0 +1,4 @@
+from transformers import AutoTokenizer
+
+
+def update_galactica_tokenizer():
\ No newline at end of file

From bcd5c52b3b370a217042b2ccb1983e113ecf6193 Mon Sep 17 00:00:00 2001
From: theblackcat102 <theblackcat102@github.com>
Date: Sat, 31 Dec 2022 03:02:10 +0000
Subject: [PATCH 02/41] [feature] working trainer code

---
 .vscode/settings.json                         |   2 +-
 .../reward/instructor/experimental_dataset.py |  10 +-
 model/reward/instructor/rank_datasets.py      |  49 ++++++---
 model/reward/instructor/tests/test_dataset.py |  10 +-
 model/reward/instructor/trainer.py            | 104 +++++++++++++++++-
 model/reward/instructor/utils.py              |  23 ++++
 6 files changed, 174 insertions(+), 24 deletions(-)

diff --git a/.vscode/settings.json b/.vscode/settings.json
index 56a51f78..4c58a32f 100644
--- a/.vscode/settings.json
+++ b/.vscode/settings.json
@@ -1,4 +1,4 @@
 {
-  "python.formatting.provider": "black",
+  "python.formatting.provider": "autopep8",
   "python.analysis.extraPaths": ["${workspaceFolder}/oasst-shared"]
 }
diff --git a/model/reward/instructor/experimental_dataset.py b/model/reward/instructor/experimental_dataset.py
index 145588c4..f705ccf6 100644
--- a/model/reward/instructor/experimental_dataset.py
+++ b/model/reward/instructor/experimental_dataset.py
@@ -1,5 +1,11 @@
 '''
-    
+    HFSummary
+
+        I want to train a multi regression model on axis_evals dataset mainly we can estimate the score of these score
+
+         - {"overall": "6", "accuracy": "6", "coverage": "6", "coherence": "7"}
+
+        Should be better than just a preference score
 
 '''
 import os
@@ -9,3 +15,5 @@ import torch
 import numpy as np
 from dataset import load_dataset
 from torch.utils.data import Dataset
+
+
diff --git a/model/reward/instructor/rank_datasets.py b/model/reward/instructor/rank_datasets.py
index 7fef5ab7..e407b30f 100644
--- a/model/reward/instructor/rank_datasets.py
+++ b/model/reward/instructor/rank_datasets.py
@@ -8,32 +8,51 @@
     [ ] 
 
 '''
+from typing import Optional, Union
 import os
 import glob
 import json
+from dataclasses import dataclass
 import numpy as np
 from torch.utils.data import Dataset
+import torch
 from datasets import load_dataset
+from transformers.tokenization_utils_base import PreTrainedTokenizerBase, PaddingStrategy
 
-class CollateFN():
-    def __init__(self, tokenizer, max_length=400) -> None:
-        self.tokenizer = tokenizer
-        self.max_length = max_length
+@dataclass
+class DataCollatorForPairRank:
+    """
 
-    def __call__(self, batch):
-        prompts = []
-        pos_sentences = []
-        neg_sentences = []
-        for prompt, pairs in batch:
+    Data collator that will dynamically pad the inputs for multiple choice received.
+
+    """
+    tokenizer: PreTrainedTokenizerBase
+    num_choices: int = 2
+    padding: Union[bool, str, PaddingStrategy] = True
+    max_length: Optional[int] = None
+    pad_to_multiple_of: Optional[int] = None
+
+    def __call__(self, features):
+
+        flatten_features = []
+        batch_size = 0
+        for question, pairs in features:
             for (pos, neg) in pairs:
-                prompts.append(prompt)
-                pos_sentences.append(pos)
-                neg_sentences.append(neg)
-
-        batch = [self.tokenizer(prompts, pos_sentences, return_tensors='pt', max_length=self.max_length, padding=True, truncation=True),\
-                self.tokenizer(prompts, neg_sentences, return_tensors='pt', max_length=self.max_length, padding=True, truncation=True)]
+                flatten_features.append(self.tokenizer(question, pos, truncation=True))
+                flatten_features.append(self.tokenizer(question, neg, truncation=True))
+                batch_size += 1
+        
+        batch = self.tokenizer.pad(
+            flatten_features,
+            padding=self.padding,
+            max_length=self.max_length,
+            pad_to_multiple_of=self.pad_to_multiple_of,
+            return_tensors="pt",
+        )
+        # batch = {k: v.view(batch_size, self.num_choices, -1) for k, v in batch.items()}
         return batch
 
+
 class WebGPT(Dataset):
 
     def __init__(self) -> None:
diff --git a/model/reward/instructor/tests/test_dataset.py b/model/reward/instructor/tests/test_dataset.py
index 4dd59c16..c452786b 100644
--- a/model/reward/instructor/tests/test_dataset.py
+++ b/model/reward/instructor/tests/test_dataset.py
@@ -1,26 +1,26 @@
 from transformers import AutoTokenizer
 from torch.utils.data import DataLoader
-from rank_datasets import WebGPT, HFSummary, CollateFN
+from rank_datasets import WebGPT, HFSummary, DataCollatorForMultipleChoice
 
 
 def test_hfsummary():
     
     tokenizer = AutoTokenizer.from_pretrained("bigscience/mt0-large")
-    collate_fn = CollateFN(tokenizer)
+    collate_fn = DataCollatorForMultipleChoice(tokenizer, max_length=200)
     dataset = HFSummary()
     dataloader = DataLoader(dataset, collate_fn=collate_fn, batch_size=8)
     for batch in dataloader:
-        print(batch[0]['input_ids'].shape)
+        print(batch['input_ids'].shape)
  
 
 def test_webgpt():
     
     tokenizer = AutoTokenizer.from_pretrained("bigscience/mt0-large")
-    collate_fn = CollateFN(tokenizer)
+    collate_fn = DataCollatorForMultipleChoice(tokenizer, max_length=200)
     dataset = WebGPT()
     dataloader = DataLoader(dataset, collate_fn=collate_fn, batch_size=32)
     for batch in dataloader:
-        print(batch[0]['input_ids'].shape)
+        print(batch['input_ids'].shape)
 
 
 if __name__ == "__main__":
diff --git a/model/reward/instructor/trainer.py b/model/reward/instructor/trainer.py
index 9ee5e043..43a5f8ef 100644
--- a/model/reward/instructor/trainer.py
+++ b/model/reward/instructor/trainer.py
@@ -1,2 +1,102 @@
-import wandb
-from accelerate import Accelerator
+from typing import Callable, List, Optional, Tuple, Union, Dict
+import torch
+from torch import nn
+import numpy as np
+import evaluate
+from dataclasses import dataclass
+from torch.utils.data import Dataset
+from transformers import AutoModelForSequenceClassification, AutoModelForMultipleChoice
+from transformers import Trainer, PreTrainedModel, TrainingArguments, DataCollator, EvalPrediction, TrainerCallback, PreTrainedTokenizerBase
+from rank_datasets import DataCollatorForPairRank, WebGPT
+from utils import get_tokenizer, train_val_dataset
+
+accuracy = evaluate.load("accuracy")
+
+@dataclass
+class CustomTrainingArguments(TrainingArguments):
+    loss_function: str='rank'
+
+
+def compute_metrics(eval_pred):
+    predictions, _ = eval_pred
+    predictions = np.argmax(predictions, axis=1)
+    return accuracy.compute(predictions=predictions, references=[0]*predictions.shape[0])
+
+class RankLoss(nn.Module):
+    def __init__(self, eps=1e-8) -> None:
+        super().__init__()
+        self.eps = eps
+        self.log_sigmoid = nn.LogSigmoid()
+
+    def forward(self, pos, neg):
+        return -self.log_sigmoid(pos - neg + self.eps).mean()
+
+
+class RankTrainer(Trainer):
+    def __init__(self, model: Union[PreTrainedModel, nn.Module] = None,
+                 args: TrainingArguments = None,
+                 data_collator: Optional[DataCollator] = None,
+                 train_dataset: Optional[Dataset] = None,
+                 eval_dataset: Optional[Dataset] = None,
+                 tokenizer: Optional[PreTrainedTokenizerBase] = None,
+                 model_init: Callable[[], PreTrainedModel] = None,
+                 compute_metrics: Optional[Callable[[EvalPrediction], Dict]] = None,
+                 callbacks: Optional[List[TrainerCallback]] = None,
+                 optimizers: Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None),
+                 preprocess_logits_for_metrics: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] = None):
+        super().__init__(model, args, data_collator, train_dataset, eval_dataset, tokenizer,
+                         model_init, compute_metrics, callbacks, optimizers, preprocess_logits_for_metrics)
+        self.loss_fct = RankLoss() if args.loss_function == 'rank' else nn.CrossEntropyLoss()
+        self.loss_function = args.loss_function
+
+    def compute_loss(self, model, inputs, return_outputs=False):
+        # forward pass
+        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, outputs) if return_outputs else loss
+
+
+if __name__ == "__main__":
+    model_name = 'bigscience/bloomz-560m'
+    model_name = 'google/electra-base-discriminator'
+    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,
+        num_train_epochs=4,
+        warmup_steps=500,
+        learning_rate=3e-5,
+        # half_precision_backend="apex",
+        gradient_checkpointing=False,
+        gradient_accumulation_steps=6,
+        per_device_train_batch_size=12,
+        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'
+    )
+    dataset = WebGPT()
+    train, eval = train_val_dataset(dataset)
+    collate_fn = DataCollatorForPairRank(tokenizer, max_length=400)
+    trainer = RankTrainer(
+        model,
+        args,
+        train_dataset=train,
+        eval_dataset=eval,
+        data_collator=collate_fn,
+        tokenizer=tokenizer
+    )
+    trainer.train()
diff --git a/model/reward/instructor/utils.py b/model/reward/instructor/utils.py
index 1487947c..10f84193 100644
--- a/model/reward/instructor/utils.py
+++ b/model/reward/instructor/utils.py
@@ -1,4 +1,7 @@
 import re
+from torch.utils.data import Subset
+from sklearn.model_selection import train_test_split
+from transformers import AutoTokenizer
 
 re_reference_remove = re.compile(r'\[([0-9])+\]|\[([0-9])+,([0-9])+\]')
 
@@ -16,3 +19,23 @@ def webgpt_return_format(row):
                  'pos': re_reference_remove.sub('', row['answer_1']),
                  'neg': re_reference_remove.sub('', row['answer_0'])
             }
+
+
+def get_tokenizer(tokenizer_name):
+    tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
+    if 'galactica' in tokenizer_name:
+        tokenizer.add_special_tokens({'pad_token':'<pad>', 'eos_token': '</s>' })
+
+    return tokenizer
+
+
+
+def train_val_dataset(dataset, val_split=0.2):
+    train_idx, val_idx = train_test_split(list(range(len(dataset))), 
+        test_size=val_split, random_state=666, shuffle=True)
+    # [3879, 11479, 8341, 9177, 10798, 18177, 5735, 15669, 4837, 2760]
+    print(val_idx[:10])
+    # [13582, 5919, 11875, 7373, 19135, 13706, 8555, 15788, 15005, 15209]
+    print(train_idx[:10])
+    return Subset(dataset, train_idx), Subset(dataset, val_idx)
+

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 03/41] [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()

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 04/41] [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,

From d2572d032301cff6c4304fd54952d2f49fe1eecd Mon Sep 17 00:00:00 2001
From: theblackcat102 <theblackcat102@github.com>
Date: Sat, 31 Dec 2022 09:42:49 +0000
Subject: [PATCH 05/41] [fix] Add drop_token_type to use galactica

---
 model/reward/instructor/rank_datasets.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/model/reward/instructor/rank_datasets.py b/model/reward/instructor/rank_datasets.py
index 128baafe..41740dcf 100644
--- a/model/reward/instructor/rank_datasets.py
+++ b/model/reward/instructor/rank_datasets.py
@@ -31,6 +31,7 @@ class DataCollatorForPairRank:
     padding: Union[bool, str, PaddingStrategy] = True
     max_length: Optional[int] = None
     pad_to_multiple_of: Optional[int] = None
+    drop_token_type: bool = False
 
     def __call__(self, features):
 
@@ -51,6 +52,8 @@ class DataCollatorForPairRank:
             pad_to_multiple_of=self.pad_to_multiple_of,
             return_tensors="pt",
         )
+        if self.drop_token_type:
+            batch.pop('token_type_ids')
         # batch = {k: v.view(batch_size, self.num_choices, -1) for k, v in batch.items()}
         return batch
 

From f3c299757d89fc6913996d852e3e8563ae61b5cf Mon Sep 17 00:00:00 2001
From: theblackcat102 <theblackcat102@github.com>
Date: Sat, 31 Dec 2022 17:02:46 +0000
Subject: [PATCH 06/41] [feature] added configs argument for parameters
 training and recording

---
 model/reward/instructor/README.md        |  3 ++
 model/reward/instructor/rank_datasets.py |  2 -
 model/reward/instructor/trainer.py       | 52 ++++++++++++++++--------
 model/reward/instructor/utils.py         | 38 +++++++++++++++++
 4 files changed, 76 insertions(+), 19 deletions(-)

diff --git a/model/reward/instructor/README.md b/model/reward/instructor/README.md
index 7dbfefbc..a8b5ef33 100644
--- a/model/reward/instructor/README.md
+++ b/model/reward/instructor/README.md
@@ -1,5 +1,8 @@
+# Sections to train Reward Model (RM)
 
 
+Currently we format 
+
 
 ```bash
 
diff --git a/model/reward/instructor/rank_datasets.py b/model/reward/instructor/rank_datasets.py
index 41740dcf..aa77089c 100644
--- a/model/reward/instructor/rank_datasets.py
+++ b/model/reward/instructor/rank_datasets.py
@@ -9,13 +9,11 @@
 
 '''
 from typing import Optional, Union
-import os
 import glob
 import json
 from dataclasses import dataclass
 import numpy as np
 from torch.utils.data import Dataset
-import torch
 from datasets import load_dataset
 from transformers.tokenization_utils_base import PreTrainedTokenizerBase, PaddingStrategy
 
diff --git a/model/reward/instructor/trainer.py b/model/reward/instructor/trainer.py
index 586c8d47..06bb8098 100644
--- a/model/reward/instructor/trainer.py
+++ b/model/reward/instructor/trainer.py
@@ -1,18 +1,22 @@
 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
+import yaml
 import evaluate
+from typing import Any, Callable, List, Optional, Tuple, Union, Dict
+from torch import nn
+from argparse import ArgumentParser
+import numpy as np
 from dataclasses import dataclass
 from torch.utils.data import Dataset, ConcatDataset
-from transformers import AutoModelForSequenceClassification, AutoModelForMultipleChoice
+from transformers import AutoModelForSequenceClassification
 from transformers import Trainer, PreTrainedModel, TrainingArguments, DataCollator, EvalPrediction, TrainerCallback, PreTrainedTokenizerBase
 from rank_datasets import DataCollatorForPairRank, WebGPT, HFSummary
-from utils import get_tokenizer, train_val_dataset
+from utils import get_tokenizer, train_val_dataset, freeze_top_n_layers, argument_parsing
 
 accuracy = evaluate.load("accuracy")
+parser = ArgumentParser()
+parser.add_argument('config', type=str)
 
 @dataclass
 class CustomTrainingArguments(TrainingArguments):
@@ -87,21 +91,26 @@ class RankTrainer(Trainer):
         return (loss, logits, labels)
 
 if __name__ == "__main__":
-    model_name = 'bigscience/bloomz-560m'
-    model_name = 'google/electra-large-discriminator'
+    training_conf = argument_parsing(parser)
+    
+    model_name = training_conf['model_name']
     model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=1, problem_type='regression')
+    if 'freeze_layer' in training_conf:
+        num_layer = training_conf['freeze_layer']
+        model = freeze_top_n_layers(model, num_layer)
+
     tokenizer = get_tokenizer(model_name)
     args = CustomTrainingArguments(
         output_dir=f"{model_name}-finetuned",
-        num_train_epochs=4,
+        num_train_epochs=training_conf['num_train_epochs'],
         warmup_steps=500,
-        loss_function='rank',
-        learning_rate=3e-5,
+        loss_function=training_conf['loss'],
+        learning_rate=training_conf['learning_rate'],
         # half_precision_backend="apex",
         fp16=True,
-        gradient_checkpointing=True,
-        gradient_accumulation_steps=8,
-        per_device_train_batch_size=8,
+        gradient_checkpointing=training_conf['gradient_checkpointing'],
+        gradient_accumulation_steps=training_conf['gradient_checkpointing'],
+        per_device_train_batch_size=training_conf['per_device_train_batch_size'],
         per_device_eval_batch_size=5,
         weight_decay=0.01,
         max_grad_norm=2.0,
@@ -112,10 +121,19 @@ if __name__ == "__main__":
         save_steps=1000,
         report_to='wandb'
     )
-    dataset = WebGPT()
-    train, eval = train_val_dataset(dataset)
-    train = ConcatDataset([train, HFSummary()])
-    collate_fn = DataCollatorForPairRank(tokenizer, max_length=440)
+    train_datasets, evals = [], {}
+    if 'webgpt' in training_conf['datasets']:
+        web_dataset = WebGPT()
+        train, eval = train_val_dataset(web_dataset)
+        train_datasets.append(train)
+        evals['webgpt'] = eval
+    if 'hfsummary' in training_conf['datasets']:
+        summary_dataset = HFSummary()
+        sum_train, sum_eval = train_val_dataset(summary_dataset)
+        train_datasets.append(sum_train)
+        evals['hfsummary'] = sum_eval
+    
+    collate_fn = DataCollatorForPairRank(tokenizer, max_length=training_conf['max_length'])
     trainer = RankTrainer(
         model,
         args,
diff --git a/model/reward/instructor/utils.py b/model/reward/instructor/utils.py
index 10f84193..4867087c 100644
--- a/model/reward/instructor/utils.py
+++ b/model/reward/instructor/utils.py
@@ -1,4 +1,5 @@
 import re
+import yaml
 from torch.utils.data import Subset
 from sklearn.model_selection import train_test_split
 from transformers import AutoTokenizer
@@ -39,3 +40,40 @@ def train_val_dataset(dataset, val_split=0.2):
     print(train_idx[:10])
     return Subset(dataset, train_idx), Subset(dataset, val_idx)
 
+def freeze_top_n_layers(model, target_layers):
+    for name, param in model.name_parameters():
+        if 'embed' in name:
+            param.requires_grad = False
+        elif 'layer' in name:
+            tokens = name.split('.')
+            idx = 0
+            for token in tokens:
+                if 'layer' in token:
+                    break
+                idx += 1
+
+            layer_ = int(tokens[idx+1])
+            if layer_ < target_layers:
+                param.requires_grad = False
+    return model
+
+
+def argument_parsing(parser):
+    default_params = {
+        'num_train_epochs': 4,
+        'learning_rate': 3e-5,
+        'eval_steps': 500,
+        'loss': 'rank',
+        'max_length': 440,
+        'per_device_train_batch_size': 8,
+        'gradient_accumulation_steps': 8,
+        'gradient_checkpointing': False,
+        'datasets': ['webgpt']
+    }
+    args = parser.parse_args()
+    with open(args.config, 'r', encoding='utf-8') as f:
+        training_conf = yaml.safe_load(f.read())
+
+    return { **default_params, **training_conf }
+
+

From 24e06626f46e1f9a4bd4f112ac8c8af45556e866 Mon Sep 17 00:00:00 2001
From: theblackcat102 <theblackcat102@github.com>
Date: Sat, 31 Dec 2022 17:04:44 +0000
Subject: [PATCH 07/41] [fix] Fix missing configs

---
 model/reward/instructor/configs/electra-base-dis-webgpt.yml | 2 ++
 1 file changed, 2 insertions(+)
 create mode 100644 model/reward/instructor/configs/electra-base-dis-webgpt.yml

diff --git a/model/reward/instructor/configs/electra-base-dis-webgpt.yml b/model/reward/instructor/configs/electra-base-dis-webgpt.yml
new file mode 100644
index 00000000..5c02fab7
--- /dev/null
+++ b/model/reward/instructor/configs/electra-base-dis-webgpt.yml
@@ -0,0 +1,2 @@
+model_name: google/electra-base-discriminator
+learning_rate: 3e-5

From 918b7b7ec0446651cb724ee0909288d6a89ce71b Mon Sep 17 00:00:00 2001
From: theblackcat102 <ray.tam@ikala.tv>
Date: Sun, 1 Jan 2023 01:25:53 +0800
Subject: [PATCH 08/41] [feature] Add galactica training config

---
 model/reward/instructor/configs/galactica-125m.yml | 13 +++++++++++++
 model/reward/instructor/configs/galactica-1b.yml   |  8 ++++++++
 model/reward/instructor/trainer.py                 |  6 +++---
 model/reward/instructor/utils.py                   |  8 ++++++--
 4 files changed, 30 insertions(+), 5 deletions(-)
 create mode 100644 model/reward/instructor/configs/galactica-125m.yml
 create mode 100644 model/reward/instructor/configs/galactica-1b.yml

diff --git a/model/reward/instructor/configs/galactica-125m.yml b/model/reward/instructor/configs/galactica-125m.yml
new file mode 100644
index 00000000..55e093f5
--- /dev/null
+++ b/model/reward/instructor/configs/galactica-125m.yml
@@ -0,0 +1,13 @@
+model_name: facebook/galactica-125m
+learning_rate: 1e-5
+gradient_checkpointing: false
+gradient_accumulation_steps: 32
+per_device_train_batch_size: 2
+warmup_steps: 600
+eval_steps: 200
+save_steps: 500
+max_length: 512
+num_train_epochs: 2
+datasets:
+  - webgpt
+  - hfsummary
\ No newline at end of file
diff --git a/model/reward/instructor/configs/galactica-1b.yml b/model/reward/instructor/configs/galactica-1b.yml
new file mode 100644
index 00000000..48ad439b
--- /dev/null
+++ b/model/reward/instructor/configs/galactica-1b.yml
@@ -0,0 +1,8 @@
+model_name: facebook/galactica-1.3b
+learning_rate: 6e-6
+gradient_checkpointing: false
+gradient_accumulation_steps: 16
+per_device_train_batch_size: 4
+warmup_steps: 600
+eval_steps: 200
+save_steps: 500
\ No newline at end of file
diff --git a/model/reward/instructor/trainer.py b/model/reward/instructor/trainer.py
index 06bb8098..dbdd91ba 100644
--- a/model/reward/instructor/trainer.py
+++ b/model/reward/instructor/trainer.py
@@ -109,7 +109,7 @@ if __name__ == "__main__":
         # half_precision_backend="apex",
         fp16=True,
         gradient_checkpointing=training_conf['gradient_checkpointing'],
-        gradient_accumulation_steps=training_conf['gradient_checkpointing'],
+        gradient_accumulation_steps=training_conf['gradient_accumulation_steps'],
         per_device_train_batch_size=training_conf['per_device_train_batch_size'],
         per_device_eval_batch_size=5,
         weight_decay=0.01,
@@ -132,8 +132,8 @@ if __name__ == "__main__":
         sum_train, sum_eval = train_val_dataset(summary_dataset)
         train_datasets.append(sum_train)
         evals['hfsummary'] = sum_eval
-    
-    collate_fn = DataCollatorForPairRank(tokenizer, max_length=training_conf['max_length'])
+    train = ConcatDataset(train_datasets)
+    collate_fn = DataCollatorForPairRank(tokenizer, max_length=training_conf['max_length'], drop_token_type= 'galactica' in model_name)
     trainer = RankTrainer(
         model,
         args,
diff --git a/model/reward/instructor/utils.py b/model/reward/instructor/utils.py
index 4867087c..733e6ea7 100644
--- a/model/reward/instructor/utils.py
+++ b/model/reward/instructor/utils.py
@@ -74,6 +74,10 @@ def argument_parsing(parser):
     with open(args.config, 'r', encoding='utf-8') as f:
         training_conf = yaml.safe_load(f.read())
 
-    return { **default_params, **training_conf }
-
+    params = { **default_params, **training_conf }
+    params['gradient_accumulation_steps'] = int(params['gradient_accumulation_steps'])
+    params['num_train_epochs'] = int(params['num_train_epochs'])
+    params['per_device_train_batch_size'] = int(params['per_device_train_batch_size'])
+    params['learning_rate'] = float(params['learning_rate'])
+    return params
 

From ba336fb087d10892b47133fdbee49846e6759db4 Mon Sep 17 00:00:00 2001
From: theblackcat102 <theblackcat102@github.com>
Date: Sat, 31 Dec 2022 17:43:27 +0000
Subject: [PATCH 09/41] [fix] fix freeze top N layers

---
 model/reward/instructor/configs/galactica-1b.yml | 10 ++++++++--
 model/reward/instructor/trainer.py               |  3 +++
 model/reward/instructor/utils.py                 |  7 ++++---
 3 files changed, 15 insertions(+), 5 deletions(-)

diff --git a/model/reward/instructor/configs/galactica-1b.yml b/model/reward/instructor/configs/galactica-1b.yml
index 48ad439b..5a094520 100644
--- a/model/reward/instructor/configs/galactica-1b.yml
+++ b/model/reward/instructor/configs/galactica-1b.yml
@@ -2,7 +2,13 @@ model_name: facebook/galactica-1.3b
 learning_rate: 6e-6
 gradient_checkpointing: false
 gradient_accumulation_steps: 16
-per_device_train_batch_size: 4
+per_device_train_batch_size: 2
 warmup_steps: 600
+freeze_layer: 20
 eval_steps: 200
-save_steps: 500
\ No newline at end of file
+save_steps: 500
+max_length: 400
+num_train_epochs: 2
+datasets:
+  - webgpt
+  - hfsummary
\ No newline at end of file
diff --git a/model/reward/instructor/trainer.py b/model/reward/instructor/trainer.py
index dbdd91ba..22baf130 100644
--- a/model/reward/instructor/trainer.py
+++ b/model/reward/instructor/trainer.py
@@ -98,6 +98,9 @@ if __name__ == "__main__":
     if 'freeze_layer' in training_conf:
         num_layer = training_conf['freeze_layer']
         model = freeze_top_n_layers(model, num_layer)
+        model_parameters = filter(lambda p: p.requires_grad, model.parameters())
+        params = sum([np.prod(p.size()) for p in model_parameters])
+        print('Number of trainable : {}M'.format(int(params/1e6)))
 
     tokenizer = get_tokenizer(model_name)
     args = CustomTrainingArguments(
diff --git a/model/reward/instructor/utils.py b/model/reward/instructor/utils.py
index 733e6ea7..ef3ed98d 100644
--- a/model/reward/instructor/utils.py
+++ b/model/reward/instructor/utils.py
@@ -41,23 +41,24 @@ def train_val_dataset(dataset, val_split=0.2):
     return Subset(dataset, train_idx), Subset(dataset, val_idx)
 
 def freeze_top_n_layers(model, target_layers):
-    for name, param in model.name_parameters():
+    for name, param in model.named_parameters():
         if 'embed' in name:
             param.requires_grad = False
-        elif 'layer' in name:
+        elif '.layer' in name:
             tokens = name.split('.')
             idx = 0
             for token in tokens:
                 if 'layer' in token:
                     break
                 idx += 1
+            if idx >= len(tokens):
+                continue
 
             layer_ = int(tokens[idx+1])
             if layer_ < target_layers:
                 param.requires_grad = False
     return model
 
-
 def argument_parsing(parser):
     default_params = {
         'num_train_epochs': 4,

From c5b31d0b9e268cebd7b1f3ab8a5327541d8e6dd2 Mon Sep 17 00:00:00 2001
From: theblackcat102 <theblackcat102@github.com>
Date: Sat, 31 Dec 2022 18:20:41 +0000
Subject: [PATCH 10/41] [feature] update reamde

---
 model/reward/instructor/README.md        |  25 +++-
 model/reward/instructor/requirements.txt | 140 +++++++++++++++++++++++
 2 files changed, 163 insertions(+), 2 deletions(-)
 create mode 100644 model/reward/instructor/requirements.txt

diff --git a/model/reward/instructor/README.md b/model/reward/instructor/README.md
index a8b5ef33..29716dca 100644
--- a/model/reward/instructor/README.md
+++ b/model/reward/instructor/README.md
@@ -1,10 +1,31 @@
 # Sections to train Reward Model (RM)
 
+Trainer code based on huggingface. Should be compatible with deepspeed or accelerate
 
-Currently we format 
+
+
+Requirements
+
+```
+wandb
+evaluate
+datasets
+transformers
+torch==1.12
+```
+
+To train your model run this
 
 
 ```bash
+python trainer.py configs/electra-base-dis-webgpt.yml
+```
+
+
+## Dataset
+
+For now we only supports webgpt and summary dataset from OpenAI. Once open-asisstant dataset are available it will be added here.
+
+
 
 
-```
\ No newline at end of file
diff --git a/model/reward/instructor/requirements.txt b/model/reward/instructor/requirements.txt
new file mode 100644
index 00000000..9fc45917
--- /dev/null
+++ b/model/reward/instructor/requirements.txt
@@ -0,0 +1,140 @@
+aiohttp==3.8.3
+aiosignal==1.3.1
+anyio==3.6.2
+argon2-cffi==21.3.0
+argon2-cffi-bindings==21.2.0
+arrow==1.2.3
+asttokens==2.2.1
+async-timeout==4.0.2
+attrs==22.2.0
+autopep8==2.0.1
+backcall==0.2.0
+beautifulsoup4==4.11.1
+bleach==5.0.1
+certifi==2022.12.7
+cffi==1.15.1
+charset-normalizer==2.1.1
+click==8.1.3
+comm==0.1.2
+datasets==2.8.0
+debugpy==1.6.4
+decorator==5.1.1
+defusedxml==0.7.1
+dill==0.3.6
+docker-pycreds==0.4.0
+entrypoints==0.4
+evaluate==0.4.0
+exceptiongroup==1.1.0
+executing==1.2.0
+fastjsonschema==2.16.2
+filelock==3.9.0
+fqdn==1.5.1
+frozenlist==1.3.3
+fsspec==2022.11.0
+gitdb==4.0.10
+GitPython==3.1.30
+huggingface-hub==0.11.1
+idna==3.4
+iniconfig==1.1.1
+ipykernel==6.19.4
+ipython==8.7.0
+ipython-genutils==0.2.0
+ipywidgets==8.0.4
+isoduration==20.11.0
+jedi==0.18.2
+Jinja2==3.1.2
+joblib==1.2.0
+jsonpointer==2.3
+jsonschema==4.17.3
+jupyter==1.0.0
+jupyter-console==6.4.4
+jupyter-events==0.5.0
+jupyter_client==7.4.8
+jupyter_core==5.1.1
+jupyter_server==2.0.6
+jupyter_server_terminals==0.4.3
+jupyterlab-pygments==0.2.2
+jupyterlab-widgets==3.0.5
+lightning-utilities==0.5.0
+MarkupSafe==2.1.1
+matplotlib-inline==0.1.6
+mistune==2.0.4
+multidict==6.0.4
+multiprocess==0.70.14
+nbclassic==0.4.8
+nbclient==0.7.2
+nbconvert==7.2.7
+nbformat==5.7.1
+nest-asyncio==1.5.6
+notebook==6.5.2
+notebook_shim==0.2.2
+numpy==1.24.1
+packaging==22.0
+pandas==1.5.2
+pandocfilters==1.5.0
+parso==0.8.3
+pathtools==0.1.2
+pexpect==4.8.0
+pickleshare==0.7.5
+platformdirs==2.6.2
+pluggy==1.0.0
+prometheus-client==0.15.0
+promise==2.3
+prompt-toolkit==3.0.36
+protobuf==3.20.1
+psutil==5.9.4
+ptyprocess==0.7.0
+pure-eval==0.2.2
+pyarrow==10.0.1
+pycodestyle==2.10.0
+pycparser==2.21
+Pygments==2.13.0
+pyrsistent==0.19.3
+pytest==7.2.0
+python-dateutil==2.8.2
+python-json-logger==2.0.4
+pytorch-lightning==1.8.6
+pytz==2022.7
+PyYAML==6.0
+pyzmq==24.0.1
+qtconsole==5.4.0
+QtPy==2.3.0
+regex==2022.10.31
+requests==2.28.1
+responses==0.18.0
+rfc3339-validator==0.1.4
+rfc3986-validator==0.1.1
+scikit-learn==1.2.0
+scipy==1.9.3
+Send2Trash==1.8.0
+sentry-sdk==1.12.1
+setproctitle==1.3.2
+shortuuid==1.0.11
+six==1.16.0
+smmap==5.0.0
+sniffio==1.3.0
+soupsieve==2.3.2.post1
+stack-data==0.6.2
+tensorboardX==2.5.1
+terminado==0.17.1
+threadpoolctl==3.1.0
+tinycss2==1.2.1
+tokenizers==0.13.2
+tomli==2.0.1
+torch==1.12.1+cu116
+torchmetrics==0.11.0
+tornado==6.2
+tqdm==4.64.1
+traitlets==5.8.0
+transformers==4.25.1
+typing_extensions==4.4.0
+uri-template==1.2.0
+urllib3==1.26.13
+wandb==0.13.7
+wcwidth==0.2.5
+webcolors==1.12
+webencodings==0.5.1
+websocket-client==1.4.2
+widgetsnbextension==4.0.5
+xxhash==3.2.0
+yarl==1.8.2

From 8a42ed32950cb7028bb50531281f6c736b2ac4e3 Mon Sep 17 00:00:00 2001
From: mrcabbage972 <mayvic@gmail.com>
Date: Sat, 31 Dec 2022 16:44:04 -0500
Subject: [PATCH 11/41] Adding a file for listing relevant research papers

---
 docs/research/README.md | 21 +++++++++++++++++++++
 1 file changed, 21 insertions(+)
 create mode 100644 docs/research/README.md

diff --git a/docs/research/README.md b/docs/research/README.md
new file mode 100644
index 00000000..bf4461e7
--- /dev/null
+++ b/docs/research/README.md
@@ -0,0 +1,21 @@
+# Research
+This page lists research papers that are relevant to the project.
+
+## Automatically Generating Instruction Data for Training
+This line of work is about significantly reducing the need for manually annotated data for the purpose of training [instruction-aligned](https://openai.com/blog/instruction-following/) language models.
+### SELF-INSTRUCT: Aligning Language Model with Self Generated Instructions [[ArXiv](https://arxiv.org/pdf/2212.10560.pdf)], [[Github](https://github.com/yizhongw/self-instruct)].
+
+> We introduce SELF-INSTRUCT, a framework for improving the instruction-following capabilities of pretrained language models by bootstrapping off its own generations. 
+> Our pipeline generates instruction, input, and output samples from a language model, then prunes them before using them to finetune the original model. 
+> Applying our method to vanilla GPT3, we demonstrate a 33% absolute improvement over the original model on SuperNaturalInstructions, on par with the performance of InstructGPT-0011, which is trained with private user data and human annotations.
+
+
+### Tuning Language Models with (Almost) No Human Labor. [[ArXiv](https://arxiv.org/pdf/2212.09689.pdf)], [[Github](https://github.com/orhonovich/unnatural-instructions)].
+
+> In this work, we introduce
+Unnatural Instructions: a large dataset of creative and diverse instructions, collected with virtually no human labor. 
+> We collect 64,000 examples by prompting a language model with three seed examples of instructions and eliciting a fourth. 
+> This set is then expanded by prompting the model to rephrase each instruction, creating a total of approximately 240,000 examples of instructions, inputs, and outputs.
+> Experiments show that despite containing a fair amount of noise, training on Unnatural Instructions rivals the effectiveness of training
+on open-source manually-curated datasets, surpassing the performance of models such as
+T0++ and Tk-Instruct across various benchmarks.
\ No newline at end of file

From 35f4c2f0d900e5bb1bada6a8748b89f3a2b36367 Mon Sep 17 00:00:00 2001
From: mrcabbage972 <mayvic@gmail.com>
Date: Sat, 31 Dec 2022 19:28:40 -0500
Subject: [PATCH 12/41] Adding missing line break

---
 docs/research/README.md | 16 +++++++++-------
 1 file changed, 9 insertions(+), 7 deletions(-)

diff --git a/docs/research/README.md b/docs/research/README.md
index bf4461e7..498a858b 100644
--- a/docs/research/README.md
+++ b/docs/research/README.md
@@ -1,21 +1,23 @@
 # Research
+
 This page lists research papers that are relevant to the project.
 
 ## Automatically Generating Instruction Data for Training
+
 This line of work is about significantly reducing the need for manually annotated data for the purpose of training [instruction-aligned](https://openai.com/blog/instruction-following/) language models.
+
 ### SELF-INSTRUCT: Aligning Language Model with Self Generated Instructions [[ArXiv](https://arxiv.org/pdf/2212.10560.pdf)], [[Github](https://github.com/yizhongw/self-instruct)].
 
-> We introduce SELF-INSTRUCT, a framework for improving the instruction-following capabilities of pretrained language models by bootstrapping off its own generations. 
-> Our pipeline generates instruction, input, and output samples from a language model, then prunes them before using them to finetune the original model. 
+> We introduce SELF-INSTRUCT, a framework for improving the instruction-following capabilities of pretrained language models by bootstrapping off its own generations.
+> Our pipeline generates instruction, input, and output samples from a language model, then prunes them before using them to finetune the original model.
 > Applying our method to vanilla GPT3, we demonstrate a 33% absolute improvement over the original model on SuperNaturalInstructions, on par with the performance of InstructGPT-0011, which is trained with private user data and human annotations.
 
-
 ### Tuning Language Models with (Almost) No Human Labor. [[ArXiv](https://arxiv.org/pdf/2212.09689.pdf)], [[Github](https://github.com/orhonovich/unnatural-instructions)].
 
 > In this work, we introduce
-Unnatural Instructions: a large dataset of creative and diverse instructions, collected with virtually no human labor. 
-> We collect 64,000 examples by prompting a language model with three seed examples of instructions and eliciting a fourth. 
+> Unnatural Instructions: a large dataset of creative and diverse instructions, collected with virtually no human labor.
+> We collect 64,000 examples by prompting a language model with three seed examples of instructions and eliciting a fourth.
 > This set is then expanded by prompting the model to rephrase each instruction, creating a total of approximately 240,000 examples of instructions, inputs, and outputs.
 > Experiments show that despite containing a fair amount of noise, training on Unnatural Instructions rivals the effectiveness of training
-on open-source manually-curated datasets, surpassing the performance of models such as
-T0++ and Tk-Instruct across various benchmarks.
\ No newline at end of file
+> on open-source manually-curated datasets, surpassing the performance of models such as
+> T0++ and Tk-Instruct across various benchmarks.

From 0119ee666b64b7de779d440976ec367e688a1594 Mon Sep 17 00:00:00 2001
From: theblackcat102 <theblackcat102@github.com>
Date: Sun, 1 Jan 2023 02:09:21 +0000
Subject: [PATCH 13/41] [feature] Add support for bloomz

---
 model/reward/instructor/README.md                 | 15 ++++++++++++---
 model/reward/instructor/configs/bloomz-560m.yml   | 10 ++++++++++
 .../configs/electra-base-dis-webgpt.yml           |  3 ++-
 model/reward/instructor/rank_datasets.py          |  8 +++++++-
 model/reward/instructor/utils.py                  | 15 +++++++++++++--
 5 files changed, 44 insertions(+), 7 deletions(-)
 create mode 100644 model/reward/instructor/configs/bloomz-560m.yml

diff --git a/model/reward/instructor/README.md b/model/reward/instructor/README.md
index 29716dca..5992dbc0 100644
--- a/model/reward/instructor/README.md
+++ b/model/reward/instructor/README.md
@@ -1,7 +1,6 @@
 # Sections to train Reward Model (RM)
 
-Trainer code based on huggingface. Should be compatible with deepspeed or accelerate
-
+Trainer code based on huggingface. Compatible with deepspeed or accelerate
 
 
 Requirements
@@ -14,7 +13,7 @@ transformers
 torch==1.12
 ```
 
-To train your model run this
+Start training
 
 
 ```bash
@@ -26,6 +25,16 @@ python trainer.py configs/electra-base-dis-webgpt.yml
 
 For now we only supports webgpt and summary dataset from OpenAI. Once open-asisstant dataset are available it will be added here.
 
+## Model
 
+Check out configs
 
+```
+Open-Assistant/model/reward/instructor/configs/
+    bloomz-560m.yml
+    electra-base-dis-webgpt.yml
+    galactica-125m.yml
+    galactica-1b.yml
+```
 
+You can add new huggingface model as you want.
diff --git a/model/reward/instructor/configs/bloomz-560m.yml b/model/reward/instructor/configs/bloomz-560m.yml
new file mode 100644
index 00000000..c8f55746
--- /dev/null
+++ b/model/reward/instructor/configs/bloomz-560m.yml
@@ -0,0 +1,10 @@
+model_name: bigscience/bloomz-560m
+learning_rate: 3e-5
+gradient_accumulation_steps: 16
+per_device_train_batch_size: 2
+max_length: 600
+freeze_layer: 12
+num_train_epochs: 2
+datasets:
+  - webgpt
+  - hfsummary
\ No newline at end of file
diff --git a/model/reward/instructor/configs/electra-base-dis-webgpt.yml b/model/reward/instructor/configs/electra-base-dis-webgpt.yml
index 5c02fab7..fc168b63 100644
--- a/model/reward/instructor/configs/electra-base-dis-webgpt.yml
+++ b/model/reward/instructor/configs/electra-base-dis-webgpt.yml
@@ -1,2 +1,3 @@
-model_name: google/electra-base-discriminator
+model_name: google/electra-large-discriminator
 learning_rate: 3e-5
+max_length: 300
\ No newline at end of file
diff --git a/model/reward/instructor/rank_datasets.py b/model/reward/instructor/rank_datasets.py
index aa77089c..3d122915 100644
--- a/model/reward/instructor/rank_datasets.py
+++ b/model/reward/instructor/rank_datasets.py
@@ -1,6 +1,12 @@
 '''
     author: theblackcat102
 
+    Dataset output format from __getitem__
+
+     - question / prompt : string
+
+     - answers / rows : list of tuple pair. The first element in the tuple pair must be the positive pair (rank higher than the second element)
+
     A list of rank based dataset for training using rank loss
 
     Some nice features to have
@@ -105,7 +111,7 @@ class HFSummary(Dataset):
 
             >> azcopy copy "https://openaipublic.blob.core.windows.net/summarize-from-feedback/dataset/*" . --recursive
         
-        choice : 0 or 1
+        labeling method : pair comparison, 0 or 1
 
     '''
     def __init__(self, split='train',
diff --git a/model/reward/instructor/utils.py b/model/reward/instructor/utils.py
index ef3ed98d..f26add55 100644
--- a/model/reward/instructor/utils.py
+++ b/model/reward/instructor/utils.py
@@ -41,14 +41,16 @@ def train_val_dataset(dataset, val_split=0.2):
     return Subset(dataset, train_idx), Subset(dataset, val_idx)
 
 def freeze_top_n_layers(model, target_layers):
+    # its possible we can simply detect which module is a ModuleList
+    # and simply freeze the module without doing string parsing
     for name, param in model.named_parameters():
         if 'embed' in name:
             param.requires_grad = False
-        elif '.layer' in name:
+        elif '.layer' in name or '.h.' in name:
             tokens = name.split('.')
             idx = 0
             for token in tokens:
-                if 'layer' in token:
+                if 'layer' in token or token == 'h':
                     break
                 idx += 1
             if idx >= len(tokens):
@@ -56,6 +58,7 @@ def freeze_top_n_layers(model, target_layers):
 
             layer_ = int(tokens[idx+1])
             if layer_ < target_layers:
+                # print('freeze ', layer_, name)
                 param.requires_grad = False
     return model
 
@@ -82,3 +85,11 @@ def argument_parsing(parser):
     params['learning_rate'] = float(params['learning_rate'])
     return params
 
+
+
+if __name__ == "__main__":
+    from transformers import AutoModelForSequenceClassification
+
+    model = AutoModelForSequenceClassification.from_pretrained('bigscience/bloomz-560m')
+    freeze_top_n_layers(model, 10)
+    print(model.state_dict().keys())
\ No newline at end of file

From e27a3eb3c75e6b3193e712e3cfd76298e0dc6bc6 Mon Sep 17 00:00:00 2001
From: theblackcat102 <theblackcat102@github.com>
Date: Sun, 1 Jan 2023 02:22:57 +0000
Subject: [PATCH 14/41] [fix] Tidy up todo and trainer comments

---
 model/reward/instructor/TODO.md                     | 13 ++++++++++++-
 .../instructor/configs/bloomz-560m-summary.yml      |  9 +++++++++
 model/reward/instructor/trainer.py                  |  2 +-
 model/utils.py                                      |  4 ----
 4 files changed, 22 insertions(+), 6 deletions(-)
 create mode 100644 model/reward/instructor/configs/bloomz-560m-summary.yml
 delete mode 100644 model/utils.py

diff --git a/model/reward/instructor/TODO.md b/model/reward/instructor/TODO.md
index 33bc6595..ec23b7c3 100644
--- a/model/reward/instructor/TODO.md
+++ b/model/reward/instructor/TODO.md
@@ -1,12 +1,23 @@
 
 Some other reward features we can use
 
+0. Finish classifcation feature 
 
-Summaries from human feedback
+1. Summaries from human feedback
 
 * use `confidence` score into the RM learning, ensure the output rank score correlates with confidence
 
 * each labeling has a labeling `note`, basically comments by labeler, not sure what else we can use
 
+* Use the score for "overall", "accuracy", "coverage", "coherence" from axis/evals to train an addition model (rank additional aspect of the policy model)
+
+    * this should be placed under experimental_dataset.py
+
+
+2. Add support for anthropic dataset
+
+* anthropic dataset is more like a conversation tree which is much complex than simply question-answer schema
+
+    * this is basically a MCTS from alphazero.
 
 
diff --git a/model/reward/instructor/configs/bloomz-560m-summary.yml b/model/reward/instructor/configs/bloomz-560m-summary.yml
new file mode 100644
index 00000000..a02f4e4a
--- /dev/null
+++ b/model/reward/instructor/configs/bloomz-560m-summary.yml
@@ -0,0 +1,9 @@
+model_name: bigscience/bloomz-560m
+learning_rate: 3e-5
+gradient_accumulation_steps: 16
+per_device_train_batch_size: 2
+max_length: 600
+freeze_layer: 12
+num_train_epochs: 2
+datasets:
+  - hfsummary
\ No newline at end of file
diff --git a/model/reward/instructor/trainer.py b/model/reward/instructor/trainer.py
index 22baf130..de0b011a 100644
--- a/model/reward/instructor/trainer.py
+++ b/model/reward/instructor/trainer.py
@@ -92,7 +92,7 @@ class RankTrainer(Trainer):
 
 if __name__ == "__main__":
     training_conf = argument_parsing(parser)
-    
+
     model_name = training_conf['model_name']
     model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=1, problem_type='regression')
     if 'freeze_layer' in training_conf:
diff --git a/model/utils.py b/model/utils.py
deleted file mode 100644
index 579b3f6e..00000000
--- a/model/utils.py
+++ /dev/null
@@ -1,4 +0,0 @@
-from transformers import AutoTokenizer
-
-
-def update_galactica_tokenizer():
\ No newline at end of file

From a5a2625e2d15f327d89ee89708284971ba96e59f Mon Sep 17 00:00:00 2001
From: theblackcat102 <theblackcat102@github.com>
Date: Sun, 1 Jan 2023 02:55:54 +0000
Subject: [PATCH 15/41] [merge] most of the bugs should be fixed. #77

---
 model/reward/instructor/cls_dataset.py            | 15 ++++-----------
 .../test-galactica-125m-classification.yml        | 14 ++++++++++++++
 model/reward/instructor/rank_datasets.py          |  8 ++++++--
 model/reward/instructor/trainer.py                |  5 ++++-
 4 files changed, 28 insertions(+), 14 deletions(-)
 create mode 100644 model/reward/instructor/configs/test-galactica-125m-classification.yml

diff --git a/model/reward/instructor/cls_dataset.py b/model/reward/instructor/cls_dataset.py
index 54bbd19e..ff824d19 100644
--- a/model/reward/instructor/cls_dataset.py
+++ b/model/reward/instructor/cls_dataset.py
@@ -24,20 +24,10 @@ class WebGPTDataset(Dataset):
         '''
         os.makedirs('dataset', exist_ok=True)
         dataset = load_dataset("openai/webgpt_comparisons")
-        if os.path.exists(index_cache):
-            train_idx = torch.load(index_cache)
-        else:
-            train_idx = np.random.choice(range(len(dataset['train'])), int(len(dataset['train'])*0.8), replace=False)
-            torch.save(set(train_idx.tolist()), index_cache)
         self.dataset = []
         self.dataset_index = []
         for idx, row in enumerate(dataset['train']):
-            if mode == 'train' and idx in train_idx:
-                self.dataset.append(webgpt_return_format(row))
-                self.dataset_index.append(idx)
-            elif idx not in train_idx and mode != 'train':
-                self.dataset.append(webgpt_return_format(row))
-                self.dataset_index.append(idx)
+            self.dataset.append(webgpt_return_format(row))
 
         # since this dataset was generated from 176B GPT-3
         # we needed some more sample generated from the starting model
@@ -71,3 +61,6 @@ class WebGPTDataset(Dataset):
 
         gen_neg = random.choice(self.additional[self.dataset_index[index]])
         return row['question'], row['pos'], row['neg'], gen_neg
+
+
+
diff --git a/model/reward/instructor/configs/test-galactica-125m-classification.yml b/model/reward/instructor/configs/test-galactica-125m-classification.yml
new file mode 100644
index 00000000..1ad1f47c
--- /dev/null
+++ b/model/reward/instructor/configs/test-galactica-125m-classification.yml
@@ -0,0 +1,14 @@
+model_name: facebook/galactica-125m
+learning_rate: 1e-5
+gradient_checkpointing: false
+gradient_accumulation_steps: 10
+per_device_train_batch_size: 6
+warmup_steps: 600
+loss: cls
+eval_steps: 200
+save_steps: 500
+max_length: 128
+num_train_epochs: 2
+datasets:
+  - webgpt
+  - hfsummary
\ No newline at end of file
diff --git a/model/reward/instructor/rank_datasets.py b/model/reward/instructor/rank_datasets.py
index 3d122915..4ba6293c 100644
--- a/model/reward/instructor/rank_datasets.py
+++ b/model/reward/instructor/rank_datasets.py
@@ -11,7 +11,11 @@
 
     Some nice features to have
 
-    [ ] 
+    [] support additional negative samples generated from other models.
+
+        For example we can use galactica-125m to generate a TLDR and assume it was
+        inferior than the human perference one
+
 
 '''
 from typing import Optional, Union
@@ -35,7 +39,7 @@ class DataCollatorForPairRank:
     padding: Union[bool, str, PaddingStrategy] = True
     max_length: Optional[int] = None
     pad_to_multiple_of: Optional[int] = None
-    drop_token_type: bool = False
+    drop_token_type: bool = False # galactica
 
     def __call__(self, features):
 
diff --git a/model/reward/instructor/trainer.py b/model/reward/instructor/trainer.py
index de0b011a..48fc4e8d 100644
--- a/model/reward/instructor/trainer.py
+++ b/model/reward/instructor/trainer.py
@@ -77,7 +77,10 @@ class RankTrainer(Trainer):
 
         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]]:
+    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

From 4b7f1f25a138e614ab9f385f08913878a8a21bbb Mon Sep 17 00:00:00 2001
From: theblackcat102 <theblackcat102@github.com>
Date: Sun, 1 Jan 2023 03:07:40 +0000
Subject: [PATCH 16/41] [fix] Use official split for eval

---
 model/reward/instructor/trainer.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/model/reward/instructor/trainer.py b/model/reward/instructor/trainer.py
index 48fc4e8d..391464c6 100644
--- a/model/reward/instructor/trainer.py
+++ b/model/reward/instructor/trainer.py
@@ -134,12 +134,14 @@ if __name__ == "__main__":
         train_datasets.append(train)
         evals['webgpt'] = eval
     if 'hfsummary' in training_conf['datasets']:
-        summary_dataset = HFSummary()
-        sum_train, sum_eval = train_val_dataset(summary_dataset)
+        sum_train = HFSummary(split='train')
         train_datasets.append(sum_train)
+        sum_eval = HFSummary(split='valid1')
+        assert len(sum_eval) > 0
         evals['hfsummary'] = sum_eval
     train = ConcatDataset(train_datasets)
     collate_fn = DataCollatorForPairRank(tokenizer, max_length=training_conf['max_length'], drop_token_type= 'galactica' in model_name)
+    assert len(evals) > 0
     trainer = RankTrainer(
         model,
         args,

From 8b1553642f8f51e6d61f05a9a8c9302691d1ef25 Mon Sep 17 00:00:00 2001
From: theblackcat102 <theblackcat102@github.com>
Date: Sun, 1 Jan 2023 08:22:30 +0000
Subject: [PATCH 17/41] [feature] remove dependency to download hfsummary
 manually

---
 .../reward/instructor/experimental_dataset.py |  1 +
 model/reward/instructor/rank_datasets.py      | 58 +++++++++----------
 model/reward/instructor/tests/test_dataset.py |  9 +--
 3 files changed, 34 insertions(+), 34 deletions(-)

diff --git a/model/reward/instructor/experimental_dataset.py b/model/reward/instructor/experimental_dataset.py
index f705ccf6..85f0c899 100644
--- a/model/reward/instructor/experimental_dataset.py
+++ b/model/reward/instructor/experimental_dataset.py
@@ -17,3 +17,4 @@ from dataset import load_dataset
 from torch.utils.data import Dataset
 
 
+class 
\ No newline at end of file
diff --git a/model/reward/instructor/rank_datasets.py b/model/reward/instructor/rank_datasets.py
index 4ba6293c..2f2260c2 100644
--- a/model/reward/instructor/rank_datasets.py
+++ b/model/reward/instructor/rank_datasets.py
@@ -112,51 +112,49 @@ class HFSummary(Dataset):
     '''
         Human feedback data from OpenAI
         https://github.com/openai/summarize-from-feedback
-
-            >> azcopy copy "https://openaipublic.blob.core.windows.net/summarize-from-feedback/dataset/*" . --recursive
         
         labeling method : pair comparison, 0 or 1
 
     '''
     def __init__(self, split='train',
-        path='summarize-from-feedback/comparisons/*.json',
         conf_threshold=-1,
-        max_comparison_per_sample=5) -> None:
+        max_comparison_per_sample=3) -> None:
         super().__init__()
-        assert split in ('train', 'valid1', 'valid2', 'test')
+        assert split in ('train', 'validation')
         summaries = {}
         # using prompt as our index will allows us
         # to add additional generated prompt later
         self.index2summary = {}
         self.max_comparison_per_sample = max_comparison_per_sample
-        for jsonl_file in glob.glob(path):
-            with open(jsonl_file, 'r') as f:
-                for line in f:
-                    data = json.loads(line)
-                    if data['split'] != split:
-                        continue
-                    if 'extra' in data and \
-                        'confidence' in data['extra'] and \
-                        conf_threshold > data['extra']['confidence']:
-                        print('skipping {}'.format(data['info']['id']))
-                        continue
+        dataset = load_dataset('Tristan/summarize_from_feedback', 'comparisons')[split]
+        for data in dataset:
+            if 'extra' in data and \
+                'confidence' in data['extra'] and \
+                data['extra']['confidence'] is not None and \
+                conf_threshold > data['extra']['confidence']:
+                print('skipping {}'.format(data['info']['id']))
+                continue
 
-                    if 'article' in data['info']:
-                        context = data['info']['article']
-                    elif 'post' in data['info']:
-                        context = data['info']['post']
+            if 'article' in data['info'] and \
+                data['info']['article'] is not None:
+                context = data['info']['article']
+            elif 'post' in data['info']:
+                context = data['info']['post']
 
-                    if context not in self.index2summary:
-                        self.index2summary[len(self.index2summary)] = context
-                    
-                    if context not in summaries:
-                        summaries[context] = []
+            if context is None:
+                continue
 
-                    pos, neg = (0, 1) if data['choice'] == 0 else (1, 0)
-                    summaries[context].append((
-                        data['summaries'][pos]['text'],
-                        data['summaries'][neg]['text']
-                    ))
+            if context not in self.index2summary:
+                self.index2summary[len(self.index2summary)] = context
+            
+            if context not in summaries:
+                summaries[context] = []
+
+            pos, neg = (0, 1) if data['choice'] == 0 else (1, 0)
+            summaries[context].append((
+                data['summaries'][pos]['text'],
+                data['summaries'][neg]['text']
+            ))
 
         self.summaries = summaries
 
diff --git a/model/reward/instructor/tests/test_dataset.py b/model/reward/instructor/tests/test_dataset.py
index c452786b..7b432fd3 100644
--- a/model/reward/instructor/tests/test_dataset.py
+++ b/model/reward/instructor/tests/test_dataset.py
@@ -1,22 +1,23 @@
 from transformers import AutoTokenizer
 from torch.utils.data import DataLoader
-from rank_datasets import WebGPT, HFSummary, DataCollatorForMultipleChoice
+from rank_datasets import WebGPT, HFSummary, DataCollatorForPairRank
 
 
 def test_hfsummary():
     
     tokenizer = AutoTokenizer.from_pretrained("bigscience/mt0-large")
-    collate_fn = DataCollatorForMultipleChoice(tokenizer, max_length=200)
+    collate_fn = DataCollatorForPairRank(tokenizer, max_length=200)
     dataset = HFSummary()
+    print(len(dataset))
     dataloader = DataLoader(dataset, collate_fn=collate_fn, batch_size=8)
     for batch in dataloader:
-        print(batch['input_ids'].shape)
+        batch['input_ids'].shape
  
 
 def test_webgpt():
     
     tokenizer = AutoTokenizer.from_pretrained("bigscience/mt0-large")
-    collate_fn = DataCollatorForMultipleChoice(tokenizer, max_length=200)
+    collate_fn = DataCollatorForPairRank(tokenizer, max_length=200)
     dataset = WebGPT()
     dataloader = DataLoader(dataset, collate_fn=collate_fn, batch_size=32)
     for batch in dataloader:

From 1197dccf11cc1eac750109b6f856969ab4db8361 Mon Sep 17 00:00:00 2001
From: theblackcat102 <theblackcat102@github.com>
Date: Sun, 1 Jan 2023 08:25:09 +0000
Subject: [PATCH 18/41] [fix] dataset split name

---
 model/reward/instructor/rank_datasets.py      | 8 ++++++--
 model/reward/instructor/tests/test_dataset.py | 2 +-
 2 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/model/reward/instructor/rank_datasets.py b/model/reward/instructor/rank_datasets.py
index 2f2260c2..c2b7e58f 100644
--- a/model/reward/instructor/rank_datasets.py
+++ b/model/reward/instructor/rank_datasets.py
@@ -120,13 +120,14 @@ class HFSummary(Dataset):
         conf_threshold=-1,
         max_comparison_per_sample=3) -> None:
         super().__init__()
-        assert split in ('train', 'validation')
+        assert split in ('train', 'valid1', 'valid2', 'test')
         summaries = {}
         # using prompt as our index will allows us
         # to add additional generated prompt later
         self.index2summary = {}
         self.max_comparison_per_sample = max_comparison_per_sample
-        dataset = load_dataset('Tristan/summarize_from_feedback', 'comparisons')[split]
+        major_split = split if 'train' == split else 'validation'
+        dataset = load_dataset('Tristan/summarize_from_feedback', 'comparisons')[major_split]
         for data in dataset:
             if 'extra' in data and \
                 'confidence' in data['extra'] and \
@@ -135,6 +136,9 @@ class HFSummary(Dataset):
                 print('skipping {}'.format(data['info']['id']))
                 continue
 
+            if split != 'train' and split != data['split']:
+                continue
+
             if 'article' in data['info'] and \
                 data['info']['article'] is not None:
                 context = data['info']['article']
diff --git a/model/reward/instructor/tests/test_dataset.py b/model/reward/instructor/tests/test_dataset.py
index 7b432fd3..5765cd43 100644
--- a/model/reward/instructor/tests/test_dataset.py
+++ b/model/reward/instructor/tests/test_dataset.py
@@ -7,7 +7,7 @@ def test_hfsummary():
     
     tokenizer = AutoTokenizer.from_pretrained("bigscience/mt0-large")
     collate_fn = DataCollatorForPairRank(tokenizer, max_length=200)
-    dataset = HFSummary()
+    dataset = HFSummary('train')
     print(len(dataset))
     dataloader = DataLoader(dataset, collate_fn=collate_fn, batch_size=8)
     for batch in dataloader:

From 168e9ca6b3863fa7cf09691e04ce0a575c537bfd Mon Sep 17 00:00:00 2001
From: theblackcat102 <theblackcat102@github.com>
Date: Sun, 1 Jan 2023 10:19:25 +0000
Subject: [PATCH 19/41] [feature] added summary quality rater

---
 model/reward/instructor/README.md             |  18 ++-
 model/reward/instructor/TODO.md               |   2 +-
 .../reward/instructor/experimental_dataset.py |  90 +++++++++++-
 model/reward/instructor/rank_datasets.py      |   4 -
 .../instructor/summary_quality_trainer.py     | 132 ++++++++++++++++++
 model/reward/instructor/tests/test_dataset.py |  15 +-
 model/reward/instructor/trainer.py            |   4 +-
 model/reward/instructor/utils.py              |   1 +
 8 files changed, 251 insertions(+), 15 deletions(-)
 create mode 100644 model/reward/instructor/summary_quality_trainer.py

diff --git a/model/reward/instructor/README.md b/model/reward/instructor/README.md
index 5992dbc0..31c25371 100644
--- a/model/reward/instructor/README.md
+++ b/model/reward/instructor/README.md
@@ -13,7 +13,7 @@ transformers
 torch==1.12
 ```
 
-Start training
+Start training reward model
 
 
 ```bash
@@ -21,6 +21,22 @@ python trainer.py configs/electra-base-dis-webgpt.yml
 ```
 
 
+Additional axis labeling, this outputs a 4 summary quality evaluation metrics (score are normalized to 0-1 )
+
+```bash
+python summary_quality_trainer.py configs/test-bloomz-560m-quality.yml
+```
+
+The four summary are :
+
+* overall
+
+* accuracy
+
+* coverage
+
+* coherence
+
 ## Dataset
 
 For now we only supports webgpt and summary dataset from OpenAI. Once open-asisstant dataset are available it will be added here.
diff --git a/model/reward/instructor/TODO.md b/model/reward/instructor/TODO.md
index ec23b7c3..1e653922 100644
--- a/model/reward/instructor/TODO.md
+++ b/model/reward/instructor/TODO.md
@@ -9,7 +9,7 @@ Some other reward features we can use
 
 * each labeling has a labeling `note`, basically comments by labeler, not sure what else we can use
 
-* Use the score for "overall", "accuracy", "coverage", "coherence" from axis/evals to train an addition model (rank additional aspect of the policy model)
+* ~~Use the score for "overall", "accuracy", "coverage", "coherence" from axis/evals to train an addition model (rank additional aspect of the policy model)~~
 
     * this should be placed under experimental_dataset.py
 
diff --git a/model/reward/instructor/experimental_dataset.py b/model/reward/instructor/experimental_dataset.py
index 85f0c899..47d20d64 100644
--- a/model/reward/instructor/experimental_dataset.py
+++ b/model/reward/instructor/experimental_dataset.py
@@ -8,13 +8,93 @@
         Should be better than just a preference score
 
 '''
-import os
-import json
-import random
 import torch
+from typing import Optional, Union
 import numpy as np
-from dataset import load_dataset
+from collections import defaultdict
+from datasets import load_dataset
+from dataclasses import dataclass
 from torch.utils.data import Dataset
+from transformers.tokenization_utils_base import PreTrainedTokenizerBase, PaddingStrategy
+
+
+@dataclass
+class DataCollatorForSummaryScore:
+    """
+
+    Data collator that will dynamically pad the inputs for multiple choice received.
+
+    """
+    tokenizer: PreTrainedTokenizerBase
+    num_choices: int = 2
+    padding: Union[bool, str, PaddingStrategy] = True
+    max_length: Optional[int] = None
+    pad_to_multiple_of: Optional[int] = None
+    drop_token_type: bool = False # galactica
+
+    def __call__(self, batch):
+
+        features = []
+        labels = []
+        for feature, label in batch:
+            features.append(feature)
+            labels.append(label)
+
+        batch_feature = self.tokenizer.pad(
+            features,
+            padding=self.padding,
+            max_length=self.max_length,
+            pad_to_multiple_of=self.pad_to_multiple_of,
+            return_tensors="pt",
+        )
+        if self.drop_token_type:
+            batch_feature.pop('token_type_ids')
+        # batch = {k: v.view(batch_size, self.num_choices, -1) for k, v in batch.items()}
+        batch_feature['labels'] = torch.from_numpy(np.array(labels)).float()
+        return batch_feature
+
+
+class HFSummaryQuality(Dataset):
+    def __init__(self, split, tokenizer, max_length=300) -> None:
+        super().__init__()
+        assert split in ('validation', 'test')
+        dataset = load_dataset('Tristan/summarize_from_feedback', 'axis')[split]
+        self.max_length = max_length
+        mean_scores = defaultdict(list)
+        self.contexts = []
+        self.responses = []
+        self.labels = []
+        for data in dataset:
+
+            if 'article' in data['info'] and \
+                data['info']['article'] is not None:
+                context = data['info']['article']
+            elif 'post' in data['info']:
+                context = data['info']['post']
+            self.contexts.append(context)
+
+            response = data['summary']['text']
+            self.responses.append(response)
+            self.labels.append(data['summary']['axes'])
+            for axis, score in data['summary']['axes'].items():
+                if score is not None:
+                    mean_scores[axis].append(score)
+
+        self.label2idx = { key: idx for idx, key in enumerate(mean_scores.keys()) }
+        self.label2mean = { key: np.mean(scores) for key, scores in mean_scores.items() }
+        self.tokenizer = tokenizer
+        print(self.label2idx)
+
+    def __len__(self):
+        return len(self.responses)
+
+    def __getitem__(self, index):
+        context = self.contexts[index]
+        # return pairs of comparison
+        response = self.responses[index]
+        labels = np.zeros(len(self.label2idx))
+        for key, score in self.labels[index].items():
+            labels[self.label2idx[key]] = (self.label2mean[key] if score is None else score)/10
+        return self.tokenizer(context, response, truncation=True, max_length=self.max_length), labels
 
 
-class 
\ No newline at end of file
diff --git a/model/reward/instructor/rank_datasets.py b/model/reward/instructor/rank_datasets.py
index c2b7e58f..f38885e4 100644
--- a/model/reward/instructor/rank_datasets.py
+++ b/model/reward/instructor/rank_datasets.py
@@ -19,8 +19,6 @@
 
 '''
 from typing import Optional, Union
-import glob
-import json
 from dataclasses import dataclass
 import numpy as np
 from torch.utils.data import Dataset
@@ -145,8 +143,6 @@ class HFSummary(Dataset):
             elif 'post' in data['info']:
                 context = data['info']['post']
 
-            if context is None:
-                continue
 
             if context not in self.index2summary:
                 self.index2summary[len(self.index2summary)] = context
diff --git a/model/reward/instructor/summary_quality_trainer.py b/model/reward/instructor/summary_quality_trainer.py
new file mode 100644
index 00000000..a6604819
--- /dev/null
+++ b/model/reward/instructor/summary_quality_trainer.py
@@ -0,0 +1,132 @@
+import os
+os.environ['WANDB_PROJECT'] = 'quality-scoring'
+import torch
+import yaml
+import evaluate
+from typing import Any, Callable, List, Optional, Tuple, Union, Dict
+from torch import nn
+from argparse import ArgumentParser
+import numpy as np
+from torch.utils.data import Dataset
+from transformers import AutoModelForSequenceClassification
+from transformers import Trainer, PreTrainedModel, TrainingArguments, DataCollator, EvalPrediction, TrainerCallback, PreTrainedTokenizerBase
+from experimental_dataset import HFSummaryQuality, DataCollatorForSummaryScore
+from utils import get_tokenizer, train_val_dataset, freeze_top_n_layers, argument_parsing
+
+parser = ArgumentParser()
+parser.add_argument('config', type=str)
+
+accuracy = evaluate.load("mse")
+def compute_metrics(eval_pred):
+    predictions, labels = eval_pred
+    return accuracy.compute(predictions=predictions.flatten(), references=labels.flatten())
+
+
+class QualityTrainer(Trainer):
+    def __init__(self, model: Union[PreTrainedModel, nn.Module] = None,
+                 args: TrainingArguments = None,
+                 data_collator: Optional[DataCollator] = None,
+                 train_dataset: Optional[Dataset] = None,
+                 eval_dataset: Optional[Dataset] = None,
+                 tokenizer: Optional[PreTrainedTokenizerBase] = None,
+                 model_init: Callable[[], PreTrainedModel] = None,
+                 compute_metrics: Optional[Callable[[EvalPrediction], Dict]] = None,
+                 callbacks: Optional[List[TrainerCallback]] = None,
+                 optimizers: Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None),
+                 preprocess_logits_for_metrics: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] = None):
+        super().__init__(model, args, data_collator, train_dataset, eval_dataset, tokenizer,
+                         model_init, compute_metrics, callbacks, optimizers, preprocess_logits_for_metrics)
+        self.loss_fct = nn.L1Loss()
+        self.sigmoid = nn.Sigmoid()
+
+    def compute_loss(self, model, inputs, return_outputs=False):
+        labels = inputs.pop('labels')
+        # forward pass
+        outputs = model(**inputs)
+        logits = self.sigmoid(outputs.get("logits"))
+        loss = self.loss_fct(logits, labels)
+
+        return (loss, outputs) if return_outputs else loss
+
+    def _compute_loss(self, model, inputs):
+        inputs = self._prepare_inputs(inputs)
+        labels = inputs.pop('labels')
+        outputs = model(**inputs)
+        logits = self.sigmoid(outputs.get("logits"))
+        loss = self.loss_fct(logits, labels)
+
+        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 = inputs['labels']
+        if self.args.prediction_loss_only:
+            return (loss, None, None)
+
+        return (loss, logits, labels)
+
+if __name__ == "__main__":
+    training_conf = argument_parsing(parser)
+
+    model_name = training_conf['model_name']
+    tokenizer = get_tokenizer(model_name)
+    collate_fn = DataCollatorForSummaryScore(tokenizer,
+        max_length=training_conf['max_length'],
+        drop_token_type= 'galactica' in model_name
+    )
+    train = HFSummaryQuality(split='validation',
+        tokenizer=tokenizer,
+        max_length=training_conf['max_length']
+        )
+    eval = HFSummaryQuality(split='test',
+            tokenizer=tokenizer,
+            max_length=training_conf['max_length']
+        )
+    model = AutoModelForSequenceClassification.from_pretrained(model_name,
+        num_labels=len(train.label2idx), problem_type='regression')
+
+    if 'freeze_layer' in training_conf:
+        num_layer = training_conf['freeze_layer']
+        model = freeze_top_n_layers(model, num_layer)
+        model_parameters = filter(lambda p: p.requires_grad, model.parameters())
+        params = sum([np.prod(p.size()) for p in model_parameters])
+        print('Number of trainable : {}M'.format(int(params/1e6)))
+
+    args = TrainingArguments(
+        output_dir=f"{model_name}-finetuned",
+        num_train_epochs=training_conf['num_train_epochs'],
+        warmup_steps=500,
+        learning_rate=training_conf['learning_rate'],
+        # half_precision_backend="apex",
+        fp16=True,
+        gradient_checkpointing=training_conf['gradient_checkpointing'],
+        gradient_accumulation_steps=training_conf['gradient_accumulation_steps'],
+        per_device_train_batch_size=training_conf['per_device_train_batch_size'],
+        per_device_eval_batch_size=training_conf['per_device_eval_batch_size'],
+        weight_decay=0.01,
+        max_grad_norm=2.0,
+        logging_steps=10,
+        save_total_limit=4,
+        evaluation_strategy='steps',
+        eval_steps=training_conf['eval_steps'],
+        save_steps=1000,
+        report_to='wandb'
+    )
+    trainer = QualityTrainer(
+        model,
+        args,
+        train_dataset=train,
+        eval_dataset=eval,
+        data_collator=collate_fn,
+        tokenizer=tokenizer,
+        compute_metrics=compute_metrics
+    )
+    trainer.train()
diff --git a/model/reward/instructor/tests/test_dataset.py b/model/reward/instructor/tests/test_dataset.py
index 5765cd43..271db83c 100644
--- a/model/reward/instructor/tests/test_dataset.py
+++ b/model/reward/instructor/tests/test_dataset.py
@@ -1,7 +1,7 @@
 from transformers import AutoTokenizer
 from torch.utils.data import DataLoader
 from rank_datasets import WebGPT, HFSummary, DataCollatorForPairRank
-
+from experimental_dataset import HFSummaryQuality, DataCollatorForSummaryScore
 
 def test_hfsummary():
     
@@ -24,6 +24,17 @@ def test_webgpt():
         print(batch['input_ids'].shape)
 
 
+def test_hf_quality():
+
+    tokenizer = AutoTokenizer.from_pretrained("bigscience/mt0-large")
+    collate_fn = DataCollatorForSummaryScore(tokenizer, max_length=200)
+    dataset = HFSummaryQuality('validation', tokenizer)
+    dataloader = DataLoader(dataset, collate_fn=collate_fn, batch_size=32)
+    for batch in dataloader:
+        print(batch['input_ids'].shape)
+
+
+
 if __name__ == "__main__":
-    test_hfsummary()
+    test_hf_quality()
     # test_webgpt()
\ No newline at end of file
diff --git a/model/reward/instructor/trainer.py b/model/reward/instructor/trainer.py
index 391464c6..c8063cf7 100644
--- a/model/reward/instructor/trainer.py
+++ b/model/reward/instructor/trainer.py
@@ -117,13 +117,13 @@ if __name__ == "__main__":
         gradient_checkpointing=training_conf['gradient_checkpointing'],
         gradient_accumulation_steps=training_conf['gradient_accumulation_steps'],
         per_device_train_batch_size=training_conf['per_device_train_batch_size'],
-        per_device_eval_batch_size=5,
+        per_device_eval_batch_size=training_conf['per_device_eval_batch_size'],
         weight_decay=0.01,
         max_grad_norm=2.0,
         logging_steps=10,
         save_total_limit=4,
         evaluation_strategy='steps',
-        eval_steps=500,
+        eval_steps=training_conf['eval_steps'],
         save_steps=1000,
         report_to='wandb'
     )
diff --git a/model/reward/instructor/utils.py b/model/reward/instructor/utils.py
index f26add55..d59bb13c 100644
--- a/model/reward/instructor/utils.py
+++ b/model/reward/instructor/utils.py
@@ -69,6 +69,7 @@ def argument_parsing(parser):
         'eval_steps': 500,
         'loss': 'rank',
         'max_length': 440,
+        'per_device_eval_batch_size': 5,
         'per_device_train_batch_size': 8,
         'gradient_accumulation_steps': 8,
         'gradient_checkpointing': False,

From 5a18fba9e6e80f69e0efe88db6dd268463c4c5d9 Mon Sep 17 00:00:00 2001
From: Graeme Harris <harris.graemeza@gmail.com>
Date: Sat, 31 Dec 2022 01:48:37 +0200
Subject: [PATCH 20/41] Added basic redis + redis insights config for docker
 and ansible

---
 ansible/dev.yaml    | 18 ++++++++++++++++++
 docker-compose.yaml | 27 +++++++++++++++++++++++++++
 redis.conf          |  2 ++
 3 files changed, 47 insertions(+)
 create mode 100644 redis.conf

diff --git a/ansible/dev.yaml b/ansible/dev.yaml
index d022ba3c..53fdc611 100644
--- a/ansible/dev.yaml
+++ b/ansible/dev.yaml
@@ -32,6 +32,24 @@
         - name: oasst-postgres
         - name: oasst-postgres-web
 
+    - name: Setup redis
+      community.docker.docker_container:
+        name: "{{ item.name }}"
+        image: redis
+        state: started
+        restart_policy: always
+        network_mode: oasst
+        volumes:
+          - redis:/data
+        healthcheck:
+          test: ["CMD-SHELL", "redis-cli ping | grep PONG"]
+          interval: 2s
+          timeout: 2s
+          retries: 10
+        ports:
+          - 6379:6379
+        command: redis-server /usr/local/etc/redis/redis.conf
+
     - name: Set up maildev
       community.docker.docker_container:
         name: oasst-maildev
diff --git a/docker-compose.yaml b/docker-compose.yaml
index d329c780..a5b90498 100644
--- a/docker-compose.yaml
+++ b/docker-compose.yaml
@@ -27,6 +27,29 @@ services:
       timeout: 2s
       retries: 10
 
+  # Redis - caching + rate limiting on BE
+  redis:
+    image: redis
+    restart: always
+    ports:
+      - 6379:6379
+    healthcheck:
+      test: ["CMD-SHELL", "redis-cli ping | grep PONG"]
+      interval: 2s
+      timeout: 2s
+      retries : 10
+    command: redis-server /usr/local/etc/redis/redis.conf
+    volumes:
+      - ./redis.conf:/usr/local/etc/redis/redis.conf
+      - redis:/data
+  # insights host - redis:6379
+  redis-insights:
+    image: redislabs/redisinsight:latest
+    ports:
+      - 8001:8001
+    volumes:
+      - redis:/data
+
   # This DB is for Web Authentication and data caching.
   webdb:
     image: postgres
@@ -100,3 +123,7 @@ services:
     ports:
       - "3000:3000"
     command: bash wait-for-postgres.sh node server.js
+
+volumes:
+  redis:
+    driver: local
diff --git a/redis.conf b/redis.conf
new file mode 100644
index 00000000..c06efe47
--- /dev/null
+++ b/redis.conf
@@ -0,0 +1,2 @@
+maxmemory 100mb
+maxmemory-policy allkeys-lru
\ No newline at end of file

From 743de0bb18da1f06061630fbcf4eb9cdaa781b90 Mon Sep 17 00:00:00 2001
From: Graeme Harris <harris.graemeza@gmail.com>
Date: Sat, 31 Dec 2022 01:49:34 +0200
Subject: [PATCH 21/41] Linting + prettier

---
 docker-compose.yaml | 2 +-
 redis.conf          | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/docker-compose.yaml b/docker-compose.yaml
index a5b90498..27f5d741 100644
--- a/docker-compose.yaml
+++ b/docker-compose.yaml
@@ -37,7 +37,7 @@ services:
       test: ["CMD-SHELL", "redis-cli ping | grep PONG"]
       interval: 2s
       timeout: 2s
-      retries : 10
+      retries: 10
     command: redis-server /usr/local/etc/redis/redis.conf
     volumes:
       - ./redis.conf:/usr/local/etc/redis/redis.conf
diff --git a/redis.conf b/redis.conf
index c06efe47..58da1e05 100644
--- a/redis.conf
+++ b/redis.conf
@@ -1,2 +1,2 @@
 maxmemory 100mb
-maxmemory-policy allkeys-lru
\ No newline at end of file
+maxmemory-policy allkeys-lru

From 3bb0b04f92c8040ac26b5aebf061d9b911245af6 Mon Sep 17 00:00:00 2001
From: Graeme Harris <harris.graemeza@gmail.com>
Date: Sat, 31 Dec 2022 12:56:09 +0200
Subject: [PATCH 22/41] Removed local volume from redis image

---
 docker-compose.yaml | 7 -------
 1 file changed, 7 deletions(-)

diff --git a/docker-compose.yaml b/docker-compose.yaml
index 27f5d741..ed72c820 100644
--- a/docker-compose.yaml
+++ b/docker-compose.yaml
@@ -41,14 +41,11 @@ services:
     command: redis-server /usr/local/etc/redis/redis.conf
     volumes:
       - ./redis.conf:/usr/local/etc/redis/redis.conf
-      - redis:/data
   # insights host - redis:6379
   redis-insights:
     image: redislabs/redisinsight:latest
     ports:
       - 8001:8001
-    volumes:
-      - redis:/data
 
   # This DB is for Web Authentication and data caching.
   webdb:
@@ -123,7 +120,3 @@ services:
     ports:
       - "3000:3000"
     command: bash wait-for-postgres.sh node server.js
-
-volumes:
-  redis:
-    driver: local

From 1ddd9155f91ae8801b4c4567ffb7bc1131a8dc0a Mon Sep 17 00:00:00 2001
From: theblackcat102 <theblackcat102@github.com>
Date: Sun, 1 Jan 2023 11:35:49 +0000
Subject: [PATCH 23/41] [fix] remove vscode settings

---
 .vscode/settings.json | 4 ----
 1 file changed, 4 deletions(-)
 delete mode 100644 .vscode/settings.json

diff --git a/.vscode/settings.json b/.vscode/settings.json
deleted file mode 100644
index 4c58a32f..00000000
--- a/.vscode/settings.json
+++ /dev/null
@@ -1,4 +0,0 @@
-{
-  "python.formatting.provider": "autopep8",
-  "python.analysis.extraPaths": ["${workspaceFolder}/oasst-shared"]
-}

From 10724411eb139204a7bf435df1312fd5fc16fde3 Mon Sep 17 00:00:00 2001
From: Graeme Harris <harris.graemeza@gmail.com>
Date: Sun, 1 Jan 2023 13:37:01 +0200
Subject: [PATCH 24/41] Removed redis from ansible for now

---
 ansible/dev.yaml | 18 ------------------
 1 file changed, 18 deletions(-)

diff --git a/ansible/dev.yaml b/ansible/dev.yaml
index 53fdc611..d022ba3c 100644
--- a/ansible/dev.yaml
+++ b/ansible/dev.yaml
@@ -32,24 +32,6 @@
         - name: oasst-postgres
         - name: oasst-postgres-web
 
-    - name: Setup redis
-      community.docker.docker_container:
-        name: "{{ item.name }}"
-        image: redis
-        state: started
-        restart_policy: always
-        network_mode: oasst
-        volumes:
-          - redis:/data
-        healthcheck:
-          test: ["CMD-SHELL", "redis-cli ping | grep PONG"]
-          interval: 2s
-          timeout: 2s
-          retries: 10
-        ports:
-          - 6379:6379
-        command: redis-server /usr/local/etc/redis/redis.conf
-
     - name: Set up maildev
       community.docker.docker_container:
         name: oasst-maildev

From fe99b46f2e02f48a55165684a61fafa4dc5c823a Mon Sep 17 00:00:00 2001
From: theblackcat102 <theblackcat102@github.com>
Date: Sun, 1 Jan 2023 11:43:15 +0000
Subject: [PATCH 25/41] [fix] pre-commit update

---
 model/reward/instructor/README.md             |  11 +-
 model/reward/instructor/TODO.md               |  18 +--
 model/reward/instructor/cls_dataset.py        |  37 +++--
 .../configs/bloomz-560m-summary.yml           |   2 +-
 .../reward/instructor/configs/bloomz-560m.yml |   2 +-
 .../configs/electra-base-dis-webgpt.yml       |   2 +-
 .../instructor/configs/galactica-125m.yml     |   2 +-
 .../instructor/configs/galactica-1b.yml       |   2 +-
 .../test-galactica-125m-classification.yml    |   2 +-
 .../reward/instructor/experimental_dataset.py |  50 +++---
 model/reward/instructor/rank_datasets.py      | 104 ++++++------
 .../instructor/summary_quality_trainer.py     | 140 ++++++++++-------
 model/reward/instructor/tests/test_dataset.py |  27 ++--
 model/reward/instructor/trainer.py            | 148 +++++++++++-------
 model/reward/instructor/utils.py              |  84 +++++-----
 15 files changed, 337 insertions(+), 294 deletions(-)

diff --git a/model/reward/instructor/README.md b/model/reward/instructor/README.md
index 31c25371..73a872a0 100644
--- a/model/reward/instructor/README.md
+++ b/model/reward/instructor/README.md
@@ -2,7 +2,6 @@
 
 Trainer code based on huggingface. Compatible with deepspeed or accelerate
 
-
 Requirements
 
 ```
@@ -15,12 +14,10 @@ torch==1.12
 
 Start training reward model
 
-
 ```bash
 python trainer.py configs/electra-base-dis-webgpt.yml
 ```
 
-
 Additional axis labeling, this outputs a 4 summary quality evaluation metrics (score are normalized to 0-1 )
 
 ```bash
@@ -29,13 +26,13 @@ python summary_quality_trainer.py configs/test-bloomz-560m-quality.yml
 
 The four summary are :
 
-* overall
+- overall
 
-* accuracy
+- accuracy
 
-* coverage
+- coverage
 
-* coherence
+- coherence
 
 ## Dataset
 
diff --git a/model/reward/instructor/TODO.md b/model/reward/instructor/TODO.md
index 1e653922..ed33b3c0 100644
--- a/model/reward/instructor/TODO.md
+++ b/model/reward/instructor/TODO.md
@@ -1,23 +1,19 @@
-
 Some other reward features we can use
 
-0. Finish classifcation feature 
+0. Finish classifcation feature
 
 1. Summaries from human feedback
 
-* use `confidence` score into the RM learning, ensure the output rank score correlates with confidence
+- use `confidence` score into the RM learning, ensure the output rank score correlates with confidence
 
-* each labeling has a labeling `note`, basically comments by labeler, not sure what else we can use
+- each labeling has a labeling `note`, basically comments by labeler, not sure what else we can use
 
-* ~~Use the score for "overall", "accuracy", "coverage", "coherence" from axis/evals to train an addition model (rank additional aspect of the policy model)~~
-
-    * this should be placed under experimental_dataset.py
+- ~~Use the score for "overall", "accuracy", "coverage", "coherence" from axis/evals to train an addition model (rank additional aspect of the policy model)~~
 
+  - this should be placed under experimental_dataset.py
 
 2. Add support for anthropic dataset
 
-* anthropic dataset is more like a conversation tree which is much complex than simply question-answer schema
-
-    * this is basically a MCTS from alphazero.
-
+- anthropic dataset is more like a conversation tree which is much complex than simply question-answer schema
 
+  - this is basically a MCTS from alphazero.
diff --git a/model/reward/instructor/cls_dataset.py b/model/reward/instructor/cls_dataset.py
index ff824d19..09aa821b 100644
--- a/model/reward/instructor/cls_dataset.py
+++ b/model/reward/instructor/cls_dataset.py
@@ -1,32 +1,34 @@
-'''
+# -*- coding: utf-8 -*-
+"""
 
     classification based ranking
 
-'''
-import os
+"""
 import json
+import os
 import random
-import torch
-import numpy as np
+
 from dataset import load_dataset
 from torch.utils.data import Dataset
+
 from .utils import webgpt_return_format
 
+
 class WebGPTDataset(Dataset):
-    def __init__(self, mode='train', index_cache='dataset/webgpt_train_idx.pt', additional_dataset=None) -> None:
+    def __init__(self, mode="train", index_cache="dataset/webgpt_train_idx.pt", additional_dataset=None) -> None:
         super().__init__()
-        '''
+        """
             mode : train or val, used for validation purpose, has nothing to do with original split
             additional_dataset : a list of jsonline format with idx, question and texts (generate candidates)
                 idx : must match the index you iterate from comparison enumerate order
                 question : for validation purpose
                 texts : list of K generate results from the question prompt
-        '''
-        os.makedirs('dataset', exist_ok=True)
+        """
+        os.makedirs("dataset", exist_ok=True)
         dataset = load_dataset("openai/webgpt_comparisons")
         self.dataset = []
         self.dataset_index = []
-        for idx, row in enumerate(dataset['train']):
+        for idx, row in enumerate(dataset["train"]):
             self.dataset.append(webgpt_return_format(row))
 
         # since this dataset was generated from 176B GPT-3
@@ -36,17 +38,17 @@ class WebGPTDataset(Dataset):
         if additional_dataset is not None:
             self.sample_additional = True
             self.additional = {}
-            with open(additional_dataset, 'r') as f:
+            with open(additional_dataset, "r") as f:
                 for line in f:
                     row = json.loads(line)
-                    if row['idx'] in self.dataset_index:
-                        self.additional[row['idx']] = row['negatives']
+                    if row["idx"] in self.dataset_index:
+                        self.additional[row["idx"]] = row["negatives"]
             if len(self.additional) != len(self.dataset_index):
                 for match_idx in self.dataset_index:
                     if match_idx in self.additional:
                         continue
 
-                    idx = match_idx-900
+                    idx = match_idx - 900
                     while idx not in self.additional:
                         idx -= 1
                     self.additional[match_idx] = self.additional[idx]
@@ -57,10 +59,7 @@ class WebGPTDataset(Dataset):
     def __getitem__(self, index):
         row = self.dataset[index]
         if not self.sample_additional:
-            return row['question'], row['pos'], row['neg']
+            return row["question"], row["pos"], row["neg"]
 
         gen_neg = random.choice(self.additional[self.dataset_index[index]])
-        return row['question'], row['pos'], row['neg'], gen_neg
-
-
-
+        return row["question"], row["pos"], row["neg"], gen_neg
diff --git a/model/reward/instructor/configs/bloomz-560m-summary.yml b/model/reward/instructor/configs/bloomz-560m-summary.yml
index a02f4e4a..55ed6cd1 100644
--- a/model/reward/instructor/configs/bloomz-560m-summary.yml
+++ b/model/reward/instructor/configs/bloomz-560m-summary.yml
@@ -6,4 +6,4 @@ max_length: 600
 freeze_layer: 12
 num_train_epochs: 2
 datasets:
-  - hfsummary
\ No newline at end of file
+  - hfsummary
diff --git a/model/reward/instructor/configs/bloomz-560m.yml b/model/reward/instructor/configs/bloomz-560m.yml
index c8f55746..bf3f14dd 100644
--- a/model/reward/instructor/configs/bloomz-560m.yml
+++ b/model/reward/instructor/configs/bloomz-560m.yml
@@ -7,4 +7,4 @@ freeze_layer: 12
 num_train_epochs: 2
 datasets:
   - webgpt
-  - hfsummary
\ No newline at end of file
+  - hfsummary
diff --git a/model/reward/instructor/configs/electra-base-dis-webgpt.yml b/model/reward/instructor/configs/electra-base-dis-webgpt.yml
index fc168b63..89200fe1 100644
--- a/model/reward/instructor/configs/electra-base-dis-webgpt.yml
+++ b/model/reward/instructor/configs/electra-base-dis-webgpt.yml
@@ -1,3 +1,3 @@
 model_name: google/electra-large-discriminator
 learning_rate: 3e-5
-max_length: 300
\ No newline at end of file
+max_length: 300
diff --git a/model/reward/instructor/configs/galactica-125m.yml b/model/reward/instructor/configs/galactica-125m.yml
index 55e093f5..13dbdfbe 100644
--- a/model/reward/instructor/configs/galactica-125m.yml
+++ b/model/reward/instructor/configs/galactica-125m.yml
@@ -10,4 +10,4 @@ max_length: 512
 num_train_epochs: 2
 datasets:
   - webgpt
-  - hfsummary
\ No newline at end of file
+  - hfsummary
diff --git a/model/reward/instructor/configs/galactica-1b.yml b/model/reward/instructor/configs/galactica-1b.yml
index 5a094520..8ffd74e9 100644
--- a/model/reward/instructor/configs/galactica-1b.yml
+++ b/model/reward/instructor/configs/galactica-1b.yml
@@ -11,4 +11,4 @@ max_length: 400
 num_train_epochs: 2
 datasets:
   - webgpt
-  - hfsummary
\ No newline at end of file
+  - hfsummary
diff --git a/model/reward/instructor/configs/test-galactica-125m-classification.yml b/model/reward/instructor/configs/test-galactica-125m-classification.yml
index 1ad1f47c..e36efcf3 100644
--- a/model/reward/instructor/configs/test-galactica-125m-classification.yml
+++ b/model/reward/instructor/configs/test-galactica-125m-classification.yml
@@ -11,4 +11,4 @@ max_length: 128
 num_train_epochs: 2
 datasets:
   - webgpt
-  - hfsummary
\ No newline at end of file
+  - hfsummary
diff --git a/model/reward/instructor/experimental_dataset.py b/model/reward/instructor/experimental_dataset.py
index 47d20d64..28f62967 100644
--- a/model/reward/instructor/experimental_dataset.py
+++ b/model/reward/instructor/experimental_dataset.py
@@ -1,4 +1,5 @@
-'''
+# -*- coding: utf-8 -*-
+"""
     HFSummary
 
         I want to train a multi regression model on axis_evals dataset mainly we can estimate the score of these score
@@ -7,15 +8,16 @@
 
         Should be better than just a preference score
 
-'''
-import torch
-from typing import Optional, Union
-import numpy as np
+"""
 from collections import defaultdict
-from datasets import load_dataset
 from dataclasses import dataclass
+from typing import Optional, Union
+
+import numpy as np
+import torch
+from datasets import load_dataset
 from torch.utils.data import Dataset
-from transformers.tokenization_utils_base import PreTrainedTokenizerBase, PaddingStrategy
+from transformers.tokenization_utils_base import PaddingStrategy, PreTrainedTokenizerBase
 
 
 @dataclass
@@ -25,12 +27,13 @@ class DataCollatorForSummaryScore:
     Data collator that will dynamically pad the inputs for multiple choice received.
 
     """
+
     tokenizer: PreTrainedTokenizerBase
     num_choices: int = 2
     padding: Union[bool, str, PaddingStrategy] = True
     max_length: Optional[int] = None
     pad_to_multiple_of: Optional[int] = None
-    drop_token_type: bool = False # galactica
+    drop_token_type: bool = False  # galactica
 
     def __call__(self, batch):
 
@@ -48,17 +51,17 @@ class DataCollatorForSummaryScore:
             return_tensors="pt",
         )
         if self.drop_token_type:
-            batch_feature.pop('token_type_ids')
+            batch_feature.pop("token_type_ids")
         # batch = {k: v.view(batch_size, self.num_choices, -1) for k, v in batch.items()}
-        batch_feature['labels'] = torch.from_numpy(np.array(labels)).float()
+        batch_feature["labels"] = torch.from_numpy(np.array(labels)).float()
         return batch_feature
 
 
 class HFSummaryQuality(Dataset):
     def __init__(self, split, tokenizer, max_length=300) -> None:
         super().__init__()
-        assert split in ('validation', 'test')
-        dataset = load_dataset('Tristan/summarize_from_feedback', 'axis')[split]
+        assert split in ("validation", "test")
+        dataset = load_dataset("Tristan/summarize_from_feedback", "axis")[split]
         self.max_length = max_length
         mean_scores = defaultdict(list)
         self.contexts = []
@@ -66,22 +69,21 @@ class HFSummaryQuality(Dataset):
         self.labels = []
         for data in dataset:
 
-            if 'article' in data['info'] and \
-                data['info']['article'] is not None:
-                context = data['info']['article']
-            elif 'post' in data['info']:
-                context = data['info']['post']
+            if "article" in data["info"] and data["info"]["article"] is not None:
+                context = data["info"]["article"]
+            elif "post" in data["info"]:
+                context = data["info"]["post"]
             self.contexts.append(context)
 
-            response = data['summary']['text']
+            response = data["summary"]["text"]
             self.responses.append(response)
-            self.labels.append(data['summary']['axes'])
-            for axis, score in data['summary']['axes'].items():
+            self.labels.append(data["summary"]["axes"])
+            for axis, score in data["summary"]["axes"].items():
                 if score is not None:
                     mean_scores[axis].append(score)
 
-        self.label2idx = { key: idx for idx, key in enumerate(mean_scores.keys()) }
-        self.label2mean = { key: np.mean(scores) for key, scores in mean_scores.items() }
+        self.label2idx = {key: idx for idx, key in enumerate(mean_scores.keys())}
+        self.label2mean = {key: np.mean(scores) for key, scores in mean_scores.items()}
         self.tokenizer = tokenizer
         print(self.label2idx)
 
@@ -94,7 +96,5 @@ class HFSummaryQuality(Dataset):
         response = self.responses[index]
         labels = np.zeros(len(self.label2idx))
         for key, score in self.labels[index].items():
-            labels[self.label2idx[key]] = (self.label2mean[key] if score is None else score)/10
+            labels[self.label2idx[key]] = (self.label2mean[key] if score is None else score) / 10
         return self.tokenizer(context, response, truncation=True, max_length=self.max_length), labels
-
-
diff --git a/model/reward/instructor/rank_datasets.py b/model/reward/instructor/rank_datasets.py
index f38885e4..99ba9955 100644
--- a/model/reward/instructor/rank_datasets.py
+++ b/model/reward/instructor/rank_datasets.py
@@ -1,4 +1,5 @@
-'''
+# -*- coding: utf-8 -*-
+"""
     author: theblackcat102
 
     Dataset output format from __getitem__
@@ -17,13 +18,15 @@
         inferior than the human perference one
 
 
-'''
-from typing import Optional, Union
+"""
 from dataclasses import dataclass
+from typing import Optional, Union
+
 import numpy as np
-from torch.utils.data import Dataset
 from datasets import load_dataset
-from transformers.tokenization_utils_base import PreTrainedTokenizerBase, PaddingStrategy
+from torch.utils.data import Dataset
+from transformers.tokenization_utils_base import PaddingStrategy, PreTrainedTokenizerBase
+
 
 @dataclass
 class DataCollatorForPairRank:
@@ -32,12 +35,13 @@ class DataCollatorForPairRank:
     Data collator that will dynamically pad the inputs for multiple choice received.
 
     """
+
     tokenizer: PreTrainedTokenizerBase
     num_choices: int = 2
     padding: Union[bool, str, PaddingStrategy] = True
     max_length: Optional[int] = None
     pad_to_multiple_of: Optional[int] = None
-    drop_token_type: bool = False # galactica
+    drop_token_type: bool = False  # galactica
 
     def __call__(self, features):
 
@@ -45,12 +49,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, max_length=self.max_length))
-                flatten_features.append(self.tokenizer(question, neg,
-                    truncation=True, max_length=self.max_length))
+                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(
             flatten_features,
             padding=self.padding,
@@ -59,13 +61,12 @@ class DataCollatorForPairRank:
             return_tensors="pt",
         )
         if self.drop_token_type:
-            batch.pop('token_type_ids')
+            batch.pop("token_type_ids")
         # batch = {k: v.view(batch_size, self.num_choices, -1) for k, v in batch.items()}
         return batch
 
 
 class WebGPT(Dataset):
-
     def __init__(self) -> None:
         super().__init__()
 
@@ -74,23 +75,19 @@ class WebGPT(Dataset):
         # using prompt as our index will allows us
         # to add additional generated prompt later
         self.index2question = {}
-        for row in dataset['train']:
-            question = row['question']['full_text']
+        for row in dataset["train"]:
+            question = row["question"]["full_text"]
             if question not in self.index2question:
                 self.index2question[len(self.index2question)] = question
 
             if question not in questions:
                 questions[question] = []
 
-            if row['score_0'] > row['score_1']:
+            if row["score_0"] > row["score_1"]:
                 # not going to risk it
-                questions[question].append((
-                    row['answer_0'], row['answer_1']
-                ))
+                questions[question].append((row["answer_0"], row["answer_1"]))
             else:
-                questions[question].append((
-                    row['answer_1'], row['answer_0']
-                ))
+                questions[question].append((row["answer_1"], row["answer_0"]))
 
         self.questions = questions
 
@@ -104,61 +101,55 @@ class WebGPT(Dataset):
         return question, rows
 
 
-
-
 class HFSummary(Dataset):
-    '''
-        Human feedback data from OpenAI
-        https://github.com/openai/summarize-from-feedback
-        
-        labeling method : pair comparison, 0 or 1
+    """
+    Human feedback data from OpenAI
+    https://github.com/openai/summarize-from-feedback
 
-    '''
-    def __init__(self, split='train',
-        conf_threshold=-1,
-        max_comparison_per_sample=3) -> None:
+    labeling method : pair comparison, 0 or 1
+
+    """
+
+    def __init__(self, split="train", conf_threshold=-1, max_comparison_per_sample=3) -> None:
         super().__init__()
-        assert split in ('train', 'valid1', 'valid2', 'test')
+        assert split in ("train", "valid1", "valid2", "test")
         summaries = {}
         # using prompt as our index will allows us
         # to add additional generated prompt later
         self.index2summary = {}
         self.max_comparison_per_sample = max_comparison_per_sample
-        major_split = split if 'train' == split else 'validation'
-        dataset = load_dataset('Tristan/summarize_from_feedback', 'comparisons')[major_split]
+        major_split = split if "train" == split else "validation"
+        dataset = load_dataset("Tristan/summarize_from_feedback", "comparisons")[major_split]
         for data in dataset:
-            if 'extra' in data and \
-                'confidence' in data['extra'] and \
-                data['extra']['confidence'] is not None and \
-                conf_threshold > data['extra']['confidence']:
-                print('skipping {}'.format(data['info']['id']))
+            if (
+                "extra" in data
+                and "confidence" in data["extra"]
+                and data["extra"]["confidence"] is not None
+                and conf_threshold > data["extra"]["confidence"]
+            ):
+                print("skipping {}".format(data["info"]["id"]))
                 continue
 
-            if split != 'train' and split != data['split']:
+            if split != "train" and split != data["split"]:
                 continue
 
-            if 'article' in data['info'] and \
-                data['info']['article'] is not None:
-                context = data['info']['article']
-            elif 'post' in data['info']:
-                context = data['info']['post']
-
+            if "article" in data["info"] and data["info"]["article"] is not None:
+                context = data["info"]["article"]
+            elif "post" in data["info"]:
+                context = data["info"]["post"]
 
             if context not in self.index2summary:
                 self.index2summary[len(self.index2summary)] = context
-            
+
             if context not in summaries:
                 summaries[context] = []
 
-            pos, neg = (0, 1) if data['choice'] == 0 else (1, 0)
-            summaries[context].append((
-                data['summaries'][pos]['text'],
-                data['summaries'][neg]['text']
-            ))
+            pos, neg = (0, 1) if data["choice"] == 0 else (1, 0)
+            summaries[context].append((data["summaries"][pos]["text"], data["summaries"][neg]["text"]))
 
         self.summaries = summaries
 
-        self.postfix_prompt = ' TLDR;'
+        self.postfix_prompt = " TLDR;"
 
     def __len__(self):
         return len(self.index2summary)
@@ -172,5 +163,4 @@ 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+self.postfix_prompt, [ 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/summary_quality_trainer.py b/model/reward/instructor/summary_quality_trainer.py
index a6604819..88bf1abf 100644
--- a/model/reward/instructor/summary_quality_trainer.py
+++ b/model/reward/instructor/summary_quality_trainer.py
@@ -1,46 +1,72 @@
+# -*- coding: utf-8 -*-
 import os
-os.environ['WANDB_PROJECT'] = 'quality-scoring'
-import torch
-import yaml
-import evaluate
-from typing import Any, Callable, List, Optional, Tuple, Union, Dict
-from torch import nn
 from argparse import ArgumentParser
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+
+import evaluate
 import numpy as np
+import torch
+from experimental_dataset import DataCollatorForSummaryScore, HFSummaryQuality
+from torch import nn
 from torch.utils.data import Dataset
-from transformers import AutoModelForSequenceClassification
-from transformers import Trainer, PreTrainedModel, TrainingArguments, DataCollator, EvalPrediction, TrainerCallback, PreTrainedTokenizerBase
-from experimental_dataset import HFSummaryQuality, DataCollatorForSummaryScore
-from utils import get_tokenizer, train_val_dataset, freeze_top_n_layers, argument_parsing
+from transformers import (
+    AutoModelForSequenceClassification,
+    DataCollator,
+    EvalPrediction,
+    PreTrainedModel,
+    PreTrainedTokenizerBase,
+    Trainer,
+    TrainerCallback,
+    TrainingArguments,
+)
+from utils import argument_parsing, freeze_top_n_layers, get_tokenizer
+
+os.environ["WANDB_PROJECT"] = "quality-scoring"
 
 parser = ArgumentParser()
-parser.add_argument('config', type=str)
+parser.add_argument("config", type=str)
 
 accuracy = evaluate.load("mse")
+
+
 def compute_metrics(eval_pred):
     predictions, labels = eval_pred
     return accuracy.compute(predictions=predictions.flatten(), references=labels.flatten())
 
 
 class QualityTrainer(Trainer):
-    def __init__(self, model: Union[PreTrainedModel, nn.Module] = None,
-                 args: TrainingArguments = None,
-                 data_collator: Optional[DataCollator] = None,
-                 train_dataset: Optional[Dataset] = None,
-                 eval_dataset: Optional[Dataset] = None,
-                 tokenizer: Optional[PreTrainedTokenizerBase] = None,
-                 model_init: Callable[[], PreTrainedModel] = None,
-                 compute_metrics: Optional[Callable[[EvalPrediction], Dict]] = None,
-                 callbacks: Optional[List[TrainerCallback]] = None,
-                 optimizers: Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None),
-                 preprocess_logits_for_metrics: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] = None):
-        super().__init__(model, args, data_collator, train_dataset, eval_dataset, tokenizer,
-                         model_init, compute_metrics, callbacks, optimizers, preprocess_logits_for_metrics)
+    def __init__(
+        self,
+        model: Union[PreTrainedModel, nn.Module] = None,
+        args: TrainingArguments = None,
+        data_collator: Optional[DataCollator] = None,
+        train_dataset: Optional[Dataset] = None,
+        eval_dataset: Optional[Dataset] = None,
+        tokenizer: Optional[PreTrainedTokenizerBase] = None,
+        model_init: Callable[[], PreTrainedModel] = None,
+        compute_metrics: Optional[Callable[[EvalPrediction], Dict]] = None,
+        callbacks: Optional[List[TrainerCallback]] = None,
+        optimizers: Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None),
+        preprocess_logits_for_metrics: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] = None,
+    ):
+        super().__init__(
+            model,
+            args,
+            data_collator,
+            train_dataset,
+            eval_dataset,
+            tokenizer,
+            model_init,
+            compute_metrics,
+            callbacks,
+            optimizers,
+            preprocess_logits_for_metrics,
+        )
         self.loss_fct = nn.L1Loss()
         self.sigmoid = nn.Sigmoid()
 
     def compute_loss(self, model, inputs, return_outputs=False):
-        labels = inputs.pop('labels')
+        labels = inputs.pop("labels")
         # forward pass
         outputs = model(**inputs)
         logits = self.sigmoid(outputs.get("logits"))
@@ -50,75 +76,73 @@ class QualityTrainer(Trainer):
 
     def _compute_loss(self, model, inputs):
         inputs = self._prepare_inputs(inputs)
-        labels = inputs.pop('labels')
+        labels = inputs.pop("labels")
         outputs = model(**inputs)
         logits = self.sigmoid(outputs.get("logits"))
         loss = self.loss_fct(logits, labels)
 
         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]]:
+    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 = inputs['labels']
+        labels = inputs["labels"]
         if self.args.prediction_loss_only:
             return (loss, None, None)
 
         return (loss, logits, labels)
 
+
 if __name__ == "__main__":
     training_conf = argument_parsing(parser)
 
-    model_name = training_conf['model_name']
+    model_name = training_conf["model_name"]
     tokenizer = get_tokenizer(model_name)
-    collate_fn = DataCollatorForSummaryScore(tokenizer,
-        max_length=training_conf['max_length'],
-        drop_token_type= 'galactica' in model_name
+    collate_fn = DataCollatorForSummaryScore(
+        tokenizer, max_length=training_conf["max_length"], drop_token_type="galactica" in model_name
+    )
+    train = HFSummaryQuality(split="validation", tokenizer=tokenizer, max_length=training_conf["max_length"])
+    eval = HFSummaryQuality(split="test", tokenizer=tokenizer, max_length=training_conf["max_length"])
+    model = AutoModelForSequenceClassification.from_pretrained(
+        model_name, num_labels=len(train.label2idx), problem_type="regression"
     )
-    train = HFSummaryQuality(split='validation',
-        tokenizer=tokenizer,
-        max_length=training_conf['max_length']
-        )
-    eval = HFSummaryQuality(split='test',
-            tokenizer=tokenizer,
-            max_length=training_conf['max_length']
-        )
-    model = AutoModelForSequenceClassification.from_pretrained(model_name,
-        num_labels=len(train.label2idx), problem_type='regression')
 
-    if 'freeze_layer' in training_conf:
-        num_layer = training_conf['freeze_layer']
+    if "freeze_layer" in training_conf:
+        num_layer = training_conf["freeze_layer"]
         model = freeze_top_n_layers(model, num_layer)
         model_parameters = filter(lambda p: p.requires_grad, model.parameters())
         params = sum([np.prod(p.size()) for p in model_parameters])
-        print('Number of trainable : {}M'.format(int(params/1e6)))
+        print("Number of trainable : {}M".format(int(params / 1e6)))
 
     args = TrainingArguments(
         output_dir=f"{model_name}-finetuned",
-        num_train_epochs=training_conf['num_train_epochs'],
+        num_train_epochs=training_conf["num_train_epochs"],
         warmup_steps=500,
-        learning_rate=training_conf['learning_rate'],
+        learning_rate=training_conf["learning_rate"],
         # half_precision_backend="apex",
         fp16=True,
-        gradient_checkpointing=training_conf['gradient_checkpointing'],
-        gradient_accumulation_steps=training_conf['gradient_accumulation_steps'],
-        per_device_train_batch_size=training_conf['per_device_train_batch_size'],
-        per_device_eval_batch_size=training_conf['per_device_eval_batch_size'],
+        gradient_checkpointing=training_conf["gradient_checkpointing"],
+        gradient_accumulation_steps=training_conf["gradient_accumulation_steps"],
+        per_device_train_batch_size=training_conf["per_device_train_batch_size"],
+        per_device_eval_batch_size=training_conf["per_device_eval_batch_size"],
         weight_decay=0.01,
         max_grad_norm=2.0,
         logging_steps=10,
         save_total_limit=4,
-        evaluation_strategy='steps',
-        eval_steps=training_conf['eval_steps'],
+        evaluation_strategy="steps",
+        eval_steps=training_conf["eval_steps"],
         save_steps=1000,
-        report_to='wandb'
+        report_to="wandb",
     )
     trainer = QualityTrainer(
         model,
@@ -127,6 +151,6 @@ if __name__ == "__main__":
         eval_dataset=eval,
         data_collator=collate_fn,
         tokenizer=tokenizer,
-        compute_metrics=compute_metrics
+        compute_metrics=compute_metrics,
     )
     trainer.train()
diff --git a/model/reward/instructor/tests/test_dataset.py b/model/reward/instructor/tests/test_dataset.py
index 271db83c..f367a50d 100644
--- a/model/reward/instructor/tests/test_dataset.py
+++ b/model/reward/instructor/tests/test_dataset.py
@@ -1,40 +1,41 @@
-from transformers import AutoTokenizer
+# -*- coding: utf-8 -*-
+from experimental_dataset import DataCollatorForSummaryScore, HFSummaryQuality
+from rank_datasets import DataCollatorForPairRank, HFSummary, WebGPT
 from torch.utils.data import DataLoader
-from rank_datasets import WebGPT, HFSummary, DataCollatorForPairRank
-from experimental_dataset import HFSummaryQuality, DataCollatorForSummaryScore
+from transformers import AutoTokenizer
+
 
 def test_hfsummary():
-    
+
     tokenizer = AutoTokenizer.from_pretrained("bigscience/mt0-large")
     collate_fn = DataCollatorForPairRank(tokenizer, max_length=200)
-    dataset = HFSummary('train')
+    dataset = HFSummary("train")
     print(len(dataset))
     dataloader = DataLoader(dataset, collate_fn=collate_fn, batch_size=8)
     for batch in dataloader:
-        batch['input_ids'].shape
- 
+        batch["input_ids"].shape
+
 
 def test_webgpt():
-    
+
     tokenizer = AutoTokenizer.from_pretrained("bigscience/mt0-large")
     collate_fn = DataCollatorForPairRank(tokenizer, max_length=200)
     dataset = WebGPT()
     dataloader = DataLoader(dataset, collate_fn=collate_fn, batch_size=32)
     for batch in dataloader:
-        print(batch['input_ids'].shape)
+        print(batch["input_ids"].shape)
 
 
 def test_hf_quality():
 
     tokenizer = AutoTokenizer.from_pretrained("bigscience/mt0-large")
     collate_fn = DataCollatorForSummaryScore(tokenizer, max_length=200)
-    dataset = HFSummaryQuality('validation', tokenizer)
+    dataset = HFSummaryQuality("validation", tokenizer)
     dataloader = DataLoader(dataset, collate_fn=collate_fn, batch_size=32)
     for batch in dataloader:
-        print(batch['input_ids'].shape)
-
+        print(batch["input_ids"].shape)
 
 
 if __name__ == "__main__":
     test_hf_quality()
-    # test_webgpt()
\ No newline at end of file
+    # test_webgpt()
diff --git a/model/reward/instructor/trainer.py b/model/reward/instructor/trainer.py
index c8063cf7..0e98e4c5 100644
--- a/model/reward/instructor/trainer.py
+++ b/model/reward/instructor/trainer.py
@@ -1,32 +1,44 @@
+# -*- coding: utf-8 -*-
 import os
-os.environ['WANDB_PROJECT'] = 'reward-model'
-import torch
-import yaml
-import evaluate
-from typing import Any, Callable, List, Optional, Tuple, Union, Dict
-from torch import nn
 from argparse import ArgumentParser
-import numpy as np
 from dataclasses import dataclass
-from torch.utils.data import Dataset, ConcatDataset
-from transformers import AutoModelForSequenceClassification
-from transformers import Trainer, PreTrainedModel, TrainingArguments, DataCollator, EvalPrediction, TrainerCallback, PreTrainedTokenizerBase
-from rank_datasets import DataCollatorForPairRank, WebGPT, HFSummary
-from utils import get_tokenizer, train_val_dataset, freeze_top_n_layers, argument_parsing
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+
+import evaluate
+import numpy as np
+import torch
+from rank_datasets import DataCollatorForPairRank, HFSummary, WebGPT
+from torch import nn
+from torch.utils.data import ConcatDataset, Dataset
+from transformers import (
+    AutoModelForSequenceClassification,
+    DataCollator,
+    EvalPrediction,
+    PreTrainedModel,
+    PreTrainedTokenizerBase,
+    Trainer,
+    TrainerCallback,
+    TrainingArguments,
+)
+from utils import argument_parsing, freeze_top_n_layers, get_tokenizer, train_val_dataset
+
+os.environ["WANDB_PROJECT"] = "reward-model"
 
 accuracy = evaluate.load("accuracy")
 parser = ArgumentParser()
-parser.add_argument('config', type=str)
+parser.add_argument("config", type=str)
+
 
 @dataclass
 class CustomTrainingArguments(TrainingArguments):
-    loss_function: str='rank'
+    loss_function: str = "rank"
 
 
 def compute_metrics(eval_pred):
     predictions, _ = eval_pred
     predictions = np.argmax(predictions, axis=1)
-    return accuracy.compute(predictions=predictions, references=[0]*predictions.shape[0])
+    return accuracy.compute(predictions=predictions, references=[0] * predictions.shape[0])
+
 
 class RankLoss(nn.Module):
     def __init__(self, eps=1e-8) -> None:
@@ -39,27 +51,41 @@ class RankLoss(nn.Module):
 
 
 class RankTrainer(Trainer):
-    def __init__(self, model: Union[PreTrainedModel, nn.Module] = None,
-                 args: TrainingArguments = None,
-                 data_collator: Optional[DataCollator] = None,
-                 train_dataset: Optional[Dataset] = None,
-                 eval_dataset: Optional[Dataset] = None,
-                 tokenizer: Optional[PreTrainedTokenizerBase] = None,
-                 model_init: Callable[[], PreTrainedModel] = None,
-                 compute_metrics: Optional[Callable[[EvalPrediction], Dict]] = None,
-                 callbacks: Optional[List[TrainerCallback]] = None,
-                 optimizers: Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None),
-                 preprocess_logits_for_metrics: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] = None):
-        super().__init__(model, args, data_collator, train_dataset, eval_dataset, tokenizer,
-                         model_init, compute_metrics, callbacks, optimizers, preprocess_logits_for_metrics)
-        self.loss_fct = RankLoss() if args.loss_function == 'rank' else nn.CrossEntropyLoss()
+    def __init__(
+        self,
+        model: Union[PreTrainedModel, nn.Module] = None,
+        args: TrainingArguments = None,
+        data_collator: Optional[DataCollator] = None,
+        train_dataset: Optional[Dataset] = None,
+        eval_dataset: Optional[Dataset] = None,
+        tokenizer: Optional[PreTrainedTokenizerBase] = None,
+        model_init: Callable[[], PreTrainedModel] = None,
+        compute_metrics: Optional[Callable[[EvalPrediction], Dict]] = None,
+        callbacks: Optional[List[TrainerCallback]] = None,
+        optimizers: Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None),
+        preprocess_logits_for_metrics: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] = None,
+    ):
+        super().__init__(
+            model,
+            args,
+            data_collator,
+            train_dataset,
+            eval_dataset,
+            tokenizer,
+            model_init,
+            compute_metrics,
+            callbacks,
+            optimizers,
+            preprocess_logits_for_metrics,
+        )
+        self.loss_fct = RankLoss() if args.loss_function == "rank" else nn.CrossEntropyLoss()
         self.loss_function = args.loss_function
 
     def compute_loss(self, model, inputs, return_outputs=False):
         # forward pass
         outputs = model(**inputs)
         logits = outputs.get("logits").view(-1, 2)
-        if self.loss_function == 'rank':
+        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))
@@ -70,17 +96,20 @@ class RankTrainer(Trainer):
         inputs = self._prepare_inputs(inputs)
         outputs = model(**inputs)
         logits = outputs.get("logits").view(-1, 2)
-        if self.loss_function == 'rank':
+        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]]:
+    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
@@ -93,54 +122,57 @@ class RankTrainer(Trainer):
 
         return (loss, logits, labels)
 
+
 if __name__ == "__main__":
     training_conf = argument_parsing(parser)
 
-    model_name = training_conf['model_name']
-    model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=1, problem_type='regression')
-    if 'freeze_layer' in training_conf:
-        num_layer = training_conf['freeze_layer']
+    model_name = training_conf["model_name"]
+    model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=1, problem_type="regression")
+    if "freeze_layer" in training_conf:
+        num_layer = training_conf["freeze_layer"]
         model = freeze_top_n_layers(model, num_layer)
         model_parameters = filter(lambda p: p.requires_grad, model.parameters())
         params = sum([np.prod(p.size()) for p in model_parameters])
-        print('Number of trainable : {}M'.format(int(params/1e6)))
+        print("Number of trainable : {}M".format(int(params / 1e6)))
 
     tokenizer = get_tokenizer(model_name)
     args = CustomTrainingArguments(
         output_dir=f"{model_name}-finetuned",
-        num_train_epochs=training_conf['num_train_epochs'],
+        num_train_epochs=training_conf["num_train_epochs"],
         warmup_steps=500,
-        loss_function=training_conf['loss'],
-        learning_rate=training_conf['learning_rate'],
+        loss_function=training_conf["loss"],
+        learning_rate=training_conf["learning_rate"],
         # half_precision_backend="apex",
         fp16=True,
-        gradient_checkpointing=training_conf['gradient_checkpointing'],
-        gradient_accumulation_steps=training_conf['gradient_accumulation_steps'],
-        per_device_train_batch_size=training_conf['per_device_train_batch_size'],
-        per_device_eval_batch_size=training_conf['per_device_eval_batch_size'],
+        gradient_checkpointing=training_conf["gradient_checkpointing"],
+        gradient_accumulation_steps=training_conf["gradient_accumulation_steps"],
+        per_device_train_batch_size=training_conf["per_device_train_batch_size"],
+        per_device_eval_batch_size=training_conf["per_device_eval_batch_size"],
         weight_decay=0.01,
         max_grad_norm=2.0,
         logging_steps=10,
         save_total_limit=4,
-        evaluation_strategy='steps',
-        eval_steps=training_conf['eval_steps'],
+        evaluation_strategy="steps",
+        eval_steps=training_conf["eval_steps"],
         save_steps=1000,
-        report_to='wandb'
+        report_to="wandb",
     )
     train_datasets, evals = [], {}
-    if 'webgpt' in training_conf['datasets']:
+    if "webgpt" in training_conf["datasets"]:
         web_dataset = WebGPT()
         train, eval = train_val_dataset(web_dataset)
         train_datasets.append(train)
-        evals['webgpt'] = eval
-    if 'hfsummary' in training_conf['datasets']:
-        sum_train = HFSummary(split='train')
+        evals["webgpt"] = eval
+    if "hfsummary" in training_conf["datasets"]:
+        sum_train = HFSummary(split="train")
         train_datasets.append(sum_train)
-        sum_eval = HFSummary(split='valid1')
+        sum_eval = HFSummary(split="valid1")
         assert len(sum_eval) > 0
-        evals['hfsummary'] = sum_eval
+        evals["hfsummary"] = sum_eval
     train = ConcatDataset(train_datasets)
-    collate_fn = DataCollatorForPairRank(tokenizer, max_length=training_conf['max_length'], drop_token_type= 'galactica' in model_name)
+    collate_fn = DataCollatorForPairRank(
+        tokenizer, max_length=training_conf["max_length"], drop_token_type="galactica" in model_name
+    )
     assert len(evals) > 0
     trainer = RankTrainer(
         model,
@@ -149,6 +181,6 @@ if __name__ == "__main__":
         eval_dataset=eval,
         data_collator=collate_fn,
         tokenizer=tokenizer,
-        compute_metrics=compute_metrics
+        compute_metrics=compute_metrics,
     )
     trainer.train()
diff --git a/model/reward/instructor/utils.py b/model/reward/instructor/utils.py
index d59bb13c..9441ddb9 100644
--- a/model/reward/instructor/utils.py
+++ b/model/reward/instructor/utils.py
@@ -1,96 +1,100 @@
+# -*- coding: utf-8 -*-
 import re
+
 import yaml
-from torch.utils.data import Subset
 from sklearn.model_selection import train_test_split
+from torch.utils.data import Subset
 from transformers import AutoTokenizer
 
-re_reference_remove = re.compile(r'\[([0-9])+\]|\[([0-9])+,([0-9])+\]')
+re_reference_remove = re.compile(r"\[([0-9])+\]|\[([0-9])+,([0-9])+\]")
+
 
 def webgpt_return_format(row):
-    if row['score_0'] >= row['score_1']:
+    if row["score_0"] >= row["score_1"]:
         # remove this to prevent information leak, since we are not using reference
         return {
-                'question': row['question']['full_text'],
-                     'pos': re_reference_remove.sub('', row['answer_0']),
-                     'neg': re_reference_remove.sub('', row['answer_1'])
-                }
+            "question": row["question"]["full_text"],
+            "pos": re_reference_remove.sub("", row["answer_0"]),
+            "neg": re_reference_remove.sub("", row["answer_1"]),
+        }
 
     return {
-            'question': row['question']['full_text'],
-                 'pos': re_reference_remove.sub('', row['answer_1']),
-                 'neg': re_reference_remove.sub('', row['answer_0'])
-            }
+        "question": row["question"]["full_text"],
+        "pos": re_reference_remove.sub("", row["answer_1"]),
+        "neg": re_reference_remove.sub("", row["answer_0"]),
+    }
 
 
 def get_tokenizer(tokenizer_name):
     tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
-    if 'galactica' in tokenizer_name:
-        tokenizer.add_special_tokens({'pad_token':'<pad>', 'eos_token': '</s>' })
+    if "galactica" in tokenizer_name:
+        tokenizer.add_special_tokens({"pad_token": "<pad>", "eos_token": "</s>"})
 
     return tokenizer
 
 
-
 def train_val_dataset(dataset, val_split=0.2):
-    train_idx, val_idx = train_test_split(list(range(len(dataset))), 
-        test_size=val_split, random_state=666, shuffle=True)
+    train_idx, val_idx = train_test_split(
+        list(range(len(dataset))), test_size=val_split, random_state=666, shuffle=True
+    )
     # [3879, 11479, 8341, 9177, 10798, 18177, 5735, 15669, 4837, 2760]
     print(val_idx[:10])
     # [13582, 5919, 11875, 7373, 19135, 13706, 8555, 15788, 15005, 15209]
     print(train_idx[:10])
     return Subset(dataset, train_idx), Subset(dataset, val_idx)
 
+
 def freeze_top_n_layers(model, target_layers):
     # its possible we can simply detect which module is a ModuleList
     # and simply freeze the module without doing string parsing
     for name, param in model.named_parameters():
-        if 'embed' in name:
+        if "embed" in name:
             param.requires_grad = False
-        elif '.layer' in name or '.h.' in name:
-            tokens = name.split('.')
+        elif ".layer" in name or ".h." in name:
+            tokens = name.split(".")
             idx = 0
             for token in tokens:
-                if 'layer' in token or token == 'h':
+                if "layer" in token or token == "h":
                     break
                 idx += 1
             if idx >= len(tokens):
                 continue
 
-            layer_ = int(tokens[idx+1])
+            layer_ = int(tokens[idx + 1])
             if layer_ < target_layers:
                 # print('freeze ', layer_, name)
                 param.requires_grad = False
     return model
 
+
 def argument_parsing(parser):
     default_params = {
-        'num_train_epochs': 4,
-        'learning_rate': 3e-5,
-        'eval_steps': 500,
-        'loss': 'rank',
-        'max_length': 440,
-        'per_device_eval_batch_size': 5,
-        'per_device_train_batch_size': 8,
-        'gradient_accumulation_steps': 8,
-        'gradient_checkpointing': False,
-        'datasets': ['webgpt']
+        "num_train_epochs": 4,
+        "learning_rate": 3e-5,
+        "eval_steps": 500,
+        "loss": "rank",
+        "max_length": 440,
+        "per_device_eval_batch_size": 5,
+        "per_device_train_batch_size": 8,
+        "gradient_accumulation_steps": 8,
+        "gradient_checkpointing": False,
+        "datasets": ["webgpt"],
     }
     args = parser.parse_args()
-    with open(args.config, 'r', encoding='utf-8') as f:
+    with open(args.config, "r", encoding="utf-8") as f:
         training_conf = yaml.safe_load(f.read())
 
-    params = { **default_params, **training_conf }
-    params['gradient_accumulation_steps'] = int(params['gradient_accumulation_steps'])
-    params['num_train_epochs'] = int(params['num_train_epochs'])
-    params['per_device_train_batch_size'] = int(params['per_device_train_batch_size'])
-    params['learning_rate'] = float(params['learning_rate'])
+    params = {**default_params, **training_conf}
+    params["gradient_accumulation_steps"] = int(params["gradient_accumulation_steps"])
+    params["num_train_epochs"] = int(params["num_train_epochs"])
+    params["per_device_train_batch_size"] = int(params["per_device_train_batch_size"])
+    params["learning_rate"] = float(params["learning_rate"])
     return params
 
 
-
 if __name__ == "__main__":
     from transformers import AutoModelForSequenceClassification
 
-    model = AutoModelForSequenceClassification.from_pretrained('bigscience/bloomz-560m')
+    model = AutoModelForSequenceClassification.from_pretrained("bigscience/bloomz-560m")
     freeze_top_n_layers(model, 10)
-    print(model.state_dict().keys())
\ No newline at end of file
+    print(model.state_dict().keys())

From 4d01704618e5588cd55a09857558355bf99abc10 Mon Sep 17 00:00:00 2001
From: theblackcat102 <theblackcat102@github.com>
Date: Sun, 1 Jan 2023 11:56:54 +0000
Subject: [PATCH 26/41] [fix] rerun pre-commit

---
 model/reward/instructor/README.md             |  11 +-
 model/reward/instructor/TODO.md               |  18 +--
 model/reward/instructor/cls_dataset.py        |  37 +++--
 .../configs/bloomz-560m-summary.yml           |   2 +-
 .../reward/instructor/configs/bloomz-560m.yml |   2 +-
 .../configs/electra-base-dis-webgpt.yml       |   2 +-
 .../instructor/configs/galactica-125m.yml     |   2 +-
 .../instructor/configs/galactica-1b.yml       |   2 +-
 .../test-galactica-125m-classification.yml    |   2 +-
 .../reward/instructor/experimental_dataset.py |  50 +++---
 model/reward/instructor/rank_datasets.py      | 104 ++++++------
 .../instructor/summary_quality_trainer.py     | 140 ++++++++++-------
 model/reward/instructor/tests/test_dataset.py |  27 ++--
 model/reward/instructor/trainer.py            | 148 +++++++++++-------
 model/reward/instructor/utils.py              |  84 +++++-----
 15 files changed, 337 insertions(+), 294 deletions(-)

diff --git a/model/reward/instructor/README.md b/model/reward/instructor/README.md
index 31c25371..73a872a0 100644
--- a/model/reward/instructor/README.md
+++ b/model/reward/instructor/README.md
@@ -2,7 +2,6 @@
 
 Trainer code based on huggingface. Compatible with deepspeed or accelerate
 
-
 Requirements
 
 ```
@@ -15,12 +14,10 @@ torch==1.12
 
 Start training reward model
 
-
 ```bash
 python trainer.py configs/electra-base-dis-webgpt.yml
 ```
 
-
 Additional axis labeling, this outputs a 4 summary quality evaluation metrics (score are normalized to 0-1 )
 
 ```bash
@@ -29,13 +26,13 @@ python summary_quality_trainer.py configs/test-bloomz-560m-quality.yml
 
 The four summary are :
 
-* overall
+- overall
 
-* accuracy
+- accuracy
 
-* coverage
+- coverage
 
-* coherence
+- coherence
 
 ## Dataset
 
diff --git a/model/reward/instructor/TODO.md b/model/reward/instructor/TODO.md
index 1e653922..ed33b3c0 100644
--- a/model/reward/instructor/TODO.md
+++ b/model/reward/instructor/TODO.md
@@ -1,23 +1,19 @@
-
 Some other reward features we can use
 
-0. Finish classifcation feature 
+0. Finish classifcation feature
 
 1. Summaries from human feedback
 
-* use `confidence` score into the RM learning, ensure the output rank score correlates with confidence
+- use `confidence` score into the RM learning, ensure the output rank score correlates with confidence
 
-* each labeling has a labeling `note`, basically comments by labeler, not sure what else we can use
+- each labeling has a labeling `note`, basically comments by labeler, not sure what else we can use
 
-* ~~Use the score for "overall", "accuracy", "coverage", "coherence" from axis/evals to train an addition model (rank additional aspect of the policy model)~~
-
-    * this should be placed under experimental_dataset.py
+- ~~Use the score for "overall", "accuracy", "coverage", "coherence" from axis/evals to train an addition model (rank additional aspect of the policy model)~~
 
+  - this should be placed under experimental_dataset.py
 
 2. Add support for anthropic dataset
 
-* anthropic dataset is more like a conversation tree which is much complex than simply question-answer schema
-
-    * this is basically a MCTS from alphazero.
-
+- anthropic dataset is more like a conversation tree which is much complex than simply question-answer schema
 
+  - this is basically a MCTS from alphazero.
diff --git a/model/reward/instructor/cls_dataset.py b/model/reward/instructor/cls_dataset.py
index ff824d19..09aa821b 100644
--- a/model/reward/instructor/cls_dataset.py
+++ b/model/reward/instructor/cls_dataset.py
@@ -1,32 +1,34 @@
-'''
+# -*- coding: utf-8 -*-
+"""
 
     classification based ranking
 
-'''
-import os
+"""
 import json
+import os
 import random
-import torch
-import numpy as np
+
 from dataset import load_dataset
 from torch.utils.data import Dataset
+
 from .utils import webgpt_return_format
 
+
 class WebGPTDataset(Dataset):
-    def __init__(self, mode='train', index_cache='dataset/webgpt_train_idx.pt', additional_dataset=None) -> None:
+    def __init__(self, mode="train", index_cache="dataset/webgpt_train_idx.pt", additional_dataset=None) -> None:
         super().__init__()
-        '''
+        """
             mode : train or val, used for validation purpose, has nothing to do with original split
             additional_dataset : a list of jsonline format with idx, question and texts (generate candidates)
                 idx : must match the index you iterate from comparison enumerate order
                 question : for validation purpose
                 texts : list of K generate results from the question prompt
-        '''
-        os.makedirs('dataset', exist_ok=True)
+        """
+        os.makedirs("dataset", exist_ok=True)
         dataset = load_dataset("openai/webgpt_comparisons")
         self.dataset = []
         self.dataset_index = []
-        for idx, row in enumerate(dataset['train']):
+        for idx, row in enumerate(dataset["train"]):
             self.dataset.append(webgpt_return_format(row))
 
         # since this dataset was generated from 176B GPT-3
@@ -36,17 +38,17 @@ class WebGPTDataset(Dataset):
         if additional_dataset is not None:
             self.sample_additional = True
             self.additional = {}
-            with open(additional_dataset, 'r') as f:
+            with open(additional_dataset, "r") as f:
                 for line in f:
                     row = json.loads(line)
-                    if row['idx'] in self.dataset_index:
-                        self.additional[row['idx']] = row['negatives']
+                    if row["idx"] in self.dataset_index:
+                        self.additional[row["idx"]] = row["negatives"]
             if len(self.additional) != len(self.dataset_index):
                 for match_idx in self.dataset_index:
                     if match_idx in self.additional:
                         continue
 
-                    idx = match_idx-900
+                    idx = match_idx - 900
                     while idx not in self.additional:
                         idx -= 1
                     self.additional[match_idx] = self.additional[idx]
@@ -57,10 +59,7 @@ class WebGPTDataset(Dataset):
     def __getitem__(self, index):
         row = self.dataset[index]
         if not self.sample_additional:
-            return row['question'], row['pos'], row['neg']
+            return row["question"], row["pos"], row["neg"]
 
         gen_neg = random.choice(self.additional[self.dataset_index[index]])
-        return row['question'], row['pos'], row['neg'], gen_neg
-
-
-
+        return row["question"], row["pos"], row["neg"], gen_neg
diff --git a/model/reward/instructor/configs/bloomz-560m-summary.yml b/model/reward/instructor/configs/bloomz-560m-summary.yml
index a02f4e4a..55ed6cd1 100644
--- a/model/reward/instructor/configs/bloomz-560m-summary.yml
+++ b/model/reward/instructor/configs/bloomz-560m-summary.yml
@@ -6,4 +6,4 @@ max_length: 600
 freeze_layer: 12
 num_train_epochs: 2
 datasets:
-  - hfsummary
\ No newline at end of file
+  - hfsummary
diff --git a/model/reward/instructor/configs/bloomz-560m.yml b/model/reward/instructor/configs/bloomz-560m.yml
index c8f55746..bf3f14dd 100644
--- a/model/reward/instructor/configs/bloomz-560m.yml
+++ b/model/reward/instructor/configs/bloomz-560m.yml
@@ -7,4 +7,4 @@ freeze_layer: 12
 num_train_epochs: 2
 datasets:
   - webgpt
-  - hfsummary
\ No newline at end of file
+  - hfsummary
diff --git a/model/reward/instructor/configs/electra-base-dis-webgpt.yml b/model/reward/instructor/configs/electra-base-dis-webgpt.yml
index fc168b63..89200fe1 100644
--- a/model/reward/instructor/configs/electra-base-dis-webgpt.yml
+++ b/model/reward/instructor/configs/electra-base-dis-webgpt.yml
@@ -1,3 +1,3 @@
 model_name: google/electra-large-discriminator
 learning_rate: 3e-5
-max_length: 300
\ No newline at end of file
+max_length: 300
diff --git a/model/reward/instructor/configs/galactica-125m.yml b/model/reward/instructor/configs/galactica-125m.yml
index 55e093f5..13dbdfbe 100644
--- a/model/reward/instructor/configs/galactica-125m.yml
+++ b/model/reward/instructor/configs/galactica-125m.yml
@@ -10,4 +10,4 @@ max_length: 512
 num_train_epochs: 2
 datasets:
   - webgpt
-  - hfsummary
\ No newline at end of file
+  - hfsummary
diff --git a/model/reward/instructor/configs/galactica-1b.yml b/model/reward/instructor/configs/galactica-1b.yml
index 5a094520..8ffd74e9 100644
--- a/model/reward/instructor/configs/galactica-1b.yml
+++ b/model/reward/instructor/configs/galactica-1b.yml
@@ -11,4 +11,4 @@ max_length: 400
 num_train_epochs: 2
 datasets:
   - webgpt
-  - hfsummary
\ No newline at end of file
+  - hfsummary
diff --git a/model/reward/instructor/configs/test-galactica-125m-classification.yml b/model/reward/instructor/configs/test-galactica-125m-classification.yml
index 1ad1f47c..e36efcf3 100644
--- a/model/reward/instructor/configs/test-galactica-125m-classification.yml
+++ b/model/reward/instructor/configs/test-galactica-125m-classification.yml
@@ -11,4 +11,4 @@ max_length: 128
 num_train_epochs: 2
 datasets:
   - webgpt
-  - hfsummary
\ No newline at end of file
+  - hfsummary
diff --git a/model/reward/instructor/experimental_dataset.py b/model/reward/instructor/experimental_dataset.py
index 47d20d64..28f62967 100644
--- a/model/reward/instructor/experimental_dataset.py
+++ b/model/reward/instructor/experimental_dataset.py
@@ -1,4 +1,5 @@
-'''
+# -*- coding: utf-8 -*-
+"""
     HFSummary
 
         I want to train a multi regression model on axis_evals dataset mainly we can estimate the score of these score
@@ -7,15 +8,16 @@
 
         Should be better than just a preference score
 
-'''
-import torch
-from typing import Optional, Union
-import numpy as np
+"""
 from collections import defaultdict
-from datasets import load_dataset
 from dataclasses import dataclass
+from typing import Optional, Union
+
+import numpy as np
+import torch
+from datasets import load_dataset
 from torch.utils.data import Dataset
-from transformers.tokenization_utils_base import PreTrainedTokenizerBase, PaddingStrategy
+from transformers.tokenization_utils_base import PaddingStrategy, PreTrainedTokenizerBase
 
 
 @dataclass
@@ -25,12 +27,13 @@ class DataCollatorForSummaryScore:
     Data collator that will dynamically pad the inputs for multiple choice received.
 
     """
+
     tokenizer: PreTrainedTokenizerBase
     num_choices: int = 2
     padding: Union[bool, str, PaddingStrategy] = True
     max_length: Optional[int] = None
     pad_to_multiple_of: Optional[int] = None
-    drop_token_type: bool = False # galactica
+    drop_token_type: bool = False  # galactica
 
     def __call__(self, batch):
 
@@ -48,17 +51,17 @@ class DataCollatorForSummaryScore:
             return_tensors="pt",
         )
         if self.drop_token_type:
-            batch_feature.pop('token_type_ids')
+            batch_feature.pop("token_type_ids")
         # batch = {k: v.view(batch_size, self.num_choices, -1) for k, v in batch.items()}
-        batch_feature['labels'] = torch.from_numpy(np.array(labels)).float()
+        batch_feature["labels"] = torch.from_numpy(np.array(labels)).float()
         return batch_feature
 
 
 class HFSummaryQuality(Dataset):
     def __init__(self, split, tokenizer, max_length=300) -> None:
         super().__init__()
-        assert split in ('validation', 'test')
-        dataset = load_dataset('Tristan/summarize_from_feedback', 'axis')[split]
+        assert split in ("validation", "test")
+        dataset = load_dataset("Tristan/summarize_from_feedback", "axis")[split]
         self.max_length = max_length
         mean_scores = defaultdict(list)
         self.contexts = []
@@ -66,22 +69,21 @@ class HFSummaryQuality(Dataset):
         self.labels = []
         for data in dataset:
 
-            if 'article' in data['info'] and \
-                data['info']['article'] is not None:
-                context = data['info']['article']
-            elif 'post' in data['info']:
-                context = data['info']['post']
+            if "article" in data["info"] and data["info"]["article"] is not None:
+                context = data["info"]["article"]
+            elif "post" in data["info"]:
+                context = data["info"]["post"]
             self.contexts.append(context)
 
-            response = data['summary']['text']
+            response = data["summary"]["text"]
             self.responses.append(response)
-            self.labels.append(data['summary']['axes'])
-            for axis, score in data['summary']['axes'].items():
+            self.labels.append(data["summary"]["axes"])
+            for axis, score in data["summary"]["axes"].items():
                 if score is not None:
                     mean_scores[axis].append(score)
 
-        self.label2idx = { key: idx for idx, key in enumerate(mean_scores.keys()) }
-        self.label2mean = { key: np.mean(scores) for key, scores in mean_scores.items() }
+        self.label2idx = {key: idx for idx, key in enumerate(mean_scores.keys())}
+        self.label2mean = {key: np.mean(scores) for key, scores in mean_scores.items()}
         self.tokenizer = tokenizer
         print(self.label2idx)
 
@@ -94,7 +96,5 @@ class HFSummaryQuality(Dataset):
         response = self.responses[index]
         labels = np.zeros(len(self.label2idx))
         for key, score in self.labels[index].items():
-            labels[self.label2idx[key]] = (self.label2mean[key] if score is None else score)/10
+            labels[self.label2idx[key]] = (self.label2mean[key] if score is None else score) / 10
         return self.tokenizer(context, response, truncation=True, max_length=self.max_length), labels
-
-
diff --git a/model/reward/instructor/rank_datasets.py b/model/reward/instructor/rank_datasets.py
index f38885e4..99ba9955 100644
--- a/model/reward/instructor/rank_datasets.py
+++ b/model/reward/instructor/rank_datasets.py
@@ -1,4 +1,5 @@
-'''
+# -*- coding: utf-8 -*-
+"""
     author: theblackcat102
 
     Dataset output format from __getitem__
@@ -17,13 +18,15 @@
         inferior than the human perference one
 
 
-'''
-from typing import Optional, Union
+"""
 from dataclasses import dataclass
+from typing import Optional, Union
+
 import numpy as np
-from torch.utils.data import Dataset
 from datasets import load_dataset
-from transformers.tokenization_utils_base import PreTrainedTokenizerBase, PaddingStrategy
+from torch.utils.data import Dataset
+from transformers.tokenization_utils_base import PaddingStrategy, PreTrainedTokenizerBase
+
 
 @dataclass
 class DataCollatorForPairRank:
@@ -32,12 +35,13 @@ class DataCollatorForPairRank:
     Data collator that will dynamically pad the inputs for multiple choice received.
 
     """
+
     tokenizer: PreTrainedTokenizerBase
     num_choices: int = 2
     padding: Union[bool, str, PaddingStrategy] = True
     max_length: Optional[int] = None
     pad_to_multiple_of: Optional[int] = None
-    drop_token_type: bool = False # galactica
+    drop_token_type: bool = False  # galactica
 
     def __call__(self, features):
 
@@ -45,12 +49,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, max_length=self.max_length))
-                flatten_features.append(self.tokenizer(question, neg,
-                    truncation=True, max_length=self.max_length))
+                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(
             flatten_features,
             padding=self.padding,
@@ -59,13 +61,12 @@ class DataCollatorForPairRank:
             return_tensors="pt",
         )
         if self.drop_token_type:
-            batch.pop('token_type_ids')
+            batch.pop("token_type_ids")
         # batch = {k: v.view(batch_size, self.num_choices, -1) for k, v in batch.items()}
         return batch
 
 
 class WebGPT(Dataset):
-
     def __init__(self) -> None:
         super().__init__()
 
@@ -74,23 +75,19 @@ class WebGPT(Dataset):
         # using prompt as our index will allows us
         # to add additional generated prompt later
         self.index2question = {}
-        for row in dataset['train']:
-            question = row['question']['full_text']
+        for row in dataset["train"]:
+            question = row["question"]["full_text"]
             if question not in self.index2question:
                 self.index2question[len(self.index2question)] = question
 
             if question not in questions:
                 questions[question] = []
 
-            if row['score_0'] > row['score_1']:
+            if row["score_0"] > row["score_1"]:
                 # not going to risk it
-                questions[question].append((
-                    row['answer_0'], row['answer_1']
-                ))
+                questions[question].append((row["answer_0"], row["answer_1"]))
             else:
-                questions[question].append((
-                    row['answer_1'], row['answer_0']
-                ))
+                questions[question].append((row["answer_1"], row["answer_0"]))
 
         self.questions = questions
 
@@ -104,61 +101,55 @@ class WebGPT(Dataset):
         return question, rows
 
 
-
-
 class HFSummary(Dataset):
-    '''
-        Human feedback data from OpenAI
-        https://github.com/openai/summarize-from-feedback
-        
-        labeling method : pair comparison, 0 or 1
+    """
+    Human feedback data from OpenAI
+    https://github.com/openai/summarize-from-feedback
 
-    '''
-    def __init__(self, split='train',
-        conf_threshold=-1,
-        max_comparison_per_sample=3) -> None:
+    labeling method : pair comparison, 0 or 1
+
+    """
+
+    def __init__(self, split="train", conf_threshold=-1, max_comparison_per_sample=3) -> None:
         super().__init__()
-        assert split in ('train', 'valid1', 'valid2', 'test')
+        assert split in ("train", "valid1", "valid2", "test")
         summaries = {}
         # using prompt as our index will allows us
         # to add additional generated prompt later
         self.index2summary = {}
         self.max_comparison_per_sample = max_comparison_per_sample
-        major_split = split if 'train' == split else 'validation'
-        dataset = load_dataset('Tristan/summarize_from_feedback', 'comparisons')[major_split]
+        major_split = split if "train" == split else "validation"
+        dataset = load_dataset("Tristan/summarize_from_feedback", "comparisons")[major_split]
         for data in dataset:
-            if 'extra' in data and \
-                'confidence' in data['extra'] and \
-                data['extra']['confidence'] is not None and \
-                conf_threshold > data['extra']['confidence']:
-                print('skipping {}'.format(data['info']['id']))
+            if (
+                "extra" in data
+                and "confidence" in data["extra"]
+                and data["extra"]["confidence"] is not None
+                and conf_threshold > data["extra"]["confidence"]
+            ):
+                print("skipping {}".format(data["info"]["id"]))
                 continue
 
-            if split != 'train' and split != data['split']:
+            if split != "train" and split != data["split"]:
                 continue
 
-            if 'article' in data['info'] and \
-                data['info']['article'] is not None:
-                context = data['info']['article']
-            elif 'post' in data['info']:
-                context = data['info']['post']
-
+            if "article" in data["info"] and data["info"]["article"] is not None:
+                context = data["info"]["article"]
+            elif "post" in data["info"]:
+                context = data["info"]["post"]
 
             if context not in self.index2summary:
                 self.index2summary[len(self.index2summary)] = context
-            
+
             if context not in summaries:
                 summaries[context] = []
 
-            pos, neg = (0, 1) if data['choice'] == 0 else (1, 0)
-            summaries[context].append((
-                data['summaries'][pos]['text'],
-                data['summaries'][neg]['text']
-            ))
+            pos, neg = (0, 1) if data["choice"] == 0 else (1, 0)
+            summaries[context].append((data["summaries"][pos]["text"], data["summaries"][neg]["text"]))
 
         self.summaries = summaries
 
-        self.postfix_prompt = ' TLDR;'
+        self.postfix_prompt = " TLDR;"
 
     def __len__(self):
         return len(self.index2summary)
@@ -172,5 +163,4 @@ 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+self.postfix_prompt, [ 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/summary_quality_trainer.py b/model/reward/instructor/summary_quality_trainer.py
index a6604819..88bf1abf 100644
--- a/model/reward/instructor/summary_quality_trainer.py
+++ b/model/reward/instructor/summary_quality_trainer.py
@@ -1,46 +1,72 @@
+# -*- coding: utf-8 -*-
 import os
-os.environ['WANDB_PROJECT'] = 'quality-scoring'
-import torch
-import yaml
-import evaluate
-from typing import Any, Callable, List, Optional, Tuple, Union, Dict
-from torch import nn
 from argparse import ArgumentParser
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+
+import evaluate
 import numpy as np
+import torch
+from experimental_dataset import DataCollatorForSummaryScore, HFSummaryQuality
+from torch import nn
 from torch.utils.data import Dataset
-from transformers import AutoModelForSequenceClassification
-from transformers import Trainer, PreTrainedModel, TrainingArguments, DataCollator, EvalPrediction, TrainerCallback, PreTrainedTokenizerBase
-from experimental_dataset import HFSummaryQuality, DataCollatorForSummaryScore
-from utils import get_tokenizer, train_val_dataset, freeze_top_n_layers, argument_parsing
+from transformers import (
+    AutoModelForSequenceClassification,
+    DataCollator,
+    EvalPrediction,
+    PreTrainedModel,
+    PreTrainedTokenizerBase,
+    Trainer,
+    TrainerCallback,
+    TrainingArguments,
+)
+from utils import argument_parsing, freeze_top_n_layers, get_tokenizer
+
+os.environ["WANDB_PROJECT"] = "quality-scoring"
 
 parser = ArgumentParser()
-parser.add_argument('config', type=str)
+parser.add_argument("config", type=str)
 
 accuracy = evaluate.load("mse")
+
+
 def compute_metrics(eval_pred):
     predictions, labels = eval_pred
     return accuracy.compute(predictions=predictions.flatten(), references=labels.flatten())
 
 
 class QualityTrainer(Trainer):
-    def __init__(self, model: Union[PreTrainedModel, nn.Module] = None,
-                 args: TrainingArguments = None,
-                 data_collator: Optional[DataCollator] = None,
-                 train_dataset: Optional[Dataset] = None,
-                 eval_dataset: Optional[Dataset] = None,
-                 tokenizer: Optional[PreTrainedTokenizerBase] = None,
-                 model_init: Callable[[], PreTrainedModel] = None,
-                 compute_metrics: Optional[Callable[[EvalPrediction], Dict]] = None,
-                 callbacks: Optional[List[TrainerCallback]] = None,
-                 optimizers: Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None),
-                 preprocess_logits_for_metrics: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] = None):
-        super().__init__(model, args, data_collator, train_dataset, eval_dataset, tokenizer,
-                         model_init, compute_metrics, callbacks, optimizers, preprocess_logits_for_metrics)
+    def __init__(
+        self,
+        model: Union[PreTrainedModel, nn.Module] = None,
+        args: TrainingArguments = None,
+        data_collator: Optional[DataCollator] = None,
+        train_dataset: Optional[Dataset] = None,
+        eval_dataset: Optional[Dataset] = None,
+        tokenizer: Optional[PreTrainedTokenizerBase] = None,
+        model_init: Callable[[], PreTrainedModel] = None,
+        compute_metrics: Optional[Callable[[EvalPrediction], Dict]] = None,
+        callbacks: Optional[List[TrainerCallback]] = None,
+        optimizers: Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None),
+        preprocess_logits_for_metrics: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] = None,
+    ):
+        super().__init__(
+            model,
+            args,
+            data_collator,
+            train_dataset,
+            eval_dataset,
+            tokenizer,
+            model_init,
+            compute_metrics,
+            callbacks,
+            optimizers,
+            preprocess_logits_for_metrics,
+        )
         self.loss_fct = nn.L1Loss()
         self.sigmoid = nn.Sigmoid()
 
     def compute_loss(self, model, inputs, return_outputs=False):
-        labels = inputs.pop('labels')
+        labels = inputs.pop("labels")
         # forward pass
         outputs = model(**inputs)
         logits = self.sigmoid(outputs.get("logits"))
@@ -50,75 +76,73 @@ class QualityTrainer(Trainer):
 
     def _compute_loss(self, model, inputs):
         inputs = self._prepare_inputs(inputs)
-        labels = inputs.pop('labels')
+        labels = inputs.pop("labels")
         outputs = model(**inputs)
         logits = self.sigmoid(outputs.get("logits"))
         loss = self.loss_fct(logits, labels)
 
         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]]:
+    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 = inputs['labels']
+        labels = inputs["labels"]
         if self.args.prediction_loss_only:
             return (loss, None, None)
 
         return (loss, logits, labels)
 
+
 if __name__ == "__main__":
     training_conf = argument_parsing(parser)
 
-    model_name = training_conf['model_name']
+    model_name = training_conf["model_name"]
     tokenizer = get_tokenizer(model_name)
-    collate_fn = DataCollatorForSummaryScore(tokenizer,
-        max_length=training_conf['max_length'],
-        drop_token_type= 'galactica' in model_name
+    collate_fn = DataCollatorForSummaryScore(
+        tokenizer, max_length=training_conf["max_length"], drop_token_type="galactica" in model_name
+    )
+    train = HFSummaryQuality(split="validation", tokenizer=tokenizer, max_length=training_conf["max_length"])
+    eval = HFSummaryQuality(split="test", tokenizer=tokenizer, max_length=training_conf["max_length"])
+    model = AutoModelForSequenceClassification.from_pretrained(
+        model_name, num_labels=len(train.label2idx), problem_type="regression"
     )
-    train = HFSummaryQuality(split='validation',
-        tokenizer=tokenizer,
-        max_length=training_conf['max_length']
-        )
-    eval = HFSummaryQuality(split='test',
-            tokenizer=tokenizer,
-            max_length=training_conf['max_length']
-        )
-    model = AutoModelForSequenceClassification.from_pretrained(model_name,
-        num_labels=len(train.label2idx), problem_type='regression')
 
-    if 'freeze_layer' in training_conf:
-        num_layer = training_conf['freeze_layer']
+    if "freeze_layer" in training_conf:
+        num_layer = training_conf["freeze_layer"]
         model = freeze_top_n_layers(model, num_layer)
         model_parameters = filter(lambda p: p.requires_grad, model.parameters())
         params = sum([np.prod(p.size()) for p in model_parameters])
-        print('Number of trainable : {}M'.format(int(params/1e6)))
+        print("Number of trainable : {}M".format(int(params / 1e6)))
 
     args = TrainingArguments(
         output_dir=f"{model_name}-finetuned",
-        num_train_epochs=training_conf['num_train_epochs'],
+        num_train_epochs=training_conf["num_train_epochs"],
         warmup_steps=500,
-        learning_rate=training_conf['learning_rate'],
+        learning_rate=training_conf["learning_rate"],
         # half_precision_backend="apex",
         fp16=True,
-        gradient_checkpointing=training_conf['gradient_checkpointing'],
-        gradient_accumulation_steps=training_conf['gradient_accumulation_steps'],
-        per_device_train_batch_size=training_conf['per_device_train_batch_size'],
-        per_device_eval_batch_size=training_conf['per_device_eval_batch_size'],
+        gradient_checkpointing=training_conf["gradient_checkpointing"],
+        gradient_accumulation_steps=training_conf["gradient_accumulation_steps"],
+        per_device_train_batch_size=training_conf["per_device_train_batch_size"],
+        per_device_eval_batch_size=training_conf["per_device_eval_batch_size"],
         weight_decay=0.01,
         max_grad_norm=2.0,
         logging_steps=10,
         save_total_limit=4,
-        evaluation_strategy='steps',
-        eval_steps=training_conf['eval_steps'],
+        evaluation_strategy="steps",
+        eval_steps=training_conf["eval_steps"],
         save_steps=1000,
-        report_to='wandb'
+        report_to="wandb",
     )
     trainer = QualityTrainer(
         model,
@@ -127,6 +151,6 @@ if __name__ == "__main__":
         eval_dataset=eval,
         data_collator=collate_fn,
         tokenizer=tokenizer,
-        compute_metrics=compute_metrics
+        compute_metrics=compute_metrics,
     )
     trainer.train()
diff --git a/model/reward/instructor/tests/test_dataset.py b/model/reward/instructor/tests/test_dataset.py
index 271db83c..f367a50d 100644
--- a/model/reward/instructor/tests/test_dataset.py
+++ b/model/reward/instructor/tests/test_dataset.py
@@ -1,40 +1,41 @@
-from transformers import AutoTokenizer
+# -*- coding: utf-8 -*-
+from experimental_dataset import DataCollatorForSummaryScore, HFSummaryQuality
+from rank_datasets import DataCollatorForPairRank, HFSummary, WebGPT
 from torch.utils.data import DataLoader
-from rank_datasets import WebGPT, HFSummary, DataCollatorForPairRank
-from experimental_dataset import HFSummaryQuality, DataCollatorForSummaryScore
+from transformers import AutoTokenizer
+
 
 def test_hfsummary():
-    
+
     tokenizer = AutoTokenizer.from_pretrained("bigscience/mt0-large")
     collate_fn = DataCollatorForPairRank(tokenizer, max_length=200)
-    dataset = HFSummary('train')
+    dataset = HFSummary("train")
     print(len(dataset))
     dataloader = DataLoader(dataset, collate_fn=collate_fn, batch_size=8)
     for batch in dataloader:
-        batch['input_ids'].shape
- 
+        batch["input_ids"].shape
+
 
 def test_webgpt():
-    
+
     tokenizer = AutoTokenizer.from_pretrained("bigscience/mt0-large")
     collate_fn = DataCollatorForPairRank(tokenizer, max_length=200)
     dataset = WebGPT()
     dataloader = DataLoader(dataset, collate_fn=collate_fn, batch_size=32)
     for batch in dataloader:
-        print(batch['input_ids'].shape)
+        print(batch["input_ids"].shape)
 
 
 def test_hf_quality():
 
     tokenizer = AutoTokenizer.from_pretrained("bigscience/mt0-large")
     collate_fn = DataCollatorForSummaryScore(tokenizer, max_length=200)
-    dataset = HFSummaryQuality('validation', tokenizer)
+    dataset = HFSummaryQuality("validation", tokenizer)
     dataloader = DataLoader(dataset, collate_fn=collate_fn, batch_size=32)
     for batch in dataloader:
-        print(batch['input_ids'].shape)
-
+        print(batch["input_ids"].shape)
 
 
 if __name__ == "__main__":
     test_hf_quality()
-    # test_webgpt()
\ No newline at end of file
+    # test_webgpt()
diff --git a/model/reward/instructor/trainer.py b/model/reward/instructor/trainer.py
index c8063cf7..0e98e4c5 100644
--- a/model/reward/instructor/trainer.py
+++ b/model/reward/instructor/trainer.py
@@ -1,32 +1,44 @@
+# -*- coding: utf-8 -*-
 import os
-os.environ['WANDB_PROJECT'] = 'reward-model'
-import torch
-import yaml
-import evaluate
-from typing import Any, Callable, List, Optional, Tuple, Union, Dict
-from torch import nn
 from argparse import ArgumentParser
-import numpy as np
 from dataclasses import dataclass
-from torch.utils.data import Dataset, ConcatDataset
-from transformers import AutoModelForSequenceClassification
-from transformers import Trainer, PreTrainedModel, TrainingArguments, DataCollator, EvalPrediction, TrainerCallback, PreTrainedTokenizerBase
-from rank_datasets import DataCollatorForPairRank, WebGPT, HFSummary
-from utils import get_tokenizer, train_val_dataset, freeze_top_n_layers, argument_parsing
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+
+import evaluate
+import numpy as np
+import torch
+from rank_datasets import DataCollatorForPairRank, HFSummary, WebGPT
+from torch import nn
+from torch.utils.data import ConcatDataset, Dataset
+from transformers import (
+    AutoModelForSequenceClassification,
+    DataCollator,
+    EvalPrediction,
+    PreTrainedModel,
+    PreTrainedTokenizerBase,
+    Trainer,
+    TrainerCallback,
+    TrainingArguments,
+)
+from utils import argument_parsing, freeze_top_n_layers, get_tokenizer, train_val_dataset
+
+os.environ["WANDB_PROJECT"] = "reward-model"
 
 accuracy = evaluate.load("accuracy")
 parser = ArgumentParser()
-parser.add_argument('config', type=str)
+parser.add_argument("config", type=str)
+
 
 @dataclass
 class CustomTrainingArguments(TrainingArguments):
-    loss_function: str='rank'
+    loss_function: str = "rank"
 
 
 def compute_metrics(eval_pred):
     predictions, _ = eval_pred
     predictions = np.argmax(predictions, axis=1)
-    return accuracy.compute(predictions=predictions, references=[0]*predictions.shape[0])
+    return accuracy.compute(predictions=predictions, references=[0] * predictions.shape[0])
+
 
 class RankLoss(nn.Module):
     def __init__(self, eps=1e-8) -> None:
@@ -39,27 +51,41 @@ class RankLoss(nn.Module):
 
 
 class RankTrainer(Trainer):
-    def __init__(self, model: Union[PreTrainedModel, nn.Module] = None,
-                 args: TrainingArguments = None,
-                 data_collator: Optional[DataCollator] = None,
-                 train_dataset: Optional[Dataset] = None,
-                 eval_dataset: Optional[Dataset] = None,
-                 tokenizer: Optional[PreTrainedTokenizerBase] = None,
-                 model_init: Callable[[], PreTrainedModel] = None,
-                 compute_metrics: Optional[Callable[[EvalPrediction], Dict]] = None,
-                 callbacks: Optional[List[TrainerCallback]] = None,
-                 optimizers: Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None),
-                 preprocess_logits_for_metrics: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] = None):
-        super().__init__(model, args, data_collator, train_dataset, eval_dataset, tokenizer,
-                         model_init, compute_metrics, callbacks, optimizers, preprocess_logits_for_metrics)
-        self.loss_fct = RankLoss() if args.loss_function == 'rank' else nn.CrossEntropyLoss()
+    def __init__(
+        self,
+        model: Union[PreTrainedModel, nn.Module] = None,
+        args: TrainingArguments = None,
+        data_collator: Optional[DataCollator] = None,
+        train_dataset: Optional[Dataset] = None,
+        eval_dataset: Optional[Dataset] = None,
+        tokenizer: Optional[PreTrainedTokenizerBase] = None,
+        model_init: Callable[[], PreTrainedModel] = None,
+        compute_metrics: Optional[Callable[[EvalPrediction], Dict]] = None,
+        callbacks: Optional[List[TrainerCallback]] = None,
+        optimizers: Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None),
+        preprocess_logits_for_metrics: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] = None,
+    ):
+        super().__init__(
+            model,
+            args,
+            data_collator,
+            train_dataset,
+            eval_dataset,
+            tokenizer,
+            model_init,
+            compute_metrics,
+            callbacks,
+            optimizers,
+            preprocess_logits_for_metrics,
+        )
+        self.loss_fct = RankLoss() if args.loss_function == "rank" else nn.CrossEntropyLoss()
         self.loss_function = args.loss_function
 
     def compute_loss(self, model, inputs, return_outputs=False):
         # forward pass
         outputs = model(**inputs)
         logits = outputs.get("logits").view(-1, 2)
-        if self.loss_function == 'rank':
+        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))
@@ -70,17 +96,20 @@ class RankTrainer(Trainer):
         inputs = self._prepare_inputs(inputs)
         outputs = model(**inputs)
         logits = outputs.get("logits").view(-1, 2)
-        if self.loss_function == 'rank':
+        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]]:
+    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
@@ -93,54 +122,57 @@ class RankTrainer(Trainer):
 
         return (loss, logits, labels)
 
+
 if __name__ == "__main__":
     training_conf = argument_parsing(parser)
 
-    model_name = training_conf['model_name']
-    model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=1, problem_type='regression')
-    if 'freeze_layer' in training_conf:
-        num_layer = training_conf['freeze_layer']
+    model_name = training_conf["model_name"]
+    model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=1, problem_type="regression")
+    if "freeze_layer" in training_conf:
+        num_layer = training_conf["freeze_layer"]
         model = freeze_top_n_layers(model, num_layer)
         model_parameters = filter(lambda p: p.requires_grad, model.parameters())
         params = sum([np.prod(p.size()) for p in model_parameters])
-        print('Number of trainable : {}M'.format(int(params/1e6)))
+        print("Number of trainable : {}M".format(int(params / 1e6)))
 
     tokenizer = get_tokenizer(model_name)
     args = CustomTrainingArguments(
         output_dir=f"{model_name}-finetuned",
-        num_train_epochs=training_conf['num_train_epochs'],
+        num_train_epochs=training_conf["num_train_epochs"],
         warmup_steps=500,
-        loss_function=training_conf['loss'],
-        learning_rate=training_conf['learning_rate'],
+        loss_function=training_conf["loss"],
+        learning_rate=training_conf["learning_rate"],
         # half_precision_backend="apex",
         fp16=True,
-        gradient_checkpointing=training_conf['gradient_checkpointing'],
-        gradient_accumulation_steps=training_conf['gradient_accumulation_steps'],
-        per_device_train_batch_size=training_conf['per_device_train_batch_size'],
-        per_device_eval_batch_size=training_conf['per_device_eval_batch_size'],
+        gradient_checkpointing=training_conf["gradient_checkpointing"],
+        gradient_accumulation_steps=training_conf["gradient_accumulation_steps"],
+        per_device_train_batch_size=training_conf["per_device_train_batch_size"],
+        per_device_eval_batch_size=training_conf["per_device_eval_batch_size"],
         weight_decay=0.01,
         max_grad_norm=2.0,
         logging_steps=10,
         save_total_limit=4,
-        evaluation_strategy='steps',
-        eval_steps=training_conf['eval_steps'],
+        evaluation_strategy="steps",
+        eval_steps=training_conf["eval_steps"],
         save_steps=1000,
-        report_to='wandb'
+        report_to="wandb",
     )
     train_datasets, evals = [], {}
-    if 'webgpt' in training_conf['datasets']:
+    if "webgpt" in training_conf["datasets"]:
         web_dataset = WebGPT()
         train, eval = train_val_dataset(web_dataset)
         train_datasets.append(train)
-        evals['webgpt'] = eval
-    if 'hfsummary' in training_conf['datasets']:
-        sum_train = HFSummary(split='train')
+        evals["webgpt"] = eval
+    if "hfsummary" in training_conf["datasets"]:
+        sum_train = HFSummary(split="train")
         train_datasets.append(sum_train)
-        sum_eval = HFSummary(split='valid1')
+        sum_eval = HFSummary(split="valid1")
         assert len(sum_eval) > 0
-        evals['hfsummary'] = sum_eval
+        evals["hfsummary"] = sum_eval
     train = ConcatDataset(train_datasets)
-    collate_fn = DataCollatorForPairRank(tokenizer, max_length=training_conf['max_length'], drop_token_type= 'galactica' in model_name)
+    collate_fn = DataCollatorForPairRank(
+        tokenizer, max_length=training_conf["max_length"], drop_token_type="galactica" in model_name
+    )
     assert len(evals) > 0
     trainer = RankTrainer(
         model,
@@ -149,6 +181,6 @@ if __name__ == "__main__":
         eval_dataset=eval,
         data_collator=collate_fn,
         tokenizer=tokenizer,
-        compute_metrics=compute_metrics
+        compute_metrics=compute_metrics,
     )
     trainer.train()
diff --git a/model/reward/instructor/utils.py b/model/reward/instructor/utils.py
index d59bb13c..9441ddb9 100644
--- a/model/reward/instructor/utils.py
+++ b/model/reward/instructor/utils.py
@@ -1,96 +1,100 @@
+# -*- coding: utf-8 -*-
 import re
+
 import yaml
-from torch.utils.data import Subset
 from sklearn.model_selection import train_test_split
+from torch.utils.data import Subset
 from transformers import AutoTokenizer
 
-re_reference_remove = re.compile(r'\[([0-9])+\]|\[([0-9])+,([0-9])+\]')
+re_reference_remove = re.compile(r"\[([0-9])+\]|\[([0-9])+,([0-9])+\]")
+
 
 def webgpt_return_format(row):
-    if row['score_0'] >= row['score_1']:
+    if row["score_0"] >= row["score_1"]:
         # remove this to prevent information leak, since we are not using reference
         return {
-                'question': row['question']['full_text'],
-                     'pos': re_reference_remove.sub('', row['answer_0']),
-                     'neg': re_reference_remove.sub('', row['answer_1'])
-                }
+            "question": row["question"]["full_text"],
+            "pos": re_reference_remove.sub("", row["answer_0"]),
+            "neg": re_reference_remove.sub("", row["answer_1"]),
+        }
 
     return {
-            'question': row['question']['full_text'],
-                 'pos': re_reference_remove.sub('', row['answer_1']),
-                 'neg': re_reference_remove.sub('', row['answer_0'])
-            }
+        "question": row["question"]["full_text"],
+        "pos": re_reference_remove.sub("", row["answer_1"]),
+        "neg": re_reference_remove.sub("", row["answer_0"]),
+    }
 
 
 def get_tokenizer(tokenizer_name):
     tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
-    if 'galactica' in tokenizer_name:
-        tokenizer.add_special_tokens({'pad_token':'<pad>', 'eos_token': '</s>' })
+    if "galactica" in tokenizer_name:
+        tokenizer.add_special_tokens({"pad_token": "<pad>", "eos_token": "</s>"})
 
     return tokenizer
 
 
-
 def train_val_dataset(dataset, val_split=0.2):
-    train_idx, val_idx = train_test_split(list(range(len(dataset))), 
-        test_size=val_split, random_state=666, shuffle=True)
+    train_idx, val_idx = train_test_split(
+        list(range(len(dataset))), test_size=val_split, random_state=666, shuffle=True
+    )
     # [3879, 11479, 8341, 9177, 10798, 18177, 5735, 15669, 4837, 2760]
     print(val_idx[:10])
     # [13582, 5919, 11875, 7373, 19135, 13706, 8555, 15788, 15005, 15209]
     print(train_idx[:10])
     return Subset(dataset, train_idx), Subset(dataset, val_idx)
 
+
 def freeze_top_n_layers(model, target_layers):
     # its possible we can simply detect which module is a ModuleList
     # and simply freeze the module without doing string parsing
     for name, param in model.named_parameters():
-        if 'embed' in name:
+        if "embed" in name:
             param.requires_grad = False
-        elif '.layer' in name or '.h.' in name:
-            tokens = name.split('.')
+        elif ".layer" in name or ".h." in name:
+            tokens = name.split(".")
             idx = 0
             for token in tokens:
-                if 'layer' in token or token == 'h':
+                if "layer" in token or token == "h":
                     break
                 idx += 1
             if idx >= len(tokens):
                 continue
 
-            layer_ = int(tokens[idx+1])
+            layer_ = int(tokens[idx + 1])
             if layer_ < target_layers:
                 # print('freeze ', layer_, name)
                 param.requires_grad = False
     return model
 
+
 def argument_parsing(parser):
     default_params = {
-        'num_train_epochs': 4,
-        'learning_rate': 3e-5,
-        'eval_steps': 500,
-        'loss': 'rank',
-        'max_length': 440,
-        'per_device_eval_batch_size': 5,
-        'per_device_train_batch_size': 8,
-        'gradient_accumulation_steps': 8,
-        'gradient_checkpointing': False,
-        'datasets': ['webgpt']
+        "num_train_epochs": 4,
+        "learning_rate": 3e-5,
+        "eval_steps": 500,
+        "loss": "rank",
+        "max_length": 440,
+        "per_device_eval_batch_size": 5,
+        "per_device_train_batch_size": 8,
+        "gradient_accumulation_steps": 8,
+        "gradient_checkpointing": False,
+        "datasets": ["webgpt"],
     }
     args = parser.parse_args()
-    with open(args.config, 'r', encoding='utf-8') as f:
+    with open(args.config, "r", encoding="utf-8") as f:
         training_conf = yaml.safe_load(f.read())
 
-    params = { **default_params, **training_conf }
-    params['gradient_accumulation_steps'] = int(params['gradient_accumulation_steps'])
-    params['num_train_epochs'] = int(params['num_train_epochs'])
-    params['per_device_train_batch_size'] = int(params['per_device_train_batch_size'])
-    params['learning_rate'] = float(params['learning_rate'])
+    params = {**default_params, **training_conf}
+    params["gradient_accumulation_steps"] = int(params["gradient_accumulation_steps"])
+    params["num_train_epochs"] = int(params["num_train_epochs"])
+    params["per_device_train_batch_size"] = int(params["per_device_train_batch_size"])
+    params["learning_rate"] = float(params["learning_rate"])
     return params
 
 
-
 if __name__ == "__main__":
     from transformers import AutoModelForSequenceClassification
 
-    model = AutoModelForSequenceClassification.from_pretrained('bigscience/bloomz-560m')
+    model = AutoModelForSequenceClassification.from_pretrained("bigscience/bloomz-560m")
     freeze_top_n_layers(model, 10)
-    print(model.state_dict().keys())
\ No newline at end of file
+    print(model.state_dict().keys())

From 28e0b4f77020ea9cd5317bbc4094c9008083a99f Mon Sep 17 00:00:00 2001
From: theblackcat102 <theblackcat102@github.com>
Date: Sun, 1 Jan 2023 12:03:34 +0000
Subject: [PATCH 27/41] [fix] Revert deleted vscode

---
 .vscode/settings.json                    |   4 +
 model/reward/instructor/requirements.txt | 136 +----------------------
 2 files changed, 5 insertions(+), 135 deletions(-)
 create mode 100644 .vscode/settings.json

diff --git a/.vscode/settings.json b/.vscode/settings.json
new file mode 100644
index 00000000..4c58a32f
--- /dev/null
+++ b/.vscode/settings.json
@@ -0,0 +1,4 @@
+{
+  "python.formatting.provider": "autopep8",
+  "python.analysis.extraPaths": ["${workspaceFolder}/oasst-shared"]
+}
diff --git a/model/reward/instructor/requirements.txt b/model/reward/instructor/requirements.txt
index 9fc45917..cb1a9e68 100644
--- a/model/reward/instructor/requirements.txt
+++ b/model/reward/instructor/requirements.txt
@@ -1,140 +1,6 @@
-aiohttp==3.8.3
-aiosignal==1.3.1
-anyio==3.6.2
-argon2-cffi==21.3.0
-argon2-cffi-bindings==21.2.0
-arrow==1.2.3
-asttokens==2.2.1
-async-timeout==4.0.2
-attrs==22.2.0
-autopep8==2.0.1
-backcall==0.2.0
-beautifulsoup4==4.11.1
-bleach==5.0.1
-certifi==2022.12.7
-cffi==1.15.1
-charset-normalizer==2.1.1
-click==8.1.3
-comm==0.1.2
 datasets==2.8.0
-debugpy==1.6.4
-decorator==5.1.1
-defusedxml==0.7.1
-dill==0.3.6
-docker-pycreds==0.4.0
-entrypoints==0.4
 evaluate==0.4.0
-exceptiongroup==1.1.0
-executing==1.2.0
-fastjsonschema==2.16.2
-filelock==3.9.0
-fqdn==1.5.1
-frozenlist==1.3.3
-fsspec==2022.11.0
-gitdb==4.0.10
-GitPython==3.1.30
-huggingface-hub==0.11.1
-idna==3.4
-iniconfig==1.1.1
-ipykernel==6.19.4
-ipython==8.7.0
-ipython-genutils==0.2.0
-ipywidgets==8.0.4
-isoduration==20.11.0
-jedi==0.18.2
-Jinja2==3.1.2
-joblib==1.2.0
-jsonpointer==2.3
-jsonschema==4.17.3
-jupyter==1.0.0
-jupyter-console==6.4.4
-jupyter-events==0.5.0
-jupyter_client==7.4.8
-jupyter_core==5.1.1
-jupyter_server==2.0.6
-jupyter_server_terminals==0.4.3
-jupyterlab-pygments==0.2.2
-jupyterlab-widgets==3.0.5
-lightning-utilities==0.5.0
-MarkupSafe==2.1.1
-matplotlib-inline==0.1.6
-mistune==2.0.4
-multidict==6.0.4
-multiprocess==0.70.14
-nbclassic==0.4.8
-nbclient==0.7.2
-nbconvert==7.2.7
-nbformat==5.7.1
-nest-asyncio==1.5.6
-notebook==6.5.2
-notebook_shim==0.2.2
-numpy==1.24.1
-packaging==22.0
-pandas==1.5.2
-pandocfilters==1.5.0
-parso==0.8.3
-pathtools==0.1.2
-pexpect==4.8.0
-pickleshare==0.7.5
-platformdirs==2.6.2
-pluggy==1.0.0
-prometheus-client==0.15.0
-promise==2.3
-prompt-toolkit==3.0.36
-protobuf==3.20.1
-psutil==5.9.4
-ptyprocess==0.7.0
-pure-eval==0.2.2
-pyarrow==10.0.1
-pycodestyle==2.10.0
-pycparser==2.21
-Pygments==2.13.0
-pyrsistent==0.19.3
-pytest==7.2.0
-python-dateutil==2.8.2
-python-json-logger==2.0.4
-pytorch-lightning==1.8.6
-pytz==2022.7
-PyYAML==6.0
-pyzmq==24.0.1
-qtconsole==5.4.0
-QtPy==2.3.0
-regex==2022.10.31
-requests==2.28.1
-responses==0.18.0
-rfc3339-validator==0.1.4
-rfc3986-validator==0.1.1
 scikit-learn==1.2.0
-scipy==1.9.3
-Send2Trash==1.8.0
-sentry-sdk==1.12.1
-setproctitle==1.3.2
-shortuuid==1.0.11
-six==1.16.0
-smmap==5.0.0
-sniffio==1.3.0
-soupsieve==2.3.2.post1
-stack-data==0.6.2
-tensorboardX==2.5.1
-terminado==0.17.1
-threadpoolctl==3.1.0
-tinycss2==1.2.1
-tokenizers==0.13.2
-tomli==2.0.1
 torch==1.12.1+cu116
-torchmetrics==0.11.0
-tornado==6.2
-tqdm==4.64.1
-traitlets==5.8.0
 transformers==4.25.1
-typing_extensions==4.4.0
-uri-template==1.2.0
-urllib3==1.26.13
-wandb==0.13.7
-wcwidth==0.2.5
-webcolors==1.12
-webencodings==0.5.1
-websocket-client==1.4.2
-widgetsnbextension==4.0.5
-xxhash==3.2.0
-yarl==1.8.2
+wandb==0.13.7
\ No newline at end of file

From 8f0028bc44133af8bd54c301fa8546d56cadc2bf Mon Sep 17 00:00:00 2001
From: theblackcat102 <theblackcat102@github.com>
Date: Sun, 1 Jan 2023 13:28:48 +0000
Subject: [PATCH 28/41] [fix] Fix provider

---
 .vscode/settings.json                    | 2 +-
 model/reward/instructor/requirements.txt | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/.vscode/settings.json b/.vscode/settings.json
index 4c58a32f..56a51f78 100644
--- a/.vscode/settings.json
+++ b/.vscode/settings.json
@@ -1,4 +1,4 @@
 {
-  "python.formatting.provider": "autopep8",
+  "python.formatting.provider": "black",
   "python.analysis.extraPaths": ["${workspaceFolder}/oasst-shared"]
 }
diff --git a/model/reward/instructor/requirements.txt b/model/reward/instructor/requirements.txt
index cb1a9e68..e225a2ca 100644
--- a/model/reward/instructor/requirements.txt
+++ b/model/reward/instructor/requirements.txt
@@ -3,4 +3,4 @@ evaluate==0.4.0
 scikit-learn==1.2.0
 torch==1.12.1+cu116
 transformers==4.25.1
-wandb==0.13.7
\ No newline at end of file
+wandb==0.13.7

From aa2d426bbc33e51e89ed2dd147302303104d29f6 Mon Sep 17 00:00:00 2001
From: onegunsamurai <89980150+onegunsamurai@users.noreply.github.com>
Date: Sun, 1 Jan 2023 17:57:45 +0000
Subject: [PATCH 29/41] Add oasst-shared folder to discord-bot container (#241)

* Add oasst-shared folder to discord-bot container

- Container was missing oasst-shared folder, thus was failing, due to missing package error.
- Add bash code block to README.MD (Co-authored-by: Vladyslav <v.bukhantsov@dexilon.io>)
---
 discord-bot/README.md         | 10 +++++++++-
 docker/Dockerfile.discord-bot |  6 +++---
 2 files changed, 12 insertions(+), 4 deletions(-)

diff --git a/discord-bot/README.md b/discord-bot/README.md
index 1ff47c31..d78bcecc 100644
--- a/discord-bot/README.md
+++ b/discord-bot/README.md
@@ -12,9 +12,17 @@ If you are unfamiliar with `hikari`, `lightbulb`, or `miru`, please refer to the
 
 ### Setup
 
-To run the bot
+To run the bot:
+
+Install dependency module `oasst-shared`
 
 ```bash
+cd oasst-shared
+pip install -e .
+```
+
+```bash
+cd ../discord-bot
 cp .env.example .env
 
 python -V  # 3.10
diff --git a/docker/Dockerfile.discord-bot b/docker/Dockerfile.discord-bot
index 13ae308a..09e65fb8 100644
--- a/docker/Dockerfile.discord-bot
+++ b/docker/Dockerfile.discord-bot
@@ -1,7 +1,7 @@
 FROM python:3.10-slim-bullseye
 RUN mkdir /app
-COPY ./discord-bot/requirements.txt /requirements.txt
-RUN pip install -r requirements.txt
 WORKDIR /app
 COPY ./discord-bot /app
-CMD ["python", "bot.py"]
+COPY ./oasst-shared/oasst_shared /app/oasst_shared
+RUN pip install -r requirements.txt
+CMD ["python","-m","bot"]

From b5ebab883940fcdfb7564579491b807bbc669adc Mon Sep 17 00:00:00 2001
From: Jack Michaud <jack@lomz.me>
Date: Sat, 31 Dec 2022 16:42:12 -0500
Subject: [PATCH 30/41] chore: add print_openapi_schema and
 start-mock-server.sh script

---
 backend/print_openapi_schema.py               | 10 +++++++++
 .../backend-development/start-mock-server.sh  | 22 +++++++++++++++++++
 2 files changed, 32 insertions(+)
 create mode 100644 backend/print_openapi_schema.py
 create mode 100755 scripts/backend-development/start-mock-server.sh

diff --git a/backend/print_openapi_schema.py b/backend/print_openapi_schema.py
new file mode 100644
index 00000000..23d5f6b2
--- /dev/null
+++ b/backend/print_openapi_schema.py
@@ -0,0 +1,10 @@
+import json
+from main import app
+
+
+def get_openapi_schema():
+    return json.dumps(app.openapi())
+
+
+if __name__ == "__main__":
+    print(get_openapi_schema())
diff --git a/scripts/backend-development/start-mock-server.sh b/scripts/backend-development/start-mock-server.sh
new file mode 100755
index 00000000..7a11e097
--- /dev/null
+++ b/scripts/backend-development/start-mock-server.sh
@@ -0,0 +1,22 @@
+#!/usr/bin/env bash
+parent_path=$( cd "$(dirname "${BASH_SOURCE[0]}")" ; pwd -P )
+
+# switch to backend directory
+pushd "$parent_path/../../backend"
+
+export DEBUG_SKIP_API_KEY_CHECK=True
+
+python -m print_openapi_schema > oasst-openapi.json
+
+MOCK_SERVER_PORT=8080
+
+docker run -d -it --rm \
+  -p $MOCK_SERVER_PORT:8080 \
+  --name wiremock \
+  wiremock/wiremock:2.35.0
+
+sleep 1
+
+curl -X POST -d @oasst-openapi.json http://localhost:$MOCK_SERVER_PORT/__admin/mappings/import
+
+popd

From 5c9b9e3c9f590bc32868908a483795d78452af14 Mon Sep 17 00:00:00 2001
From: Jack Michaud <jack@lomz.me>
Date: Sat, 31 Dec 2022 18:10:21 -0500
Subject: [PATCH 31/41] fix: use prism instead of wiremock

the previous iteration of this script which used wiremock threw a null pointer exception.
---
 scripts/backend-development/start-mock-server.sh | 9 ++-------
 1 file changed, 2 insertions(+), 7 deletions(-)

diff --git a/scripts/backend-development/start-mock-server.sh b/scripts/backend-development/start-mock-server.sh
index 7a11e097..a9bd0fbc 100755
--- a/scripts/backend-development/start-mock-server.sh
+++ b/scripts/backend-development/start-mock-server.sh
@@ -10,13 +10,8 @@ python -m print_openapi_schema > oasst-openapi.json
 
 MOCK_SERVER_PORT=8080
 
-docker run -d -it --rm \
-  -p $MOCK_SERVER_PORT:8080 \
-  --name wiremock \
-  wiremock/wiremock:2.35.0
+docker run --init --rm -d -p $MOCK_SERVER_PORT:4010 -v $(pwd):/tmp -P stoplight/prism:4 proxy -h 0.0.0.0 "/tmp/oasst-openapi.json"
 
-sleep 1
-
-curl -X POST -d @oasst-openapi.json http://localhost:$MOCK_SERVER_PORT/__admin/mappings/import
 
 popd
+

From 036c4b6c6c3944eae1d507f8d169d608df3df24d Mon Sep 17 00:00:00 2001
From: Jack Michaud <jack@lomz.me>
Date: Sat, 31 Dec 2022 18:56:13 -0500
Subject: [PATCH 32/41] test: add contract test for discord bot

---
 discord-bot/requirements.dev.txt                 |  2 ++
 discord-bot/tests/test_api_contract.py           | 15 +++++++++++++++
 scripts/backend-development/start-mock-server.sh |  6 +++++-
 scripts/discord-bot-development/test.sh          | 10 ++++++++++
 4 files changed, 32 insertions(+), 1 deletion(-)
 create mode 100644 discord-bot/requirements.dev.txt
 create mode 100644 discord-bot/tests/test_api_contract.py
 create mode 100755 scripts/discord-bot-development/test.sh

diff --git a/discord-bot/requirements.dev.txt b/discord-bot/requirements.dev.txt
new file mode 100644
index 00000000..ee4ba018
--- /dev/null
+++ b/discord-bot/requirements.dev.txt
@@ -0,0 +1,2 @@
+pytest
+pytest-asyncio
diff --git a/discord-bot/tests/test_api_contract.py b/discord-bot/tests/test_api_contract.py
new file mode 100644
index 00000000..fe7be6cf
--- /dev/null
+++ b/discord-bot/tests/test_api_contract.py
@@ -0,0 +1,15 @@
+import pytest
+from bot.api_client import OasstApiClient
+
+from oasst_shared.schemas import protocol as protocol_schema
+
+
+@pytest.fixture
+def oasst_api_client_mocked():
+    client = OasstApiClient(backend_url="http://localhost:8080", api_key="123")
+    yield client
+
+
+@pytest.mark.asyncio
+async def test_fetch_task(oasst_api_client_mocked: OasstApiClient):
+    assert await oasst_api_client_mocked.fetch_task(task_type=protocol_schema.TaskRequestType.random) is not None
diff --git a/scripts/backend-development/start-mock-server.sh b/scripts/backend-development/start-mock-server.sh
index a9bd0fbc..807230e8 100755
--- a/scripts/backend-development/start-mock-server.sh
+++ b/scripts/backend-development/start-mock-server.sh
@@ -10,7 +10,11 @@ python -m print_openapi_schema > oasst-openapi.json
 
 MOCK_SERVER_PORT=8080
 
-docker run --init --rm -d -p $MOCK_SERVER_PORT:4010 -v $(pwd):/tmp -P stoplight/prism:4 proxy -h 0.0.0.0 "/tmp/oasst-openapi.json"
+docker run --init --rm -d \
+  -p $MOCK_SERVER_PORT:4010 \
+  -v $(pwd):/tmp \
+  -P stoplight/prism:4 \
+  mock -h 0.0.0.0 "/tmp/oasst-openapi.json"
 
 
 popd
diff --git a/scripts/discord-bot-development/test.sh b/scripts/discord-bot-development/test.sh
new file mode 100755
index 00000000..f863f74e
--- /dev/null
+++ b/scripts/discord-bot-development/test.sh
@@ -0,0 +1,10 @@
+#!/usr/bin/env bash
+parent_path=$( cd "$(dirname "${BASH_SOURCE[0]}")" ; pwd -P )
+
+# switch to backend directory
+pushd "$parent_path/../../discord-bot"
+
+pytest .
+
+popd
+

From 904f442f5a6f394c8ec0ed5ddf2a90bb96475ee4 Mon Sep 17 00:00:00 2001
From: Jack Michaud <jack@lomz.me>
Date: Sat, 31 Dec 2022 18:56:23 -0500
Subject: [PATCH 33/41] fix: indicate that user is optional

---
 oasst-shared/oasst_shared/schemas/protocol.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/oasst-shared/oasst_shared/schemas/protocol.py b/oasst-shared/oasst_shared/schemas/protocol.py
index 8a6685c2..5f05adc3 100644
--- a/oasst-shared/oasst_shared/schemas/protocol.py
+++ b/oasst-shared/oasst_shared/schemas/protocol.py
@@ -5,7 +5,7 @@ from typing import Literal, Optional, Union
 from uuid import UUID, uuid4
 
 import pydantic
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 
 
 class TaskRequestType(str, enum.Enum):
@@ -56,7 +56,9 @@ class TaskRequest(BaseModel):
     """The frontend asks the backend for a task."""
 
     type: TaskRequestType = TaskRequestType.random
-    user: Optional[User] = None
+    # Must use Field(..., nullable=True) to indicate to the OpenAPI schema that
+    # this is optional. https://github.com/pydantic/pydantic/issues/1270
+    user: Optional[User] = Field(None, nullable=True)
     collective: bool = False
 
 

From 953d243a6fe099de8e4203f12bace22cae565fe2 Mon Sep 17 00:00:00 2001
From: Jack Michaud <jack@lomz.me>
Date: Sun, 1 Jan 2023 09:36:36 -0500
Subject: [PATCH 34/41] chore: start-mock-server waits for health

---
 .../backend-development/start-mock-server.sh  | 39 ++++++++++++++-----
 1 file changed, 30 insertions(+), 9 deletions(-)

diff --git a/scripts/backend-development/start-mock-server.sh b/scripts/backend-development/start-mock-server.sh
index 807230e8..4bd3b09c 100755
--- a/scripts/backend-development/start-mock-server.sh
+++ b/scripts/backend-development/start-mock-server.sh
@@ -4,17 +4,38 @@ parent_path=$( cd "$(dirname "${BASH_SOURCE[0]}")" ; pwd -P )
 # switch to backend directory
 pushd "$parent_path/../../backend"
 
-export DEBUG_SKIP_API_KEY_CHECK=True
-
-python -m print_openapi_schema > oasst-openapi.json
-
 MOCK_SERVER_PORT=8080
+OPENAPI_JSON_FILE_NAME=openapi.json
 
-docker run --init --rm -d \
-  -p $MOCK_SERVER_PORT:4010 \
-  -v $(pwd):/tmp \
-  -P stoplight/prism:4 \
-  mock -h 0.0.0.0 "/tmp/oasst-openapi.json"
+echo "Generating OpenAPI schema..."
+python -m print_openapi_schema > $OPENAPI_JSON_FILE_NAME
+echo "Done!"
+
+# If oasst-mock-backend docker container is already running,
+# just restart it
+if [ "$(docker ps -q -f name=oasst-mock-backend)" ]; then
+    echo "oasst-mock-backend container exists, restarting..."
+    docker restart oasst-mock-backend
+else
+    echo "Creating new oasst-mock-backend container..."
+    docker run --init --rm -d \
+      --name oasst-mock-backend \
+      -p $MOCK_SERVER_PORT:4010 \
+      -v $(pwd):/tmp \
+      -P stoplight/prism:4 \
+      mock -h 0.0.0.0 "/tmp/$OPENAPI_JSON_FILE_NAME"
+fi
+
+echo "Waiting for server to be live..."
+curl --retry-all-errors --retry 5 localhost:$MOCK_SERVER_PORT
+echo ""
+
+# if return code is successful, print successful response
+if [ $? -eq 0 ]; then
+    echo "Mock server is running at localhost:$MOCK_SERVER_PORT"
+else
+    echo "Mock server failed to start"
+fi
 
 
 popd

From b85d4feb58c4d4d5e1e908a35cd8f74be4ab1c2c Mon Sep 17 00:00:00 2001
From: Jack Michaud <jack@lomz.me>
Date: Sun, 1 Jan 2023 09:36:57 -0500
Subject: [PATCH 35/41] chore: ignore generated openapi schema

---
 .gitignore | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/.gitignore b/.gitignore
index 84512e5a..9cdabc03 100644
--- a/.gitignore
+++ b/.gitignore
@@ -5,3 +5,6 @@
 *.egg-info
 __pycache__
 .DS_Store
+
+# Generated files
+backend/oasst-openapi.json

From 5870148f0a7bffdf7728e24076746855f13a8386 Mon Sep 17 00:00:00 2001
From: Jack Michaud <jack@lomz.me>
Date: Sun, 1 Jan 2023 10:04:55 -0500
Subject: [PATCH 36/41] test: add complete tests for oasst_api_client

---
 discord-bot/tests/test_api_contract.py     | 15 -------
 discord-bot/tests/test_oasst_api_client.py | 51 ++++++++++++++++++++++
 2 files changed, 51 insertions(+), 15 deletions(-)
 delete mode 100644 discord-bot/tests/test_api_contract.py
 create mode 100644 discord-bot/tests/test_oasst_api_client.py

diff --git a/discord-bot/tests/test_api_contract.py b/discord-bot/tests/test_api_contract.py
deleted file mode 100644
index fe7be6cf..00000000
--- a/discord-bot/tests/test_api_contract.py
+++ /dev/null
@@ -1,15 +0,0 @@
-import pytest
-from bot.api_client import OasstApiClient
-
-from oasst_shared.schemas import protocol as protocol_schema
-
-
-@pytest.fixture
-def oasst_api_client_mocked():
-    client = OasstApiClient(backend_url="http://localhost:8080", api_key="123")
-    yield client
-
-
-@pytest.mark.asyncio
-async def test_fetch_task(oasst_api_client_mocked: OasstApiClient):
-    assert await oasst_api_client_mocked.fetch_task(task_type=protocol_schema.TaskRequestType.random) is not None
diff --git a/discord-bot/tests/test_oasst_api_client.py b/discord-bot/tests/test_oasst_api_client.py
new file mode 100644
index 00000000..50ea6d5b
--- /dev/null
+++ b/discord-bot/tests/test_oasst_api_client.py
@@ -0,0 +1,51 @@
+from uuid import uuid4
+import pytest
+from bot.api_client import OasstApiClient
+
+from oasst_shared.schemas import protocol as protocol_schema
+
+
+@pytest.fixture
+def oasst_api_client_mocked():
+    client = OasstApiClient(backend_url="http://localhost:8080", api_key="123")
+    yield client
+    # TODO The fixture should close this connection, but there seems to be a bug
+    # with async fixtures and pytest.
+    # Since this only results in a warning, I'm leaving this for now.
+    # await client.close()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("task_type", protocol_schema.TaskRequestType)
+async def test_can_fetch_task(task_type: protocol_schema.TaskRequestType, oasst_api_client_mocked: OasstApiClient):
+    assert await oasst_api_client_mocked.fetch_task(task_type=task_type) is not None
+
+
+@pytest.mark.asyncio
+async def test_can_ack_task(oasst_api_client_mocked: OasstApiClient):
+    await oasst_api_client_mocked.ack_task(task_id=uuid4(), message_id="123")
+
+
+@pytest.mark.asyncio
+async def test_can_nack_task(oasst_api_client_mocked: OasstApiClient):
+    await oasst_api_client_mocked.nack_task(task_id=uuid4(), reason="bad task")
+
+
+@pytest.mark.asyncio
+async def test_can_post_interaction(oasst_api_client_mocked: OasstApiClient):
+    assert (
+        await oasst_api_client_mocked.post_interaction(
+            protocol_schema.TextReplyToMessage(
+                type="text_reply_to_message",
+                message_id="123",
+                user_message_id="321",
+                text="This is my reply",
+                user=protocol_schema.User(
+                    id="123",
+                    display_name="lomz",
+                    auth_method="discord",
+                ),
+            )
+        )
+        is not None
+    )

From 898113ef8d1b15a0d94a89ec6f66860665e8bc0d Mon Sep 17 00:00:00 2001
From: Jack Michaud <jack@lomz.me>
Date: Sun, 1 Jan 2023 10:05:17 -0500
Subject: [PATCH 37/41] fix: correct typing for OasstApiClient

---
 discord-bot/bot/api_client.py | 15 ++++++++-------
 1 file changed, 8 insertions(+), 7 deletions(-)

diff --git a/discord-bot/bot/api_client.py b/discord-bot/bot/api_client.py
index 54b489b4..f97ab840 100644
--- a/discord-bot/bot/api_client.py
+++ b/discord-bot/bot/api_client.py
@@ -53,14 +53,16 @@ class OasstApiClient:
             TaskType.done: protocol_schema.TaskDone,
         }
 
-    async def post(self, path: str, data: dict[str, t.Any]) -> dict[str, t.Any]:
+    async def post(self, path: str, data: dict[str, t.Any]) -> Optional[dict[str, t.Any]]:
         """Make a POST request to the backend."""
         logger.debug(f"POST {self.backend_url}{path} DATA: {data}")
         response = await self.session.post(f"{self.backend_url}{path}", json=data, headers={"X-API-Key": self.api_key})
         response.raise_for_status()
         return await response.json()
 
-    def _parse_task(self, data: dict[str, t.Any]) -> protocol_schema.Task:
+    def _parse_task(self, data: Optional[dict[str, t.Any]]) -> protocol_schema.Task:
+        if data is None:
+            raise Exception("Cannot parse data as a task: data is none")
         task_type = TaskType(data.get("type"))
 
         model = self.task_models_map.get(task_type)
@@ -89,23 +91,22 @@ class OasstApiClient:
         logger.debug(f"Fetching random for user {user}")
         return await self.fetch_task(protocol_schema.TaskRequestType.random, user, collective)
 
-    async def ack_task(self, task_id: str | UUID, message_id: str):
+    async def ack_task(self, task_id: str | UUID, message_id: str) -> None:
         """Send an ACK for a task to the backend."""
         logger.debug(f"ACK task {task_id} with post {message_id}")
         req = protocol_schema.TaskAck(message_id=message_id)
-        return await self.post(f"/api/v1/tasks/{task_id}/ack", data=req.dict())
+        await self.post(f"/api/v1/tasks/{task_id}/ack", data=req.dict())
 
-    async def nack_task(self, task_id: str | UUID, reason: str):
+    async def nack_task(self, task_id: str | UUID, reason: str) -> None:
         """Send a NACK for a task to the backend."""
         logger.debug(f"NACK task {task_id} with reason {reason}")
         req = protocol_schema.TaskNAck(reason=reason)
-        return await self.post(f"/api/v1/tasks/{task_id}/nack", data=req.dict())
+        await self.post(f"/api/v1/tasks/{task_id}/nack", data=req.dict())
 
     async def post_interaction(self, interaction: protocol_schema.Interaction) -> protocol_schema.Task:
         """Send a completed task to the backend."""
         logger.debug(f"Interaction: {interaction}")
         resp = await self.post("/api/v1/tasks/interaction", data=interaction.dict())
-
         return self._parse_task(resp)
 
     async def close(self):

From e15e208b201e623aafb6419bb08ae88fa14b499b Mon Sep 17 00:00:00 2001
From: Jack Michaud <jack@lomz.me>
Date: Sun, 1 Jan 2023 10:05:34 -0500
Subject: [PATCH 38/41] fix: add response_model to tasks_interaction

This is required for output validation and for the mock server.
---
 backend/oasst_backend/api/v1/tasks.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/backend/oasst_backend/api/v1/tasks.py b/backend/oasst_backend/api/v1/tasks.py
index a271d5f0..570834ad 100644
--- a/backend/oasst_backend/api/v1/tasks.py
+++ b/backend/oasst_backend/api/v1/tasks.py
@@ -204,7 +204,7 @@ def tasks_acknowledge_failure(
         raise OasstError("Failed to not acknowledge task.", OasstErrorCode.TASK_NACK_FAILED)
 
 
-@router.post("/interaction")
+@router.post("/interaction", response_model=protocol_schema.TaskDone)
 def tasks_interaction(
     *,
     db: Session = Depends(deps.get_db),

From 604ad2f8e77dc331839ca2ce6c46ed48eb78f4ca Mon Sep 17 00:00:00 2001
From: Jack Michaud <jack@lomz.me>
Date: Sun, 1 Jan 2023 10:21:52 -0500
Subject: [PATCH 39/41] style: run pre-commit hooks

---
 backend/print_openapi_schema.py                  | 2 ++
 discord-bot/tests/test_oasst_api_client.py       | 3 ++-
 scripts/backend-development/start-mock-server.sh | 1 -
 scripts/discord-bot-development/test.sh          | 1 -
 4 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/backend/print_openapi_schema.py b/backend/print_openapi_schema.py
index 23d5f6b2..e724cc7e 100644
--- a/backend/print_openapi_schema.py
+++ b/backend/print_openapi_schema.py
@@ -1,4 +1,6 @@
+# -*- coding: utf-8 -*-
 import json
+
 from main import app
 
 
diff --git a/discord-bot/tests/test_oasst_api_client.py b/discord-bot/tests/test_oasst_api_client.py
index 50ea6d5b..c5cafe99 100644
--- a/discord-bot/tests/test_oasst_api_client.py
+++ b/discord-bot/tests/test_oasst_api_client.py
@@ -1,7 +1,8 @@
+# -*- coding: utf-8 -*-
 from uuid import uuid4
+
 import pytest
 from bot.api_client import OasstApiClient
-
 from oasst_shared.schemas import protocol as protocol_schema
 
 
diff --git a/scripts/backend-development/start-mock-server.sh b/scripts/backend-development/start-mock-server.sh
index 4bd3b09c..e27f532e 100755
--- a/scripts/backend-development/start-mock-server.sh
+++ b/scripts/backend-development/start-mock-server.sh
@@ -39,4 +39,3 @@ fi
 
 
 popd
-
diff --git a/scripts/discord-bot-development/test.sh b/scripts/discord-bot-development/test.sh
index f863f74e..a45adf00 100755
--- a/scripts/discord-bot-development/test.sh
+++ b/scripts/discord-bot-development/test.sh
@@ -7,4 +7,3 @@ pushd "$parent_path/../../discord-bot"
 pytest .
 
 popd
-

From ea40a22049edfc9bcc1c20f1342dca46252deda8 Mon Sep 17 00:00:00 2001
From: Jack Michaud <jack@lomz.me>
Date: Sun, 1 Jan 2023 12:27:20 -0500
Subject: [PATCH 40/41] fix: nack and ack should return empty responses

---
 backend/oasst_backend/api/v1/tasks.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/backend/oasst_backend/api/v1/tasks.py b/backend/oasst_backend/api/v1/tasks.py
index 570834ad..636f0feb 100644
--- a/backend/oasst_backend/api/v1/tasks.py
+++ b/backend/oasst_backend/api/v1/tasks.py
@@ -153,7 +153,7 @@ def request_task(
     return task
 
 
-@router.post("/{task_id}/ack")
+@router.post("/{task_id}/ack", response_model=None)
 def tasks_acknowledge(
     *,
     db: Session = Depends(deps.get_db),
@@ -179,10 +179,9 @@ def tasks_acknowledge(
     except Exception:
         logger.exception("Failed to acknowledge task.")
         raise OasstError("Failed to acknowledge task.", OasstErrorCode.TASK_ACK_FAILED)
-    return {}
 
 
-@router.post("/{task_id}/nack")
+@router.post("/{task_id}/nack", response_model=None)
 def tasks_acknowledge_failure(
     *,
     db: Session = Depends(deps.get_db),

From 29b08875d4d7f9f469167faa48bab4af700ce8c3 Mon Sep 17 00:00:00 2001
From: Jack Michaud <jack@lomz.me>
Date: Sun, 1 Jan 2023 12:34:03 -0500
Subject: [PATCH 41/41] refactor: move print_openapi_schema into main

---
 backend/main.py                               | 30 +++++++++++++++++++
 backend/print_openapi_schema.py               | 12 --------
 .../backend-development/start-mock-server.sh  |  2 +-
 3 files changed, 31 insertions(+), 13 deletions(-)
 delete mode 100644 backend/print_openapi_schema.py

diff --git a/backend/main.py b/backend/main.py
index 387d4e51..9cf43701 100644
--- a/backend/main.py
+++ b/backend/main.py
@@ -179,3 +179,33 @@ if settings.DEBUG_USE_SEED_DATA:
 
 
 app.include_router(api_router, prefix=settings.API_V1_STR)
+
+
+def get_openapi_schema():
+    return json.dumps(app.openapi())
+
+
+if __name__ == "__main__":
+    # Importing here so we don't import packages unnecessarily if we're
+    # importing main as a module.
+    import argparse
+    import json
+
+    import uvicorn
+
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument(
+        "--print-openapi-schema",
+        help="Dumps the openapi schema to stdout",
+        action=argparse.BooleanOptionalAction,
+    )
+    parser.add_argument("--host", help="The host to run the server")
+    parser.add_argument("--port", help="The port to run the server")
+
+    args = parser.parse_args()
+
+    if args.print_openapi_schema:
+        print(get_openapi_schema())
+    else:
+        uvicorn.run(app, host=args.host, port=args.port)
diff --git a/backend/print_openapi_schema.py b/backend/print_openapi_schema.py
deleted file mode 100644
index e724cc7e..00000000
--- a/backend/print_openapi_schema.py
+++ /dev/null
@@ -1,12 +0,0 @@
-# -*- coding: utf-8 -*-
-import json
-
-from main import app
-
-
-def get_openapi_schema():
-    return json.dumps(app.openapi())
-
-
-if __name__ == "__main__":
-    print(get_openapi_schema())
diff --git a/scripts/backend-development/start-mock-server.sh b/scripts/backend-development/start-mock-server.sh
index e27f532e..35a202a6 100755
--- a/scripts/backend-development/start-mock-server.sh
+++ b/scripts/backend-development/start-mock-server.sh
@@ -8,7 +8,7 @@ MOCK_SERVER_PORT=8080
 OPENAPI_JSON_FILE_NAME=openapi.json
 
 echo "Generating OpenAPI schema..."
-python -m print_openapi_schema > $OPENAPI_JSON_FILE_NAME
+python -m main --print-openapi-schema > $OPENAPI_JSON_FILE_NAME
 echo "Done!"
 
 # If oasst-mock-backend docker container is already running,