44 lines
1.3 KiB
Python
44 lines
1.3 KiB
Python
from typing import Callable
|
|
|
|
import torch
|
|
from datasets import load_dataset
|
|
from torch import Tensor
|
|
|
|
from .Dataset import Dataset
|
|
|
|
|
|
class OpenGenomeDataset(Dataset):
|
|
"""
|
|
Hugging Face: https://huggingface.co/datasets/LongSafari/open-genome
|
|
|
|
:param split Either 'train', 'test' or 'validation'
|
|
:param stage Either 'sample', 'stage1' or 'stage2'.
|
|
'sample' only provides a 'validation' split
|
|
"""
|
|
|
|
def __init__(self,
|
|
root: str | None = None,
|
|
split: str = 'train',
|
|
transform: Callable = None,
|
|
stage: str = 'stage2'):
|
|
super().__init__('open_genome', root, transform)
|
|
|
|
data = load_dataset("LongSafari/open-genome", stage)
|
|
self.__train = ''.join(data[split]['text']).encode('utf-8', errors='replace')
|
|
|
|
self.data: Tensor = torch.tensor(bytearray(self.__train), dtype=torch.long)
|
|
|
|
# Model uses fixed 128-length context
|
|
self.context_length = 128
|
|
|
|
def __len__(self):
|
|
return len(self.data) - self.context_length
|
|
|
|
def __getitem__(self, item):
|
|
x = self.data[item: item + self.context_length]
|
|
y = self.data[item + self.context_length]
|
|
|
|
if self.transform:
|
|
x = self.transform(x)
|
|
|
|
return x, y
|