data_utils.py

from __future__ import annotations

from typing import Any, Callable, Dict, Iterable, List, Optional, Sequence, Tuple

from torch.utils.data import IterableDataset, get_worker_info

import torch
import random

def build_grad_attention_mask(
    *,
    seq_lens: List[int],
    think_open_positions: List[int],
    think_close_positions: List[int],
    max_len: int,
) -> torch.Tensor:
    """
    Build additive attention mask for grad pass with shape [B, 1, T, T].
    0.0 means allowed attention; finfo.min means blocked.

    Rule:
      - Standard causal masking for all tokens.
      - For token positions strictly after </think>, allow attention only to
        tokens inside the <think>...</think> span (including markers).
    """
    bsz = len(seq_lens)
    min_val = torch.finfo(torch.float32).min
    mask = torch.full((bsz, 1, max_len, max_len), min_val, dtype=torch.float32)
    causal = torch.tril(torch.ones((max_len, max_len), dtype=torch.bool))

    for b, (seq_len, think_open_ix, think_close_ix) in enumerate(
        zip(seq_lens, think_open_positions, think_close_positions)
    ):
        allowed = causal.clone()
        if seq_len < max_len:
            allowed[seq_len:, :] = False
            allowed[:, seq_len:] = False

        post_think_start = think_close_ix + 1
        if post_think_start < seq_len:
            allowed[post_think_start:seq_len, :] = False
            span_start = max(0, think_open_ix)
            span_end = min(seq_len - 1, think_close_ix)
            if span_start <= span_end:
                allowed[post_think_start:seq_len, span_start : span_end + 1] = True

        mask[b, 0][allowed] = 0.0

    return mask


def pad_1d(seqs: List[List[int]], pad_id: int) -> Tuple[torch.Tensor, torch.Tensor]:
    max_len = max(len(s) for s in seqs) if seqs else 0
    input_ids = torch.full((len(seqs), max_len), pad_id, dtype=torch.long)
    attn = torch.zeros((len(seqs), max_len), dtype=torch.long)
    for i, s in enumerate(seqs):
        if not s:
            continue
        input_ids[i, : len(s)] = torch.tensor(s, dtype=torch.long)
        attn[i, : len(s)] = 1
    return input_ids, attn


def make_lemath_collate_fn(*, pad_id: int) -> Callable[[List[Dict[str, Any]]], Dict[str, Any]]:
    """
    Returns a DataLoader-compatible collate_fn that pads variable-length fields and
    produces a plain dict of tensors/lists.
    """

    def _collate(examples: List[Dict[str, Any]]) -> Dict[str, Any]:
        base_input_ids_list = [ex["base_input_ids"] for ex in examples]
        input_ids_list = [ex["input_ids"] for ex in examples]
        labels_list = [ex["labels"] for ex in examples]
        cross_pairs = [ex["cross_pairs"] for ex in examples]
        num_jetons_list = [ex["num_jetons"] for ex in examples]
        think_open_positions = [int(ex["think_open_ix"]) for ex in examples]
        think_close_positions = [int(ex["think_close_ix"]) for ex in examples]

        base_input_ids, base_attention_mask = pad_1d(base_input_ids_list, pad_id=pad_id)
        input_ids, attention_mask = pad_1d(input_ids_list, pad_id=pad_id)
        labels, _ = pad_1d(labels_list, pad_id=-100)
        grad_attention_mask = build_grad_attention_mask(
            seq_lens=[len(ids) for ids in input_ids_list],
            think_open_positions=think_open_positions,
            think_close_positions=think_close_positions,
            max_len=input_ids.shape[1],
        )
        return {
            "base_input_ids": base_input_ids,
            "base_attention_mask": base_attention_mask,
            "input_ids": input_ids,
            "attention_mask": attention_mask,
            "grad_attention_mask": grad_attention_mask,
            "labels": labels,
            "cross_pairs": cross_pairs,
            "num_jetons": num_jetons_list,
        }

    return _collate


class EncodedIterableDataset(IterableDataset):
    """
    Wraps an (HF) iterable of raw examples, encoding them on-the-fly.

    - Supports multi-worker DataLoader by sharding by (global_index % num_workers).
    - Filters out examples where `encode_fn` returns None.
    """

    def __init__(self, raw_examples: Iterable[Dict[str, Any]], *, encode_fn: Callable[[Dict[str, Any]], Optional[Dict[str, Any]]]):
        super().__init__()
        self.raw_examples = raw_examples
        self.encode_fn = encode_fn

    def __iter__(self):
        it = iter(self.raw_examples)
        wi = get_worker_info()
        if wi is None:
            for ex in it:
                enc = self.encode_fn(ex)
                if enc is None:
                    continue
                yield enc
            return

        # Simple deterministic sharding across workers.
        worker_id = int(wi.id)
        num_workers = int(wi.num_workers)
        for idx, ex in enumerate(it):
            if (idx % num_workers) != worker_id:
                continue
            enc = self.encode_fn(ex)
            if enc is None:
                continue
            yield enc

def apply_chat(tokenizer, problem_text: str, assistant_text: str, system_prompt: Optional[str] = None) -> Tuple[List[int], int]:
    """
    Returns (input_ids, prompt_len) such that labels can be masked with -100 for [0:prompt_len].
    Uses the tokenizer chat template when available/allowed.
    """

    messages: List[Dict[str, str]] = []
    if system_prompt:
        messages.append({"role": "system", "content": system_prompt})
    messages.append({"role": "user", "content": problem_text})
    messages.append({"role": "assistant", "content": assistant_text})

    prompt_text = tokenizer.apply_chat_template(messages[:-1], tokenize=False, add_generation_prompt=True)
    full_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=False)

    prompt_ids = tokenizer(prompt_text, add_special_tokens=False).input_ids
    input_ids = tokenizer(full_text, add_special_tokens=False).input_ids
    return input_ids, len(prompt_ids)

def encode_one_example(
    *,
    ex: Dict[str, Any],
    prompt_template: str,
    jeton_token: str,
    jeton_id: int,
    tokens_per_jeton: int,
    max_seq_len: int,
    think_open_id: int,
    think_close_id: int,
    tokenizer: Any,
    system_prompt: Optional[str] = None,
) -> Optional[Dict[str, Any]]:
    problem = ex.get("problem")
    solution = ex.get("generated_solution")

    if not isinstance(problem, str) or not isinstance(solution, str):
        return None

    user_text = prompt_template.format(problem=problem)
    # Better: build two inputs via apply_chat twice:
    #   1) original solution, cut at </think> (for base / no-grad pass)
    #   2) solution where think content is replaced by jetons (for grad pass)
    open_s = solution.find("<think>")
    close_s = solution.find("</think>", open_s + len("<think>")) if open_s != -1 else -1
    align_token = tokenizer.convert_tokens_to_ids(["\n\n"])[0]

    if open_s < 0 and close_s < 0:
        raise ValueError("No <think> tags found in solution.")

    base_solution = solution[:close_s] + "</think>"
    base_input_ids, _ = apply_chat(tokenizer, user_text, base_solution, system_prompt)
    base_input_ids = base_input_ids[:max_seq_len]
    base_think_open = base_input_ids.index(think_open_id)
    base_think_close = base_input_ids.index(think_close_id) if think_close_id in base_input_ids else len(base_input_ids)

    think_len = base_think_close - base_think_open
    effort_base = 768
    reasoning_effort = max(min(10, think_len // effort_base), 1)
    num_jetons = reasoning_effort * (effort_base // tokens_per_jeton)

    grad_solution = solution[:open_s] + "<think>\n" + jeton_token * num_jetons + "\n</think>" + solution[close_s + len("</think>"):]
    
    input_ids, _ = apply_chat(tokenizer, user_text, grad_solution, system_prompt + f"\nReasoning Effort: {reasoning_effort}")

    think_open_ix = input_ids.index(think_open_id)
    think_close_ix = input_ids.index(think_close_id)

    labels = [-100] * think_close_ix + input_ids[think_close_ix:]
    assert len(labels) == len(input_ids)

    last_jeton = len(input_ids) - 1 - input_ids[::-1].index(jeton_id)
    last_think_token = base_think_close - random.choice([1, 2, 3])
    cross_pairs = [(last_jeton, last_think_token)]

    return {
        "base_input_ids": base_input_ids,
        "input_ids": input_ids,
        "labels": labels,
        "cross_pairs": cross_pairs,
        "num_jetons": num_jetons,
        "think_open_ix": think_open_ix,
        "think_close_ix": think_close_ix,
    }