scripts/03_training/data.py

"""
Data loading pipeline for BabyLM training.
Handles text loading, optional Morfessor pre-segmentation, tokenization, and batching.
"""

import re
import random
from pathlib import Path
from typing import Optional

import torch
from torch.utils.data import Dataset, DataLoader

ROOT = Path(__file__).resolve().parent.parent.parent

# ===================================================
# Morfessor pre-segmentation (shared with tokenizer training)
# ===================================================

MIN_MORPH_LEN = 2
MIN_WORD_LEN = 3
_WORD_RE = re.compile(r'^([^a-zA-Z]*?)([a-zA-Z]+)([^a-zA-Z]*)$')


def presegment_word(word: str, morf_model) -> str:
    """Pre-segment a word using Morfessor, preserving case."""
    m = _WORD_RE.match(word)
    if not m:
        return word
    prefix, core, suffix = m.groups()
    if len(core) < MIN_WORD_LEN:
        return word
    segments = morf_model.viterbi_segment(core.lower())[0]
    if len(segments) <= 1 or not all(len(s) >= MIN_MORPH_LEN for s in segments):
        return word
    parts, pos = [], 0
    for seg in segments:
        n = len(seg)
        parts.append(core[pos:pos + n])
        pos += n
    return prefix + ' '.join(parts) + suffix


def presegment_text(text: str, morf_model) -> str:
    """Pre-segment entire text line using Morfessor."""
    return ' '.join(presegment_word(w, morf_model) for w in text.split())


def load_morfessor_model(model_path: str):
    """Load a trained Morfessor model."""
    import morfessor
    io = morfessor.MorfessorIO()
    return io.read_binary_model_file(model_path)


# ===================================================
# Text Dataset
# ===================================================

class TextLineDataset(Dataset):
    """
    Dataset that reads text lines, optionally pre-segments with Morfessor,
    tokenizes with HuggingFace tokenizer, and returns fixed-length chunks.
    """

    def __init__(
        self,
        text_path: str,
        tokenizer,
        max_seq_len: int = 128,
        morf_model=None,
    ):
        self.tokenizer = tokenizer
        self.max_seq_len = max_seq_len
        self.morf_model = morf_model

        # Read and tokenize all text into one long token sequence
        print(f"Loading and tokenizing {text_path}...")
        text_path = Path(text_path)
        if not text_path.is_absolute():
            text_path = ROOT / text_path

        all_ids = []
        with open(text_path) as f:
            for i, line in enumerate(f):
                line = line.strip()
                if not line:
                    continue
                if self.morf_model is not None:
                    line = presegment_text(line, self.morf_model)
                ids = tokenizer.encode(line, add_special_tokens=False)
                all_ids.extend(ids)
                if (i + 1) % 500000 == 0:
                    print(f"  Processed {i+1:,} lines, {len(all_ids):,} tokens so far...")

        self.all_ids = torch.tensor(all_ids, dtype=torch.long)
        # Split into non-overlapping chunks of max_seq_len
        n_chunks = len(self.all_ids) // max_seq_len
        self.all_ids = self.all_ids[:n_chunks * max_seq_len]
        self.chunks = self.all_ids.view(n_chunks, max_seq_len)
        print(f"  Total: {len(all_ids):,} tokens -> {n_chunks:,} chunks of {max_seq_len}")

    def __len__(self):
        return len(self.chunks)

    def __getitem__(self, idx):
        return self.chunks[idx]


class SentenceDataset(Dataset):
    """
    Per-sentence dataset: each sentence is an independent sample.
    Short sentences are padded to max_seq_len, long ones truncated.
    Unlike TextLineDataset (which packs all tokens into fixed chunks),
    this preserves sentence boundaries.
    """

    def __init__(self, text_path: str, tokenizer, max_seq_len: int = 128,
                 morf_model=None):
        self.tokenizer = tokenizer
        self.max_seq_len = max_seq_len
        self.pad_id = tokenizer.convert_tokens_to_ids("<pad>")
        if self.pad_id is None:
            self.pad_id = 0

        print(f"Loading sentences from {text_path}...")
        text_path = Path(text_path)
        if not text_path.is_absolute():
            text_path = ROOT / text_path

        self.sentences = []
        total_tokens = 0
        with open(text_path) as f:
            for i, line in enumerate(f):
                line = line.strip()
                if not line:
                    continue
                if morf_model is not None:
                    line = presegment_text(line, morf_model)
                ids = tokenizer.encode(line, add_special_tokens=False)
                if len(ids) < 3:
                    continue  # skip very short lines
                # Truncate to max_seq_len
                ids = ids[:max_seq_len]
                self.sentences.append(torch.tensor(ids, dtype=torch.long))
                total_tokens += len(ids)
                if (i + 1) % 500000 == 0:
                    print(f"  Processed {i+1:,} lines, {len(self.sentences):,} sentences...")

        print(f"  Total: {total_tokens:,} tokens, {len(self.sentences):,} sentences")
        # Pre-compute for FrequencyMasker compatibility
        self.chunks = self.sentences  # alias for token counting

    def __len__(self):
        return len(self.sentences)

    def __getitem__(self, idx):
        return self.sentences[idx]


def sentence_collate_fn(batch, pad_id: int = 0):
    """Collate variable-length sentences into a padded batch."""
    max_len = max(len(s) for s in batch)
    padded = torch.full((len(batch), max_len), pad_id, dtype=torch.long)
    for i, s in enumerate(batch):
        padded[i, :len(s)] = s
    return padded


# ===================================================
# Masking strategies
# ===================================================

class StandardMasker:
    """Standard random masking for MLM/MNTP with updatable mask_ratio."""

    def __init__(self, tokenizer, mask_ratio: float = 0.30):
        self.mask_token_id = tokenizer.convert_tokens_to_ids("<mask>")
        self.vocab_size = tokenizer.vocab_size
        self.mask_ratio = mask_ratio
        # Special token IDs to never mask
        self.special_ids = set()
        for name in ["bos_token", "eos_token", "pad_token", "unk_token", "mask_token"]:
            tid = getattr(tokenizer, name + "_id", None)
            if tid is not None:
                self.special_ids.add(tid)

    def set_mask_ratio(self, ratio: float):
        """Update the mask ratio (used for mask rate decay)."""
        self.mask_ratio = ratio

    def __call__(self, input_ids: torch.Tensor) -> tuple:
        """
        Apply random masking.
        Returns: (masked_input_ids, labels) where labels=-100 for non-masked positions.
        """
        labels = input_ids.clone()
        masked_ids = input_ids.clone()

        # Create mask probability tensor (ensure float even if mask_ratio is int 0)
        prob = torch.full(input_ids.shape, float(self.mask_ratio))
        # Don't mask special tokens
        for sid in self.special_ids:
            prob[input_ids == sid] = 0.0

        mask = torch.bernoulli(prob).bool()
        labels[~mask] = -100  # Only compute loss on masked tokens

        # 80% [MASK], 10% random, 10% unchanged
        indices_mask = mask & (torch.rand(input_ids.shape) < 0.8)
        indices_random = mask & ~indices_mask & (torch.rand(input_ids.shape) < 0.5)

        masked_ids[indices_mask] = self.mask_token_id
        random_tokens = torch.randint(5, self.vocab_size, input_ids.shape)
        masked_ids[indices_random] = random_tokens[indices_random]

        return masked_ids, labels


class FrequencyMasker:
    """Frequency-aware masking: low-frequency tokens get higher mask probability.

    Computes token frequencies from the training data, then assigns mask
    probabilities inversely proportional to frequency. Interpolates between
    frequency-based and uniform masking via alpha parameter.

    mask_prob[t] = alpha * normalized_inv_freq[t] + (1-alpha) * uniform
    """

    def __init__(self, tokenizer, token_counts: torch.Tensor,
                 mask_ratio: float = 0.30, alpha: float = 0.3):
        self.mask_token_id = tokenizer.convert_tokens_to_ids("<mask>")
        self.vocab_size = tokenizer.vocab_size
        self.mask_ratio = mask_ratio
        self.alpha = alpha
        self.special_ids = set()
        for name in ["bos_token", "eos_token", "pad_token", "unk_token", "mask_token"]:
            tid = getattr(tokenizer, name + "_id", None)
            if tid is not None:
                self.special_ids.add(tid)

        # Compute per-token mask probability from inverse frequency
        # token_counts: [vocab_size] tensor of token occurrence counts
        freq = token_counts.float() + 1.0  # Laplace smoothing
        inv_freq = 1.0 / freq
        # Normalize so mean = 1.0
        inv_freq = inv_freq / inv_freq.mean()
        self.per_token_weight = inv_freq  # [vocab_size]

    def set_mask_ratio(self, ratio: float):
        self.mask_ratio = ratio

    def __call__(self, input_ids: torch.Tensor) -> tuple:
        labels = input_ids.clone()
        masked_ids = input_ids.clone()

        # Per-token mask probability: weighted by inverse frequency
        uniform = torch.full(input_ids.shape, 1.0)
        freq_weight = self.per_token_weight[input_ids]  # lookup per token
        blended = self.alpha * freq_weight + (1.0 - self.alpha) * uniform
        # Scale so mean probability = mask_ratio
        prob = blended * (self.mask_ratio / blended.mean())
        prob = prob.clamp(0.0, 0.95)

        # Don't mask special tokens
        for sid in self.special_ids:
            prob[input_ids == sid] = 0.0

        mask = torch.bernoulli(prob).bool()
        labels[~mask] = -100

        indices_mask = mask & (torch.rand(input_ids.shape) < 0.8)
        indices_random = mask & ~indices_mask & (torch.rand(input_ids.shape) < 0.5)
        masked_ids[indices_mask] = self.mask_token_id
        random_tokens = torch.randint(5, self.vocab_size, input_ids.shape)
        masked_ids[indices_random] = random_tokens[indices_random]

        return masked_ids, labels


class AMLMMasker:
    """
    Adaptive Masked Language Modeling (Hard AMLM, accuracy-based).
    From Edman & Fraser 2025.

    Adjusts per-token mask rate based on model's prediction accuracy.
    Updates every `update_interval` steps using Laplace-smoothed accuracy:
        score = (correct + 0.5) / (total + 1)
    Mask probabilities are normalized so the mean equals the base mask rate.
    Lambda controls interpolation between uniform and adaptive masking.
    """

    def __init__(self, tokenizer, vocab_size: int,
                 base_mask_ratio: float = 0.30,
                 amlm_lambda: float = 0.2,
                 update_interval: int = 200):
        self.mask_token_id = tokenizer.convert_tokens_to_ids("<mask>")
        self.vocab_size = vocab_size
        self.base_mask_ratio = base_mask_ratio
        self.amlm_lambda = amlm_lambda
        self.update_interval = update_interval
        self.special_ids = set()
        for name in ["bos_token", "eos_token", "pad_token", "unk_token", "mask_token"]:
            tid = getattr(tokenizer, name + "_id", None)
            if tid is not None:
                self.special_ids.add(tid)

        # Per-token accuracy tracking (reset each update interval)
        self.token_correct = torch.zeros(vocab_size)
        self.token_total = torch.zeros(vocab_size)

        # Per-token mask probabilities (start uniform)
        self.token_mask_prob = torch.full((vocab_size,), base_mask_ratio)

        self.steps_since_update = 0

    def set_mask_ratio(self, ratio: float):
        """Update the base mask ratio (used for mask rate decay)."""
        self.base_mask_ratio = ratio

    def update_accuracy(self, token_ids: torch.Tensor, predictions: torch.Tensor):
        """
        Record whether predictions were correct for masked tokens.
        Called after each training step.

        Args:
            token_ids: ground truth token IDs [N]
            predictions: predicted token IDs [N]
        """
        with torch.no_grad():
            token_ids_flat = token_ids.flatten().cpu()
            predictions_flat = predictions.flatten().cpu()
            correct = (token_ids_flat == predictions_flat)

            for tid, is_correct in zip(token_ids_flat, correct):
                tid = tid.item()
                if 0 <= tid < self.vocab_size:
                    self.token_total[tid] += 1
                    if is_correct:
                        self.token_correct[tid] += 1

        self.steps_since_update += 1
        if self.steps_since_update >= self.update_interval:
            self._recompute_mask_probs()
            self.steps_since_update = 0
            # Reset counters
            self.token_correct.zero_()
            self.token_total.zero_()

    def _recompute_mask_probs(self):
        """Recompute per-token mask probabilities from accuracy stats."""
        # Laplace-smoothed accuracy: score = (correct + 0.5) / (total + 1)
        scores = (self.token_correct + 0.5) / (self.token_total + 1.0)

        # Higher accuracy -> higher mask probability (mask easy tokens more)
        # Invert: mask_prob proportional to (1 - score) so harder tokens get masked more?
        # Actually in AMLM, tokens the model gets RIGHT should be masked LESS,
        # tokens the model gets WRONG should be masked MORE.
        # score is accuracy, so low score = hard = mask more
        raw_prob = 1.0 - scores

        # Normalize so mean probability = base mask rate
        current_mean = raw_prob.mean()
        if current_mean > 0:
            raw_prob = raw_prob * (self.base_mask_ratio / current_mean)

        # Clamp to reasonable range
        raw_prob = raw_prob.clamp(0.01, 0.80)

        # Interpolate with uniform: lambda * adaptive + (1-lambda) * uniform
        uniform = torch.full_like(raw_prob, self.base_mask_ratio)
        self.token_mask_prob = self.amlm_lambda * raw_prob + (1.0 - self.amlm_lambda) * uniform

    def __call__(self, input_ids: torch.Tensor) -> tuple:
        """Apply adaptive masking based on per-token accuracy."""
        labels = input_ids.clone()
        masked_ids = input_ids.clone()

        # Get per-token mask probability
        prob = self.token_mask_prob[input_ids.cpu()].to(input_ids.device)

        # Don't mask special tokens
        for sid in self.special_ids:
            prob[input_ids == sid] = 0.0

        mask = torch.bernoulli(prob).bool()
        labels[~mask] = -100

        indices_mask = mask & (torch.rand(input_ids.shape, device=input_ids.device) < 0.8)
        indices_random = mask & ~indices_mask & (torch.rand(input_ids.shape, device=input_ids.device) < 0.5)

        masked_ids[indices_mask] = self.mask_token_id
        random_tokens = torch.randint(5, self.vocab_size, input_ids.shape, device=input_ids.device)
        masked_ids[indices_random] = random_tokens[indices_random]

        return masked_ids, labels


# ===================================================
# Collation functions
# ===================================================

def create_masker(masking_cfg, tokenizer):
    """Factory function to create a masker from config."""
    if masking_cfg.type == "standard":
        return StandardMasker(tokenizer, mask_ratio=masking_cfg.mask_ratio)
    elif masking_cfg.type == "amlm":
        return AMLMMasker(
            tokenizer,
            vocab_size=tokenizer.vocab_size,
            base_mask_ratio=masking_cfg.mask_ratio,
            amlm_lambda=masking_cfg.amlm_lambda,
            update_interval=masking_cfg.amlm_update_interval,
        )
    elif masking_cfg.type == "frequency":
        # Need token counts from dataset — will be set by build_dataloader
        return FrequencyMasker(
            tokenizer,
            token_counts=torch.ones(tokenizer.vocab_size),  # placeholder, updated later
            mask_ratio=masking_cfg.mask_ratio,
            alpha=getattr(masking_cfg, 'freq_alpha', 0.3),
        )
    else:
        raise ValueError(f"Unknown masking type: {masking_cfg.type}")


class GPTBertCollator:
    """
    Collator for GPT-BERT dual objective:
    - 15 MNTP batches per 1 CLM batch (15:1 ratio)
    - MNTP labels are SHIFTED: position k's label = original token at k+1
    - Supports mask rate decay over training

    From Edman & Fraser 2025 "Mask and You Shall Receive".
    """

    def __init__(self, masker, bos_token_id: int = 1, mntp_ratio: int = 15,
                 mask_ratio_start: float = 0.30, mask_ratio_end: float = 0.15,
                 total_steps: int = 0):
        self.masker = masker
        self.bos_token_id = bos_token_id
        self.mntp_ratio = mntp_ratio
        self.mask_ratio_start = mask_ratio_start
        self.mask_ratio_end = mask_ratio_end
        self.total_steps = total_steps
        self.step = 0

    def _update_mask_ratio(self):
        """Linearly decay mask ratio from start to end over training."""
        if self.total_steps > 0:
            progress = min(self.step / self.total_steps, 1.0)
            current_ratio = self.mask_ratio_start + (self.mask_ratio_end - self.mask_ratio_start) * progress
            self.masker.set_mask_ratio(current_ratio)

    def __call__(self, batch):
        input_ids = torch.stack(batch)  # [B, seq_len]
        self.step += 1
        self._update_mask_ratio()

        # Every (mntp_ratio+1) steps, one CLM batch; otherwise MNTP
        if self.step % (self.mntp_ratio + 1) == 0:
            # CLM: predict next token
            labels = input_ids.clone()
            labels[:, :-1] = input_ids[:, 1:]
            labels[:, -1] = -100
            return {
                "input_ids": input_ids,
                "labels": labels,
                "task": "clm",
            }
        else:
            # MNTP: mask + shifted labels (position k predicts token k+1)
            masked_ids, mask_labels = self.masker(input_ids)

            # Shift labels: position k's label = original token at k+1
            shifted_labels = torch.full_like(mask_labels, -100)
            # Only set shifted labels where masking occurred (mask_labels != -100)
            mask_positions = (mask_labels != -100)
            # For masked positions, the label is the NEXT token in the original sequence
            # Shift: for position k, label = input_ids[k+1]
            shifted_labels[:, :-1] = torch.where(
                mask_positions[:, :-1],
                input_ids[:, 1:],
                torch.tensor(-100, dtype=input_ids.dtype)
            )
            # Last position can't predict next token
            shifted_labels[:, -1] = -100

            return {
                "input_ids": masked_ids,
                "labels": shifted_labels,
                "task": "mntp",
            }


class CLMCollator:
    """CLM collator with optional multi-token prediction (MTP).

    When mtp_k=1 (default), standard next-token prediction.
    When mtp_k=2, each position predicts the token 2 steps ahead.
    This is used for reverse curriculum MTP:
      - First half of training: k=2 (harder task, builds long-range representation)
      - Second half: k=1 (standard, fine-grained prediction)
    """

    def __init__(self, bos_token_id: int = 1, mtp_k: int = 1):
        self.bos_token_id = bos_token_id
        self.mtp_k = mtp_k

    def set_mtp_k(self, k: int):
        """Update the prediction horizon (called by training loop for curriculum)."""
        self.mtp_k = k

    def __call__(self, batch):
        input_ids = torch.stack(batch)
        k = self.mtp_k
        labels = torch.full_like(input_ids, -100)
        # Position i predicts token at position i+k
        if k < input_ids.shape[1]:
            labels[:, :-k] = input_ids[:, k:]
        return {"input_ids": input_ids, "labels": labels, "task": "clm"}


class MLMCollator:
    """MLM collator with masking."""

    def __init__(self, masker):
        self.masker = masker

    def __call__(self, batch):
        input_ids = torch.stack(batch)
        masked_ids, labels = self.masker(input_ids)
        return {"input_ids": masked_ids, "labels": labels, "task": "mlm"}


def create_collator(objective: str, masker, tokenizer, training_cfg=None):
    """Factory to create the right collator based on training objective."""
    bos_id = tokenizer.convert_tokens_to_ids("<s>")
    if objective == "gpt_bert":
        mntp_ratio = getattr(training_cfg, 'mntp_ratio', 15) if training_cfg else 15
        # Compute total steps for mask decay if possible
        total_steps = 0  # Will be set externally if needed
        mask_ratio_start = 0.30
        mask_ratio_end = 0.15
        if hasattr(masker, 'mask_ratio'):
            mask_ratio_start = masker.mask_ratio
        return GPTBertCollator(
            masker, bos_token_id=bos_id, mntp_ratio=mntp_ratio,
            mask_ratio_start=mask_ratio_start, mask_ratio_end=mask_ratio_end,
            total_steps=total_steps,
        )
    elif objective == "clm":
        mtp_k = getattr(training_cfg, 'mtp_k_start', 1) if training_cfg and getattr(training_cfg, 'use_mtp', False) else 1
        return CLMCollator(bos_token_id=bos_id, mtp_k=mtp_k)
    elif objective in ("mlm", "mntp", "amlm"):
        return MLMCollator(masker)
    elif objective == "rtd":
        return MLMCollator(masker)  # RTD uses same masking, model handles the rest
    else:
        raise ValueError(f"Unknown objective: {objective}")


def build_dataloader(cfg, tokenizer):
    """
    Build complete DataLoader from config.

    Args:
        cfg: ExperimentConfig
        tokenizer: HuggingFace tokenizer
    Returns:
        DataLoader, masker (masker needed for AMLM updates)
    """
    # Load Morfessor if needed
    morf_model = None
    if cfg.data.tokenizer == "morfessor_bpe" and cfg.data.morfessor_model_path:
        morf_model = load_morfessor_model(cfg.data.morfessor_model_path)
        print(f"Loaded Morfessor model from {cfg.data.morfessor_model_path}")

    # Create dataset
    packing = getattr(cfg.data, 'packing', 'concat')
    if packing == "sentence":
        dataset = SentenceDataset(
            text_path=cfg.data.train_file,
            tokenizer=tokenizer,
            max_seq_len=cfg.data.max_seq_len,
            morf_model=morf_model,
        )
    else:
        dataset = TextLineDataset(
            text_path=cfg.data.train_file,
            tokenizer=tokenizer,
            max_seq_len=cfg.data.max_seq_len,
            morf_model=morf_model,
        )

    # Create masker
    masker = create_masker(cfg.masking, tokenizer)

    # For frequency masking, compute actual token counts from dataset
    if isinstance(masker, FrequencyMasker):
        token_counts = torch.zeros(tokenizer.vocab_size, dtype=torch.long)
        for ids in dataset.chunks:
            for tid in ids:
                if tid < tokenizer.vocab_size:
                    token_counts[tid] += 1
        inv_freq = 1.0 / (token_counts.float() + 1.0)
        masker.per_token_weight = inv_freq / inv_freq.mean()
        print(f"  FrequencyMasker: computed token frequencies from {token_counts.sum().item():,} tokens")

    # Create collator
    collator = create_collator(cfg.training.objective, masker, tokenizer,
                               training_cfg=cfg.training)

    # Set total steps for mask decay in GPTBertCollator
    if isinstance(collator, GPTBertCollator):
        steps_per_epoch = len(dataset) // cfg.training.batch_size
        total_steps = steps_per_epoch * cfg.training.epochs
        collator.total_steps = total_steps
        collator.mask_ratio_start = cfg.masking.mask_ratio
        collator.mask_ratio_end = cfg.masking.mask_ratio_end

    # For sentence packing, wrap collator to pad variable-length sequences first
    if packing == "sentence":
        pad_id = tokenizer.convert_tokens_to_ids("<pad>") or 0
        base_collator = collator
        def padded_collator(batch):
            padded = sentence_collate_fn(batch, pad_id=pad_id)
            # base_collator expects a list of equal-length tensors
            return base_collator([padded[i] for i in range(padded.shape[0])])
        final_collator = padded_collator
    else:
        final_collator = collator

    # Create DataLoader
    loader = DataLoader(
        dataset,
        batch_size=cfg.training.batch_size,
        shuffle=True,
        num_workers=4,
        pin_memory=True,
        collate_fn=final_collator,
        drop_last=True,
    )

    return loader, masker