tdpeuter/2025ML-project-neural_compression
tdpeuter/2025ML-project-neural_compression
Archived
1
Fork 0
Code Issues Pull requests Releases 1 Activity Actions

feat: Choose dataset with options

Browse source
This commit is contained in:
Tibo De Peuter 2025-11-30 20:19:39 +01:00
parent 20bdd4f566
commit 81c767371e
Signed by: tdpeuter
GPG key ID: 38297DE43F75FFE2
5 changed files with 67 additions and 60 deletions
Download patch file Download diff file Expand all files Collapse all files

23
CNN-model/dataset_loaders/EnWik9.py
View file

@ -1,18 +1,20 @@
from datasets import load_dataset
from torch.utils.data import Dataset
import torch
from os.path import curdir, join
from typing import Callable
import torch
from datasets import load_dataset
from .Dataset import Dataset
class EnWik9DataSet(Dataset):
def __init__(self, root: str = "data", transform: Callable | None = None):
super().__init__()
self.transform = transform
"""
Hugging Face: https://huggingface.co/datasets/haukur/enwik9
"""
def __init__(self, root: str | None = None, transform: Callable | None = None):
super().__init__('enwik9', root, transform)
# HuggingFace dataset: string text
path = join(curdir, root)
data = load_dataset("haukur/enwik9", cache_dir=path, split="train")
data = load_dataset("haukur/enwik9", cache_dir=self.root, split="train")
# Extract raw text
text = data["text"]
@ -31,7 +33,7 @@ class EnWik9DataSet(Dataset):
def __getitem__(self, idx):
# context window
x = self.data[idx : idx + self.context_length]
x = self.data[idx: idx + self.context_length]
# next byte target
y = self.data[idx + self.context_length]
@ -40,4 +42,3 @@ class EnWik9DataSet(Dataset):
x = self.transform(x)
return x, y