from __future__ import annotations import torch from torch.utils.data import DataLoader, Dataset, IterableDataset class SyntheticMLMDataset(Dataset): """Deterministic, completely offline masked-LM dataset for CI/smoke tests.""" def __init__(self, size=4096, seq_len=256, vocab_size=4096, mask_probability=0.15, seed=1234): if vocab_size < 8: raise ValueError("synthetic vocab_size must be >= 8") self.size, self.seq_len, self.vocab_size = int(size), int(seq_len), int(vocab_size) self.mask_probability, self.seed, self.mask_token_id = float(mask_probability), int(seed), 1 def __len__(self): return self.size def __getitem__(self, idx): g = torch.Generator().manual_seed(self.seed + int(idx)) original = torch.randint(2, self.vocab_size, (self.seq_len,), generator=g, dtype=torch.long) masked = torch.rand(self.seq_len, generator=g) < self.mask_probability if not masked.any(): masked[int(idx) % self.seq_len] = True labels = original.clone(); labels[~masked] = -100 input_ids = original.clone(); input_ids[masked] = self.mask_token_id return {"input_ids": input_ids, "labels": labels} class SyntheticTokenizerInfo: def __init__(self, vocab_size): self.vocab_size = int(vocab_size) def __len__(self): return self.vocab_size def build_synthetic_loader(seq_len, batch_size, vocab_size=4096, size=4096, mask_probability=0.15, seed=1234, num_workers=0): ds = SyntheticMLMDataset(size, seq_len, vocab_size, mask_probability, seed) return DataLoader(ds, batch_size=batch_size, shuffle=False, num_workers=num_workers), SyntheticTokenizerInfo(vocab_size) class PackedTokenStream(IterableDataset): """Pack a text iterable into fixed-length token blocks. Raw string columns never reach the DataLoader collator. Documents are tokenized without padding, concatenated through a rolling buffer, and emitted as exact ``seq_len`` LongTensor blocks. This avoids padding-heavy line-by-line WikiText batches and the v2.1 ``text``-column tensorization failure. """ def __init__(self, source, tokenizer, seq_len, text_column="text"): super().__init__(); self.source=source; self.tokenizer=tokenizer self.seq_len=int(seq_len); self.text_column=text_column def __iter__(self): buf=[] for ex in self.source: text=ex.get(self.text_column, "") if not isinstance(text, str) or not text.strip(): continue # Fast-tokenizer backend encoding deliberately bypasses Transformers' # model_max_length warning. This is safe here because the backend output # is never fed directly to a model: it is appended to a continuous token # stream and sliced below into exact ``seq_len`` blocks. No document # tokens are discarded merely because a source paragraph exceeds BERT's # advertised 512-token model limit. backend = getattr(self.tokenizer, "backend_tokenizer", None) if backend is not None: ids = backend.encode(text, add_special_tokens=False).ids else: # Slow/custom-tokenizer fallback: tokenize bounded character chunks. # Each chunk is intentionally small enough to stay below common # 512-token limits, while the rolling buffer still packs across chunks. ids = [] for start in range(0, len(text), 1000): piece = text[start:start + 1000] try: part = self.tokenizer( piece, add_special_tokens=False, truncation=False, verbose=False )["input_ids"] except TypeError: part = self.tokenizer( piece, add_special_tokens=False, truncation=False )["input_ids"] ids.extend(part) if not ids: continue # Separate documents when the tokenizer supplies a separator token. if self.tokenizer.sep_token_id is not None: ids = ids + [self.tokenizer.sep_token_id] buf.extend(ids) while len(buf) >= self.seq_len: block=buf[:self.seq_len]; del buf[:self.seq_len] yield {"input_ids": torch.tensor(block, dtype=torch.long)} class DeterministicMLMCollator: """MLM corruption for already-packed token blocks; no raw strings/padding.""" def __init__(self, tokenizer, mlm_probability=.15, seed=1234): if tokenizer.mask_token_id is None: raise ValueError("Tokenizer must provide a mask token") self.tok, self.p, self.seed, self.calls = tokenizer, float(mlm_probability), int(seed), 0 def __call__(self, features): x=torch.stack([f["input_ids"] for f in features]).clone(); labels=x.clone() g=torch.Generator().manual_seed(self.seed + self.calls); self.calls += 1 prob=torch.full(x.shape, self.p) # Never predict tokenizer-defined special tokens. special=torch.tensor([self.tok.get_special_tokens_mask(row.tolist(), already_has_special_tokens=True) for row in x], dtype=torch.bool) masked=(torch.rand(x.shape, generator=g) < prob) & ~special # Guarantee at least one supervised token per sequence. for i in range(masked.shape[0]): if not masked[i].any(): candidates=(~special[i]).nonzero(as_tuple=False).flatten() if len(candidates): masked[i, candidates[0]] = True labels[~masked] = -100 # BERT-style 80% [MASK], 10% random, 10% unchanged. replace=(torch.rand(x.shape, generator=g) < .8) & masked; x[replace]=self.tok.mask_token_id random_sel=(torch.rand(x.shape, generator=g) < .5) & masked & ~replace random_words=torch.randint(len(self.tok), x.shape, generator=g, dtype=torch.long); x[random_sel]=random_words[random_sel] return {"input_ids":x, "labels":labels} def build_streaming_loader(dataset, split, tokenizer_name, seq_len, batch_size, num_workers=0, shuffle_buffer=10000, dataset_config=None, seed=1234, mlm_probability=.15): from datasets import load_dataset from transformers import AutoTokenizer tok=AutoTokenizer.from_pretrained(tokenizer_name, use_fast=True) if tok.mask_token is None: raise ValueError("Tokenizer must provide a mask token for masked-LM training") ds=load_dataset(dataset, dataset_config, split=split, streaming=True) if shuffle_buffer and shuffle_buffer > 1: ds=ds.shuffle(seed=seed, buffer_size=shuffle_buffer) packed=PackedTokenStream(ds, tok, seq_len) coll=DeterministicMLMCollator(tok, mlm_probability=mlm_probability, seed=seed) # Keep zero workers for one-shard streams such as WikiText and reproducible masking. return DataLoader(packed, batch_size=batch_size, collate_fn=coll, num_workers=0), tok def build_benchmark_loader(dataset, split, tokenizer_name, seq_len, batch_size, dataset_config=None, synthetic_vocab_size=4096, synthetic_samples=4096, seed=1234): if dataset.lower() in {"synthetic", "offline-synthetic", "ci"}: return build_synthetic_loader(seq_len, batch_size, synthetic_vocab_size, synthetic_samples, seed=seed) return build_streaming_loader(dataset, split, tokenizer_name, seq_len, batch_size, dataset_config=dataset_config, seed=seed)