[feature] add rank dataset for webgpt and human feedback summary

model/reward/instructor/README.md

@@ -0,0 +1,7 @@
+
+
+
+```bash
+
+
+```

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
+
+
+

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

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

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 ]
+
+

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()

model/reward/instructor/trainer.py

@@ -0,0 +1,2 @@
+import wandb
+from accelerate import Accelerator

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'])
+            }

model/utils.py

@@ -0,0 +1,4 @@
+from transformers import AutoTokenizer
+
+
+def update_galactica_tokenizer():