File size: 939 Bytes
b6447fa | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 | import torch
from torch.utils.data import Dataset
from typing import Callable
class CharRNNDataset(Dataset):
"""Character-level dataset for training RNN language models."""
def __init__(self, text: str, seq_len: int, text_encoder_fn: Callable[[str], torch.Tensor]):
"""
Args:
text (str): Raw text corpus.
seq_len (int): Length of each training sequence.
text_encoder_fn (Callable): Function to encode text into a tensor of IDs.
"""
self.full_sequence = text_encoder_fn(text)
self.seq_len = seq_len
def __len__(self) -> int:
return len(self.full_sequence) - self.seq_len
def __getitem__(self, idx: int):
if idx < 0 or idx >= len(self):
raise IndexError("Index out of range")
end = idx + self.seq_len
x = self.full_sequence[idx:end]
y = self.full_sequence[idx + 1:end + 1]
return x, y
|