| import torch |
| from torch.utils.data import Dataset, DataLoader |
|
|
| from src.data.tokenizer import BaseTokenizer |
|
|
|
|
| class GPTDataset(Dataset): |
| """Sliding-window dataset for GPT-style next-token prediction. |
| |
| Args: |
| text: raw input text |
| tokenizer: a BaseTokenizer instance |
| max_len: number of tokens the model sees as context |
| stride: step size between windows |
| """ |
|
|
| def __init__( |
| self, text: str, tokenizer: BaseTokenizer, max_len: int, stride: int |
| ) -> None: |
| self.max_len = max_len |
| self.stride = stride |
| self.tokens = torch.tensor(tokenizer.encode(text), dtype=torch.long) |
| self.num_samples = max(0, (len(self.tokens) - max_len - 1) // stride + 1) |
|
|
| def __len__(self) -> int: |
| return self.num_samples |
|
|
| def __getitem__(self, idx: int) -> tuple[torch.Tensor, torch.Tensor]: |
| start = idx * self.stride |
| return ( |
| self.tokens[start : start + self.max_len], |
| self.tokens[start + 1 : start + self.max_len + 1], |
| ) |
|
|
|
|
| def create_gpt_dataloader( |
| text: str, |
| tokenizer: BaseTokenizer, |
| max_len: int = 3, |
| stride: int = 1, |
| batch_size: int = 10, |
| num_workers: int = 0, |
| pin_memory: bool = False, |
| ) -> DataLoader: |
| """Create a DataLoader from raw text using sliding-window chunking.""" |
| dataset = GPTDataset(text, tokenizer, max_len, stride) |
| return DataLoader( |
| dataset, |
| batch_size=batch_size, |
| shuffle=True, |
| num_workers=num_workers, |
| pin_memory=pin_memory, |
| ) |
|
|
