mirror of
https://github.com/openai/gpt-2-output-dataset
synced 2025-08-22 18:07:53 +00:00
87 lines
3.5 KiB
Python
87 lines
3.5 KiB
Python
|
import json
|
||
|
|
import numpy as np
|
||
|
|
from typing import List
|
||
|
|
|
||
|
|
import torch
|
||
|
|
from torch.utils.data import Dataset
|
||
|
|
from tqdm import tqdm
|
||
|
|
from transformers import PreTrainedTokenizer
|
||
|
|
|
||
|
|
from .download import download
|
||
|
|
|
||
|
|
|
||
|
|
def load_texts(data_file, expected_size=None):
|
||
|
|
texts = []
|
||
|
|
|
||
|
|
for line in tqdm(open(data_file), total=expected_size, desc=f'Loading {data_file}'):
|
||
|
|
texts.append(json.loads(line)['text'])
|
||
|
|
|
||
|
|
return texts
|
||
|
|
|
||
|
|
|
||
|
|
class Corpus:
|
||
|
|
def __init__(self, name, data_dir='data', skip_train=False):
|
||
|
|
download(name, data_dir=data_dir)
|
||
|
|
self.name = name
|
||
|
|
self.train = load_texts(f'{data_dir}/{name}.train.jsonl', expected_size=250000) if not skip_train else None
|
||
|
|
self.test = load_texts(f'{data_dir}/{name}.test.jsonl', expected_size=5000)
|
||
|
|
self.valid = load_texts(f'{data_dir}/{name}.valid.jsonl', expected_size=5000)
|
||
|
|
|
||
|
|
|
||
|
|
class EncodedDataset(Dataset):
|
||
|
|
def __init__(self, real_texts: List[str], fake_texts: List[str], tokenizer: PreTrainedTokenizer,
|
||
|
|
max_sequence_length: int = None, min_sequence_length: int = None, epoch_size: int = None,
|
||
|
|
token_dropout: float = None, seed: int = None):
|
||
|
|
self.real_texts = real_texts
|
||
|
|
self.fake_texts = fake_texts
|
||
|
|
self.tokenizer = tokenizer
|
||
|
|
self.max_sequence_length = max_sequence_length
|
||
|
|
self.min_sequence_length = min_sequence_length
|
||
|
|
self.epoch_size = epoch_size
|
||
|
|
self.token_dropout = token_dropout
|
||
|
|
self.random = np.random.RandomState(seed)
|
||
|
|
|
||
|
|
def __len__(self):
|
||
|
|
return self.epoch_size or len(self.real_texts) + len(self.fake_texts)
|
||
|
|
|
||
|
|
def __getitem__(self, index):
|
||
|
|
if self.epoch_size is not None:
|
||
|
|
label = self.random.randint(2)
|
||
|
|
texts = [self.fake_texts, self.real_texts][label]
|
||
|
|
text = texts[self.random.randint(len(texts))]
|
||
|
|
else:
|
||
|
|
if index < len(self.real_texts):
|
||
|
|
text = self.real_texts[index]
|
||
|
|
label = 1
|
||
|
|
else:
|
||
|
|
text = self.fake_texts[index - len(self.real_texts)]
|
||
|
|
label = 0
|
||
|
|
|
||
|
|
tokens = self.tokenizer.encode(text)
|
||
|
|
|
||
|
|
if self.max_sequence_length is None:
|
||
|
|
tokens = tokens[:self.tokenizer.max_len - 2]
|
||
|
|
else:
|
||
|
|
output_length = min(len(tokens), self.max_sequence_length)
|
||
|
|
if self.min_sequence_length:
|
||
|
|
output_length = self.random.randint(min(self.min_sequence_length, len(tokens)), output_length + 1)
|
||
|
|
start_index = 0 if len(tokens) <= output_length else self.random.randint(0, len(tokens) - output_length + 1)
|
||
|
|
end_index = start_index + output_length
|
||
|
|
tokens = tokens[start_index:end_index]
|
||
|
|
|
||
|
|
if self.token_dropout:
|
||
|
|
dropout_mask = self.random.binomial(1, self.token_dropout, len(tokens)).astype(np.bool)
|
||
|
|
tokens = np.array(tokens)
|
||
|
|
tokens[dropout_mask] = self.tokenizer.unk_token_id
|
||
|
|
tokens = tokens.tolist()
|
||
|
|
|
||
|
|
if self.max_sequence_length is None or len(tokens) == self.max_sequence_length:
|
||
|
|
mask = torch.ones(len(tokens) + 2)
|
||
|
|
return torch.tensor([self.tokenizer.bos_token_id] + tokens + [self.tokenizer.eos_token_id]), mask, label
|
||
|
|
|
||
|
|
padding = [self.tokenizer.pad_token_id] * (self.max_sequence_length - len(tokens))
|
||
|
|
tokens = torch.tensor([self.tokenizer.bos_token_id] + tokens + [self.tokenizer.eos_token_id] + padding)
|
||
|
|
mask = torch.ones(tokens.shape[0])
|
||
|
|
mask[-len(padding):] = 0
|
||
|
|
return tokens, mask, label
|