exp/exp2/dataset_utils.py

"""Dataset helpers for Experiment 2 (CoT / multi-hop faithfulness).

Named dataset_utils to avoid collision with the HF `datasets` package.
"""

from __future__ import annotations

import json
import random
import re
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, Iterable, List, Optional

from attribution_datasets import (
    AttributionExample,
    MoreHopQAAttributionDataset,
    RulerAttributionDataset,
)


@dataclass
class CachedExample:
    prompt: str
    target: Optional[str]
    indices_to_explain: Optional[List[int]]
    attr_mask_indices: Optional[List[int]]
    sink_span: Optional[List[int]]
    thinking_span: Optional[List[int]]
    metadata: Dict[str, Any]


def read_cached_jsonl(path: Path) -> List[CachedExample]:
    examples: List[CachedExample] = []
    with path.open("r", encoding="utf-8") as f:
        for line in f:
            if not line.strip():
                continue
            obj = json.loads(line)
            examples.append(
                CachedExample(
                    prompt=obj["prompt"],
                    target=obj.get("target"),
                    indices_to_explain=obj.get("indices_to_explain"),
                    attr_mask_indices=obj.get("attr_mask_indices"),
                    sink_span=obj.get("sink_span"),
                    thinking_span=obj.get("thinking_span"),
                    metadata=obj.get("metadata", {}),
                )
            )
    return examples


def load_cached(path: Path, sample: Optional[int] = None, seed: int = 42) -> List[CachedExample]:
    ex = read_cached_jsonl(path)
    if sample is not None and sample < len(ex):
        random.Random(seed).shuffle(ex)
        ex = ex[:sample]
    return ex


def load_ruler(path: Path, sample: Optional[int] = None, seed: int = 42) -> List[CachedExample]:
    ds = RulerAttributionDataset(path)
    examples: List[CachedExample] = []
    ex_iter: Iterable[AttributionExample] = ds
    if sample is not None and sample < len(ds):
        ex_iter = list(ds)
        random.Random(seed).shuffle(ex_iter)
        ex_iter = ex_iter[:sample]
    for ex in ex_iter:
        examples.append(
            CachedExample(
                prompt=ex.prompt,
                target=ex.target,
                indices_to_explain=ex.indices_to_explain,
                attr_mask_indices=ex.attr_mask_indices,
                sink_span=None,
                thinking_span=None,
                metadata=ex.metadata,
            )
        )
    return examples


def load_morehopqa(
    path: str | Path = "./data/with_human_verification.json", sample: Optional[int] = None, seed: int = 42
) -> List[CachedExample]:
    ds = MoreHopQAAttributionDataset(path)
    ex_iter: Iterable[AttributionExample] = ds
    if sample is not None and sample < len(ds):
        ex_iter = list(ds)
        random.Random(seed).shuffle(ex_iter)
        ex_iter = ex_iter[:sample]
    examples: List[CachedExample] = []
    for ex in ex_iter:
        examples.append(
            CachedExample(
                prompt=ex.prompt,
                target=None,
                indices_to_explain=ex.indices_to_explain,
                attr_mask_indices=ex.attr_mask_indices,
                sink_span=None,
                thinking_span=None,
                metadata=ex.metadata,
            )
        )
    return examples


def auto_find_ruler(task: str) -> Optional[Path]:
    length_dirs = ["4096", "8192", "16384", "32768", "65536", "131072"]
    base = Path("data/ruler_multihop")
    for ld in length_dirs:
        cand = base / ld / task / "validation.jsonl"
        if cand.exists():
            return cand
    return None


def dataset_from_name(name: str) -> Optional[Path]:
    if name == "hotpotqa_long":
        return auto_find_ruler("hotpotqa_long")
    if name.startswith("vt_"):
        return auto_find_ruler(name)
    if name.startswith("niah"):
        return auto_find_ruler(name)
    p = Path(name)
    if p.exists():
        return p
    return None


_BOX_PATTERN = re.compile(r"\\box(?:ed)?\s*[\{｛](.*?)[\}｝]", flags=re.DOTALL)


def _find_box_span(text: str) -> Optional[tuple[int, int, str]]:
    """Return (start_char, end_char, answer_text) for the last \\boxed block."""
    matches = list(_BOX_PATTERN.finditer(text))
    if not matches:
        return None
    m = matches[-1]
    return m.start(0), m.end(0), m.group(1).strip()


def extract_boxed_answer(text: str) -> Optional[str]:
    """Extract the answer string inside the last \\boxed{} block."""
    match = _find_box_span(text)
    return match[2] if match else None


def _find_answer_span(text: str, answer: str) -> Optional[tuple[int, int]]:
    """Return (start_char, end_char) for the last occurrence of `answer` in text."""
    if not answer or not text:
        return None
    start = text.rfind(answer)
    if start == -1:
        return None
    return start, start + len(answer)


def split_boxed_generation(text: str) -> Optional[tuple[str, str, str]]:
    """Return (thinking_text, boxed_segment, boxed_answer) if format matches."""
    if not text:
        return None
    match = _find_box_span(text)
    if not match:
        return None

    start_char, end_char, boxed_inner = match
    boxed_segment = text[start_char:end_char].strip()
    thinking_text = text[:start_char].strip()
    trailing = text[end_char:].strip()

    if not boxed_inner or not boxed_segment:
        return None
    if trailing:
        return None
    if not thinking_text:
        return None

    return thinking_text, boxed_segment, boxed_inner


def attach_spans_from_answer(
    example: CachedExample, tokenizer, answer_text: Optional[str] = None
) -> CachedExample:
    """Attach sink/thinking spans by locating the (plain) answer in `target`.

    `answer_text` should be the extracted boxed answer; falls back to metadata or
    parsing the target when omitted. Works even when the target no longer keeps
    the \\box{} wrapper.
    """
    tgt = example.target or ""
    answer = (answer_text or "").strip()
    if not answer:
        answer = (example.metadata.get("boxed_answer") or extract_boxed_answer(tgt) or "").strip()

    metadata = dict(example.metadata)
    if answer:
        metadata.setdefault("boxed_answer", answer)

    if tokenizer is None or not tgt or not answer:
        return CachedExample(
            prompt=example.prompt,
            target=example.target,
            indices_to_explain=example.indices_to_explain,
            attr_mask_indices=example.attr_mask_indices,
            sink_span=example.sink_span,
            thinking_span=example.thinking_span,
            metadata=metadata,
        )

    span = _find_answer_span(tgt, answer)
    if span is None:
        return CachedExample(
            prompt=example.prompt,
            target=example.target,
            indices_to_explain=example.indices_to_explain,
            attr_mask_indices=example.attr_mask_indices,
            sink_span=example.sink_span,
            thinking_span=example.thinking_span,
            metadata=metadata,
        )

    span_start_char, span_end_char = span
    gen_ids = tokenizer(tgt, add_special_tokens=False, return_offsets_mapping=True)
    sink_tokens: List[int] = []
    for idx, (s, e) in enumerate(gen_ids["offset_mapping"]):
        # include tokens that overlap the answer span
        if s < span_end_char and e > span_start_char:
            sink_tokens.append(idx)
    if not sink_tokens:
        return CachedExample(
            prompt=example.prompt,
            target=example.target,
            indices_to_explain=example.indices_to_explain,
            attr_mask_indices=example.attr_mask_indices,
            sink_span=example.sink_span,
            thinking_span=example.thinking_span,
            metadata=metadata,
        )

    sink_span = [min(sink_tokens), max(sink_tokens)]
    thinking_end = max(0, sink_span[0] - 1)
    thinking_span = [0, thinking_end] if thinking_end >= 0 else sink_span

    return CachedExample(
        prompt=example.prompt,
        target=example.target,
        indices_to_explain=example.indices_to_explain,
        attr_mask_indices=example.attr_mask_indices,
        sink_span=example.sink_span or sink_span,
        thinking_span=example.thinking_span or thinking_span,
        metadata=metadata,
    )


def attach_spans_from_boxed(example: CachedExample, tokenizer) -> CachedExample:
    """Backward-compatible wrapper that first looks for \\box{} then falls back to answer text."""
    tgt = example.target
    match = _find_box_span(tgt) if tgt else None
    boxed_answer = match[2] if match else None
    return attach_spans_from_answer(example, tokenizer, boxed_answer)


def ruler_gold_prompt_token_indices(example: CachedExample, tokenizer) -> List[int]:
    """Return token indices (prompt-side) that overlap RULER `needle_spans` in metadata.

    The returned indices are with respect to `tokenizer(" " + example.prompt, add_special_tokens=False)`,
    matching the attribution pipeline's leading-space convention.
    """
    needle_spans = (example.metadata or {}).get("needle_spans") or []
    if not isinstance(needle_spans, list) or not needle_spans:
        return []

    prompt_text = " " + (example.prompt or "")
    enc = tokenizer(prompt_text, add_special_tokens=False, return_offsets_mapping=True)
    offsets = enc.get("offset_mapping")
    if offsets is None:
        raise ValueError("Tokenizer does not provide offset_mapping; cannot map needle_spans to tokens.")

    spans: List[tuple[int, int]] = []
    for item in needle_spans:
        if not isinstance(item, dict):
            continue
        raw = item.get("span")
        if not (isinstance(raw, list) and len(raw) == 2):
            continue
        try:
            start = int(raw[0]) + 1  # shift for leading space in prompt_text
            end = int(raw[1]) + 1
        except Exception:
            continue
        if end > start:
            spans.append((start, end))

    if not spans:
        return []

    gold: set[int] = set()
    for tok_idx, off in enumerate(offsets):
        if off is None:
            continue
        try:
            s, e = int(off[0]), int(off[1])
        except Exception:
            continue
        if e <= s:
            continue
        for span_start, span_end in spans:
            if s < span_end and e > span_start:
                gold.add(tok_idx)
                break

    return sorted(gold)


class DatasetLoader:
    """Thin loader that resolves and samples datasets for exp2."""

    def __init__(self, seed: int = 42, data_root: Path | str = Path("exp/exp2/data")) -> None:
        self.seed = seed
        self.data_root = Path(data_root)

    def _sample(self, items: List[CachedExample], sample: Optional[int]) -> List[CachedExample]:
        if sample is not None and sample < len(items):
            rnd = random.Random(self.seed)
            rnd.shuffle(items)
            items = items[:sample]
        return items

    def _cached_path(self, name: str) -> Optional[Path]:
        path = self.data_root / f"{name}.jsonl"
        return path if path.exists() else None

    def load(self, name: str, sample: Optional[int] = None) -> List[CachedExample]:
        # 1) Prefer prepared cache under exp/exp2/data
        cached_path = self._cached_path(name)
        if cached_path:
            return self._sample(load_cached(cached_path), sample)

        return self.load_raw(name, sample=sample)

    def load_raw(self, name: str, sample: Optional[int] = None) -> List[CachedExample]:
        def _looks_like_json_array(path: Path) -> bool:
            try:
                with path.open("r", encoding="utf-8") as f:
                    while True:
                        ch = f.read(1)
                        if not ch:
                            return False
                        if ch.isspace():
                            continue
                        return ch == "["
            except OSError:
                return False

        # MoreHopQA
        if name == "morehopqa":
            ex = load_morehopqa()
            for item in ex:
                if "answer" in item.metadata:
                    item.metadata.setdefault("reference_answer", item.metadata["answer"])
            return self._sample(ex, sample)

        # Allow passing the raw MoreHopQA JSON path directly.
        p = Path(name)
        if p.exists() and _looks_like_json_array(p):
            ex = load_morehopqa(p)
            for item in ex:
                if "answer" in item.metadata:
                    item.metadata.setdefault("reference_answer", item.metadata["answer"])
            return self._sample(ex, sample)

        # RULER / HotpotQA / niah / vt (all go through RulerAttributionDataset)
        resolved = dataset_from_name(name)
        if resolved is None:
            raise FileNotFoundError(f"Could not resolve dataset {name}")
        ex = load_ruler(resolved)
        for item in ex:
            outputs = item.metadata.get("outputs") or []
            if outputs:
                item.metadata.setdefault("reference_answer", ", ".join(outputs))
            if item.target and "reference_answer" not in item.metadata:
                item.metadata["reference_answer"] = item.target
        return self._sample(ex, sample)