code/local_train/pack.py

from __future__ import annotations

import json
import random
import re
from collections import defaultdict
from dataclasses import asdict, dataclass, field
from pathlib import Path
from typing import Any, Iterable

import pyarrow.parquet as pq

from albedo_eval_service.remote.dataset import EvalSample, apply_submit_protocol, format_messages
from albedo_eval_service.shared.observation_format import first_bash_block
from albedo_eval_service.shared.submit_protocol import ANY_MARKER_RE, TAILS

from .constants import (
    DEFAULT_DATA_ROOT,
    DEFAULT_PACK_DIR,
    KEEP_ORIGINAL_RATIO,
    MAX_PREFIX_CHARS,
    NON_PYTHON_FRACTION,
    SOURCES,
    TOKENIZER_DIR,
)
from .think import wrap_completion

_EDIT_RE = re.compile(
    r"sed\s+-i|tee\s+[\w./-]|cat\s*>|str_replace|git apply|patch\s+-p|applypatch|"
    r"cp\s+[\w./-]|mv\s+[\w./-]|(?<![-\d&])>>?\s*(?!/dev/)[\w.][\w./-]*"
)
_DUMMY_RE = re.compile(
    r"your_command_here|cat\s+<<'EOF'\s*>\s*newfile\.py|sed\s+-i\s+.*\bfilename\.py\b"
)
_PATH_RE = re.compile(r"(?:/|\./|[\w.-]+/)[\w./-]+\.[A-Za-z0-9]{1,8}")


@dataclass
class PackedExample:
    sample_id: str
    prompt: str
    completion: str
    source: str
    phase: str
    kind: str
    family: str
    language: str
    repo: str
    submit_command: str
    submit_marker: str
    rewrite_mode: str
    gold_paths: list[str] = field(default_factory=list)
    salt: str = ""

    def as_dict(self) -> dict[str, Any]:
        return asdict(self)


def is_gold_submit(text: str) -> bool:
    command = first_bash_block(text) or ""
    hay = command or text or ""
    return bool(ANY_MARKER_RE.search(hay))


def is_edit_command(command: str) -> bool:
    return bool(command and _EDIT_RE.search(command) and not _DUMMY_RE.search(command))


def gold_paths(text: str) -> list[str]:
    command = first_bash_block(text) or text or ""
    seen: list[str] = []
    for match in _PATH_RE.finditer(command):
        path = match.group(0)
        if path not in seen and "filename.py" not in path and "newfile.py" not in path:
            seen.append(path)
    return seen


def family_of(instance_id: str, given: str = "") -> str:
    if given:
        return given
    if "." not in (instance_id or ""):
        return "pr"
    tail = instance_id.rsplit(".", 1)[-1]
    for prefix, family in (("pr_", "pr"), ("lm_", "lm"), ("combine", "combine")):
        if tail.startswith(prefix):
            return family
    return "mechanical"


def phase_for(turn_idx: int, first_edit: int) -> str:
    if first_edit <= 0:
        return "cold" if turn_idx in (1, 2) else "explore"
    if turn_idx == first_edit:
        return "at_edit"
    if turn_idx == max(1, first_edit - 2):
        return "pre_edit"
    if turn_idx in (1, 2) and turn_idx < max(1, first_edit - 2):
        return "cold"
    if turn_idx >= first_edit:
        return "post_edit"
    return "explore"


def candidate_turns(n_assistant: int, first_edit: int, golds: list[str]) -> list[tuple[int, str]]:
    """Official cold/pre_edit/at_edit cuts plus the first-edit turn and later submits."""
    out: list[tuple[int, str]] = []
    if n_assistant < 3:
        return out
    for turn_idx in (1, 2):
        if turn_idx < n_assistant:
            out.append((turn_idx, "cold"))
    if first_edit > 0:
        pre = max(1, first_edit - 2)
        if pre < n_assistant:
            out.append((pre, "pre_edit"))
        if first_edit < n_assistant:
            out.append((first_edit, "at_edit"))
        edit_idx = first_edit - 1
        if 0 <= edit_idx < n_assistant:
            out.append((edit_idx, "at_edit"))
    last_submit = next(
        (turn_idx for turn_idx in range(len(golds) - 1, 2, -1) if is_gold_submit(golds[turn_idx])),
        None,
    )
    if last_submit is not None:
        out.append((last_submit, "post_edit"))
    seen: set[int] = set()
    unique: list[tuple[int, str]] = []
    for turn_idx, phase in out:
        if turn_idx in seen or turn_idx < 0 or turn_idx >= n_assistant:
            continue
        seen.add(turn_idx)
        unique.append((turn_idx, phase_for(turn_idx, first_edit) if first_edit else phase))
    return unique


def pack(
    *,
    dataset_root: Path = DEFAULT_DATA_ROOT,
    out_dir: Path = DEFAULT_PACK_DIR,
    max_examples: int = 20_000,
    seed: str = "sft-pack",
    n_salts: int = 2,
    tokenizer_path: Path | None = None,
    submit_frac: float = 0.20,
    edit_frac: float = 0.35,
    expand_submit_salts: bool = False,
) -> Path:
    dataset_root = Path(dataset_root)
    out_dir = Path(out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    tokenizer = str(tokenizer_path or TOKENIZER_DIR)
    salts = [f"{seed}-{i}" for i in range(max(1, n_salts))]
    raw = list(_iter_raw(dataset_root, max_raw=max(max_examples * 8, 256), seed=seed))
    selected = _select(
        raw,
        max_examples=max_examples,
        seed=seed,
        submit_frac=submit_frac,
        edit_frac=edit_frac,
    )
    examples = _materialize(
        selected,
        salts=salts,
        tokenizer_path=tokenizer,
        expand_submit_salts=expand_submit_salts,
    )
    pack_path = out_dir / f"sft-{max_examples}-{seed}.jsonl"
    with pack_path.open("w") as handle:
        for example in examples:
            handle.write(json.dumps(example.as_dict(), ensure_ascii=False) + "\n")
    summary = _summary(examples)
    (out_dir / f"sft-{max_examples}-{seed}.meta.json").write_text(
        json.dumps(summary, indent=2) + "\n"
    )
    print(json.dumps(summary, indent=2), flush=True)
    print(f"pack: {pack_path}", flush=True)
    return pack_path


@dataclass
class _Raw:
    source: str
    shard: str
    row: int
    turn_idx: int
    phase: str
    kind: str
    family: str
    language: str
    repo: str
    instance_id: str
    messages: list[dict[str, str]]
    gold: str
    first_edit: int


def _iter_raw(dataset_root: Path, *, max_raw: int, seed: str) -> Iterable[_Raw]:
    rng = random.Random(seed)
    shards_by_source = {
        source: sorted((dataset_root / source / "data").glob("train-*.parquet"))
        for source in SOURCES
        if (dataset_root / source / "data").is_dir()
    }
    if not any(shards_by_source.values()):
        raise FileNotFoundError(
            f"no official shards under {dataset_root}/<source>/data/train-*.parquet"
        )
    produced = 0
    cursor = {source: 0 for source in shards_by_source}
    while produced < max_raw and any(
        cursor[source] < len(shards) for source, shards in shards_by_source.items()
    ):
        for source, shards in shards_by_source.items():
            if produced >= max_raw or cursor[source] >= len(shards):
                continue
            shard = shards[cursor[source]]
            cursor[source] += 1
            rel = f"{source}/data/{shard.name}"
            took = 0
            for raw in _rows_from_shard(source, rel, shard):
                yield raw
                produced += 1
                took += 1
                if produced >= max_raw or took >= 64:
                    break
            if produced >= max_raw:
                return


def _rows_from_shard(source: str, rel: str, shard: Path) -> Iterable[_Raw]:
    schema = pq.read_schema(shard)
    columns = [
        name
        for name in (
            "messages",
            "turns",
            "conversation",
            "instance_id",
            "first_edit",
            "family",
            "repo",
            "language",
        )
        if name in schema.names
    ]
    if not columns:
        return
    parquet = pq.ParquetFile(shard)
    row_idx = 0
    for batch in parquet.iter_batches(batch_size=256, columns=columns):
        for row in batch.to_pylist():
            current = row_idx
            row_idx += 1
            turns = _as_turns(
                row.get("messages") or row.get("turns") or row.get("conversation")
            )
            assistant = [i for i, turn in enumerate(turns) if _role(turn) == "assistant"]
            if len(assistant) < 3:
                continue
            golds = [_content(turns[i]) for i in assistant]
            given_edit = row.get("first_edit")
            first_edit = (
                int(given_edit)
                if given_edit is not None
                else _first_edit(golds)
            )
            instance_id = str(row.get("instance_id") or "")
            language = str(row.get("language") or ("rust" if source == "mini-coder-rs" else "python"))
            family = family_of(instance_id, str(row.get("family") or ""))
            repo = str(row.get("repo") or (instance_id.split(".")[0] if instance_id else source))
            for turn_idx, phase in candidate_turns(len(assistant), first_edit, golds):
                gold = golds[turn_idx]
                command = first_bash_block(gold)
                if not command or _DUMMY_RE.search(command):
                    continue
                kind = (
                    "submit"
                    if is_gold_submit(gold)
                    else "edit"
                    if is_edit_command(command)
                    else "explore"
                )
                if kind == "submit" and turn_idx <= 2:
                    continue
                prefix_turns = turns[: assistant[turn_idx]]
                messages = [
                    {"role": _chat_role(_role(turn)), "content": _content(turn)}
                    for turn in prefix_turns
                    if _content(turn)
                ]
                prefix_chars = sum(len(m["content"]) for m in messages)
                if prefix_chars > MAX_PREFIX_CHARS or not messages:
                    continue
                yield _Raw(
                    source=source,
                    shard=rel,
                    row=current,
                    turn_idx=turn_idx,
                    phase=phase,
                    kind=kind,
                    family=family,
                    language=language,
                    repo=repo,
                    instance_id=instance_id,
                    messages=messages,
                    gold=gold,
                    first_edit=first_edit,
                )


def _first_edit(golds: list[str]) -> int:
    for index, gold in enumerate(golds, start=1):
        if is_edit_command(first_bash_block(gold)):
            return index
    return 0


def _select(
    raw: list[_Raw],
    *,
    max_examples: int,
    seed: str,
    submit_frac: float = 0.20,
    edit_frac: float = 0.35,
) -> list[_Raw]:
    rng = random.Random(f"{seed}:select")
    buckets: dict[tuple[str, str, str, str], list[_Raw]] = defaultdict(list)
    for item in raw:
        lang = "other" if item.language != "python" else "python"
        buckets[(item.source, item.phase, item.kind, lang)].append(item)
    for items in buckets.values():
        rng.shuffle(items)

    want_submit = max(1, int(max_examples * submit_frac))
    want_edit = max(1, int(max_examples * edit_frac))
    want_other_lang = max(1, int(max_examples * NON_PYTHON_FRACTION))
    picked: list[_Raw] = []
    used: set[tuple[str, int, int]] = set()

    def take(predicate, limit: int) -> None:
        leftover = limit
        keys = sorted(buckets)
        while leftover > 0:
            progressed = False
            for key in keys:
                if leftover <= 0:
                    break
                items = buckets[key]
                kept: list[_Raw] = []
                found = None
                while items:
                    item = items.pop()
                    ident = (item.shard, item.row, item.turn_idx)
                    if ident in used or not predicate(item):
                        kept.append(item)
                        continue
                    found = item
                    break
                items.extend(kept)
                if found is None:
                    continue
                used.add((found.shard, found.row, found.turn_idx))
                picked.append(found)
                leftover -= 1
                progressed = True
            if not progressed:
                break

    take(lambda item: item.kind == "submit", want_submit)
    take(lambda item: item.kind == "edit" and item.phase == "at_edit", want_edit)
    take(lambda item: item.kind == "edit", want_edit - sum(1 for i in picked if i.kind == "edit"))
    take(lambda item: item.language != "python", want_other_lang)
    take(lambda _item: True, max_examples - len(picked))
    rng.shuffle(picked)
    return picked[:max_examples]


def _materialize(
    raw: list[_Raw],
    *,
    salts: list[str],
    tokenizer_path: str,
    expand_submit_salts: bool = False,
) -> list[PackedExample]:
    by_salt: dict[str, list[tuple[_Raw, EvalSample]]] = defaultdict(list)
    for index, item in enumerate(raw):
        chosen = salts if (expand_submit_salts and item.kind == "submit") else [salts[index % len(salts)]]
        for salt in chosen:
            sample_id = f"{item.shard}:{item.row}:{item.turn_idx}"
            sample = EvalSample(sample_id=sample_id, prompt="", messages=list(item.messages))
            by_salt[salt].append((item, sample))

    examples: list[PackedExample] = []
    for salt, group in by_salt.items():
        rewritten = apply_submit_protocol(
            [sample for _, sample in group],
            salt=salt,
            keep_original_ratio=KEEP_ORIGINAL_RATIO,
            tokenizer_path=tokenizer_path,
            enable_thinking=True,
        )
        for (item, _), sample in zip(group, rewritten, strict=True):
            bash = sample.submit_command if item.kind == "submit" else None
            if item.kind == "submit" and not sample.submit_command:
                continue
            completion = wrap_completion(item.gold, bash)
            if completion is None:
                continue
            if item.kind == "submit" and not ANY_MARKER_RE.search(completion):
                continue
            prompt = sample.prompt or format_messages(
                sample.messages or item.messages,
                tokenizer_path=tokenizer_path,
                enable_thinking=True,
            )
            if item.phase == "cold" and item.kind == "submit":
                continue
            extra: list[PackedExample] = []
            if (
                item.kind == "submit"
                and "cat patch.txt" in (sample.submit_command or "")
            ):
                extra.extend(_patch_prep_example(item, sample, tokenizer_path, salt))
            examples.append(
                PackedExample(
                    sample_id=sample.sample_id,
                    prompt=prompt,
                    completion=completion,
                    source=item.source,
                    phase=item.phase,
                    kind=item.kind,
                    family=item.family,
                    language=item.language,
                    repo=item.repo,
                    submit_command=sample.submit_command,
                    submit_marker=sample.submit_marker,
                    rewrite_mode=sample.rewrite_mode,
                    gold_paths=gold_paths(item.gold),
                    salt=salt,
                )
            )
            examples.extend(extra)
    return examples


def _patch_prep_example(
    item: _Raw, sample: EvalSample, tokenizer_path: str, salt: str
) -> list[PackedExample]:
    """If gold already built a patch, keep that as its own prior turn (protocol: separate commands)."""
    prev = None
    for message in reversed(item.messages):
        if message.get("role") == "assistant":
            prev = message.get("content") or ""
            break
    if not prev:
        return []
    command = first_bash_block(prev)
    if not command or "patch.txt" not in command:
        if command and command.startswith("git diff") and ">" not in command:
            command = f"{command} > patch.txt"
        else:
            return []
    completion = wrap_completion(prev, command)
    if completion is None:
        return []
    return [
        PackedExample(
            sample_id=f"{sample.sample_id}:patch-prep",
            prompt=sample.prompt,
            completion=completion,
            source=item.source,
            phase="post_edit",
            kind="edit",
            family=item.family,
            language=item.language,
            repo=item.repo,
            submit_command=sample.submit_command,
            submit_marker=sample.submit_marker,
            rewrite_mode=sample.rewrite_mode,
            gold_paths=gold_paths(prev),
            salt=salt,
        )
    ]


def _summary(examples: list[PackedExample]) -> dict[str, Any]:
    def count(field: str) -> dict[str, int]:
        out: dict[str, int] = defaultdict(int)
        for example in examples:
            out[str(getattr(example, field))] += 1
        return dict(sorted(out.items()))

    return {
        "n": len(examples),
        "source": count("source"),
        "phase": count("phase"),
        "kind": count("kind"),
        "family": count("family"),
        "language": count("language"),
        "rewrite_mode": count("rewrite_mode"),
        "markers": count("submit_marker"),
        "tails": _tail_counts(examples),
    }


def _tail_counts(examples: list[PackedExample]) -> dict[str, int]:
    out: dict[str, int] = defaultdict(int)
    for example in examples:
        command = example.submit_command or ""
        if "cat patch.txt" in command:
            out["patchtxt"] += 1
        elif "git diff --cached" in command:
            out["gitdiff"] += 1
        elif command:
            out["bare"] += 1
        else:
            out["none"] += 1
    out["known_tails"] = len(TAILS)
    return dict(out)


def _as_turns(value: Any) -> list[Any]:
    parsed = value
    if isinstance(value, str):
        try:
            parsed = json.loads(value)
        except json.JSONDecodeError:
            return []
    if isinstance(parsed, dict):
        for key in ("messages", "turns", "conversation"):
            if isinstance(parsed.get(key), list):
                return parsed[key]
        return []
    return parsed if isinstance(parsed, list) else []


def _role(turn: Any) -> str:
    if not isinstance(turn, dict):
        return ""
    return str(turn.get("role") or turn.get("speaker") or turn.get("from") or "").lower()


def _content(turn: Any) -> str:
    if not isinstance(turn, dict):
        return str(turn or "")
    for key in ("content", "text", "value", "message"):
        value = turn.get(key)
        if value:
            return str(value)
    return ""


def _chat_role(role: str) -> str:
    if role in {"assistant", "system", "user"}:
        return role
    if role in {"human", "prompter"}:
        return "user"
    if role in {"gpt", "bot", "model"}:
        return "assistant"
    return "user"