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