| """ |
| Dataset module for Part 2: Pretraining on browndw/human-ai-parallel-corpus. |
| Implements next-token prediction streaming and chunking for Chinchilla-style training dynamics. |
| """ |
|
|
| from typing import List, Dict, Any, Optional, Tuple |
| import torch |
| from torch.utils.data import Dataset, DataLoader |
| from transformers import PreTrainedTokenizerFast |
|
|
|
|
| class PretrainDataset(Dataset): |
| """ |
| Dataset for next-token prediction language modeling. |
| Extracts fixed-length token chunks from continuous document texts. |
| """ |
| def __init__( |
| self, |
| texts: List[str], |
| tokenizer: PreTrainedTokenizerFast, |
| seq_len: int = 256, |
| prefix_eval_ratio: float = 0.5 |
| ): |
| self.seq_len = seq_len |
| self.tokenizer = tokenizer |
| self.prefix_eval_ratio = prefix_eval_ratio |
| self.samples: List[torch.Tensor] = [] |
| self.eval_pairs: List[Tuple[str, str]] = [] |
|
|
| for text in texts: |
| if not text or not text.strip(): |
| continue |
| enc = tokenizer(text, truncation=False, return_tensors="pt")["input_ids"].squeeze(0) |
| if len(enc) < 8: |
| continue |
|
|
| |
| for i in range(0, len(enc) - seq_len, seq_len): |
| chunk = enc[i : i + seq_len] |
| self.samples.append(chunk) |
|
|
| |
| split_idx = int(len(text) * prefix_eval_ratio) |
| prompt = text[:split_idx].strip() |
| continuation = text[split_idx:].strip() |
| if prompt and continuation: |
| self.eval_pairs.append((prompt, continuation)) |
|
|
| def __len__(self) -> int: |
| return len(self.samples) |
|
|
| def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]: |
| chunk = self.samples[idx] |
| return { |
| "input_ids": chunk, |
| "targets": chunk |
| } |
|
|
