mirror of
https://github.com/wassname/Open-Assistant.git
synced 2026-08-10 11:50:44 +08:00
67 lines
2.5 KiB
Python
67 lines
2.5 KiB
Python
'''
|
|
|
|
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")
|
|
self.dataset = []
|
|
self.dataset_index = []
|
|
for idx, row in enumerate(dataset['train']):
|
|
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
|
|
# 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
|
|
|
|
|
|
|