File size: 1,935 Bytes
f90f02b | 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 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 | """
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
# Split into chunks of size seq_len
for i in range(0, len(enc) - seq_len, seq_len):
chunk = enc[i : i + seq_len]
self.samples.append(chunk)
# Store prefix and continuation for BLEU evaluation
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
}
|