src/rsil/data.py

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)