canvas_sweep.py

# /// script
# requires-python = ">=3.11"
# dependencies = [
#   "transformers>=5.11,<6",
#   "torch>=2.8,<2.10",  # a100-large job image driver is CUDA 12.9; torch 2.10+ ships cu13 wheels
#   "torchvision",  # Gemma4Processor imports it even for text-only use
#   "accelerate",
#   "huggingface-hub",
#   "pillow",
#   "pyzipper",
#   "requests",
# ]
# ///
"""Canvas-rescue sweep: can an OCR-seeded canvas be made to EDIT instead of copy?

Background (v1, experiments/2026-06-10_v1-bln600-text): seeding DiffusionGemma's
denoising canvas with the OCR text collapses to copy-through — the entropy-bound
sampler accepts the real text wholesale in 2-5 steps (fix rate 0.6%, CER worse
than doing nothing). Three rescue axes, the first two suggested by João Gante
(DeepMind):
  1. t_max        — raise the initial temperature of the schedule (default 0.8)
                    to "wiggle the model out of the initial OCR estimate"
  2. entropy_bound — tighten the EntropyBoundSampler bound (default 0.1) so
                    fewer tokens lock early
  3. noise_p      — SDEdit-style: corrupt a fraction p of the seeded canvas
                    tokens to uniform random, putting it at an intermediate
                    noise level on the training distribution

GENERATION ONLY — metrics computed offline by `metrics.py sweep`.

Stages (HF Jobs, a100-large):
    --stage smoke    3 passages x 3 cells (anchor / default-seeded / extreme):
                     verifies the config overrides actually bite before
                     burning a grid on silently-ignored knobs
    --stage screen   20 passages x (27-cell factorial + random-canvas anchor)
    --stage confirm  75 passages x top cells (--cells p0.25_t1.2_e0.03,...) x 3 sampler seeds
"""

import argparse
import copy
import difflib
import gc
import json
import random
import re
import time
from pathlib import Path

import requests

DG_MODEL = "google/diffusiongemma-26B-A4B-it"
TOK_MODEL = "google/gemma-4-E4B-it"  # same tokenizer family; keeps v1 passage selection identical

BLN600_API = "https://api.figshare.com/v2/articles/25439023"
BLN600_ZIP_PASSWORD = b"BLN600"
MAX_PASSAGE_TOKENS = 220

P_VALUES = [0.0, 0.25, 0.5]
T_MAX_VALUES = [0.8, 1.2, 1.6]
ENTROPY_BOUNDS = [0.1, 0.03, 0.01]

PROMPT_TEMPLATE = """\
Correct the OCR errors in the following text from a 19th-century English newspaper.
Fix only recognition errors (wrong, missing, or extra characters). Do not modernise \
spelling, do not rephrase, and do not add or remove content. Preserve the original \
punctuation unless it is clearly an OCR error.
Output only the corrected text, with no commentary or preamble.

OCR text:
{ocr}"""

STOP_MARKERS = ("<turn|>", "<eos>", "<end_of_turn>", "<pad>")


# ------------------------------------------------- shared with benchmark.py
# (duplicated: HF Jobs uploads only this file, so the script must be self-contained)


def extract_answer(raw: str) -> tuple[str, str]:
    stops = [i for m in STOP_MARKERS if (i := raw.find(m)) != -1]
    if stops:
        raw = raw[: min(stops)]
    thought = ""
    if "<channel|>" in raw:
        head, _, raw = raw.rpartition("<channel|>")
        m = re.search(r"<\|channel>thought(.*)$", head, flags=re.DOTALL)
        if m:
            thought = m.group(1).strip()
    return raw.strip(), thought


def clean_output(text: str) -> str:
    return re.sub(r"^\s*corrected text:?\s*", "", text.strip(), flags=re.IGNORECASE)


def count_generated_tokens(generated_ids, tokenizer) -> int:
    ids = generated_ids.tolist()
    stop_ids = {tokenizer.eos_token_id, tokenizer.pad_token_id}
    count = 0
    for tid in ids:
        if tid in stop_ids:
            break
        count += 1
    return count


def _download(url: str, dest: Path) -> Path:
    dest.parent.mkdir(parents=True, exist_ok=True)
    if dest.exists():
        return dest
    print(f"downloading {url} -> {dest}")
    with requests.get(url, stream=True, timeout=600) as r:
        r.raise_for_status()
        with dest.open("wb") as f:
            for chunk in r.iter_content(chunk_size=1 << 20):
                f.write(chunk)
    return dest


def download_bln600(workdir: Path) -> list[dict]:
    import pyzipper

    meta = requests.get(BLN600_API, timeout=60).json()
    zips = [f for f in meta["files"] if f["name"].lower().endswith(".zip")]
    zip_path = _download(zips[0]["download_url"], workdir / zips[0]["name"])
    extract_dir = workdir / "bln600"
    if not extract_dir.exists():
        with pyzipper.AESZipFile(zip_path) as zf:
            zf.setpassword(BLN600_ZIP_PASSWORD)
            zf.extractall(extract_dir)
    ocr_files = {p.stem: p for p in extract_dir.rglob("*.txt") if "ocr" in str(p.parent).lower()}
    gold_files = {
        p.stem: p
        for p in extract_dir.rglob("*.txt")
        if "ground" in str(p.parent).lower() or "gold" in str(p.parent).lower()
    }
    common = sorted(set(ocr_files) & set(gold_files))
    print(f"BLN600: {len(common)} aligned pairs")
    return [
        {
            "id": f"bln600/{stem}",
            "ocr_input": ocr_files[stem].read_text(errors="replace"),
            "gold": gold_files[stem].read_text(errors="replace"),
        }
        for stem in common
    ]


def trim_pair(ocr: str, gold: str, n_tokens, max_tokens: int) -> tuple[str, str] | None:
    if n_tokens(ocr) <= max_tokens and n_tokens(gold) <= max_tokens:
        return ocr, gold
    sm = difflib.SequenceMatcher(None, ocr, gold, autojunk=False)
    candidates = [
        (i1 + m.start(), j1 + m.start())
        for op, i1, i2, j1, _j2 in sm.get_opcodes()
        if op == "equal"
        for m in re.finditer(r"\s", ocr[i1:i2])
    ]
    if not candidates:
        return None

    def fits(idx: int) -> bool:
        i_cut, j_cut = candidates[idx]
        return n_tokens(ocr[:i_cut]) <= max_tokens and n_tokens(gold[:j_cut]) <= max_tokens

    if not fits(0):
        return None
    lo, hi = 0, len(candidates) - 1
    while lo < hi:
        mid = (lo + hi + 1) // 2
        if fits(mid):
            lo = mid
        else:
            hi = mid - 1
    i_cut, j_cut = candidates[lo]
    return ocr[:i_cut].rstrip(), gold[:j_cut].rstrip()


def sample_passages(passages: list[dict], n: int, seed: int) -> list[dict]:
    from transformers import AutoTokenizer

    tok = AutoTokenizer.from_pretrained(TOK_MODEL)

    def n_tokens(text: str) -> int:
        return len(tok(text)["input_ids"])

    chosen = random.Random(seed).sample(passages, len(passages))
    out: list[dict] = []
    for p in chosen:
        if len(out) >= n:
            break
        trimmed = trim_pair(p["ocr_input"], p["gold"], n_tokens, MAX_PASSAGE_TOKENS)
        if trimmed is None or len(trimmed[1]) < 200:
            continue
        out.append({"id": p["id"], "ocr_input": trimmed[0], "gold": trimmed[1]})
    print(f"sampled {len(out)} passages (seed {seed})")
    return out


# ------------------------------------------------- sweep


def make_conditions(stage: str, cells_arg: str | None) -> list[dict]:
    """A condition: canvas ('random' or 'ocr'), noise_p, t_max, entropy_bound, sampler_seed."""
    anchor = {"key": "anchor_random_canvas", "canvas": "random", "noise_p": None,
              "t_max": 0.8, "entropy_bound": 0.1, "sampler_seed": 0}
    grid = [
        {
            "key": f"p{p}_t{t}_e{e}",
            "canvas": "ocr", "noise_p": p, "t_max": t, "entropy_bound": e, "sampler_seed": 0,
        }
        for p in P_VALUES
        for t in T_MAX_VALUES
        for e in ENTROPY_BOUNDS
    ]
    if stage == "smoke":
        by_key = {c["key"]: c for c in grid}
        return [anchor, by_key["p0.0_t0.8_e0.1"], by_key["p0.5_t1.6_e0.01"]]
    if stage == "screen":
        return [anchor] + grid
    # confirm: named cells x 3 sampler seeds
    if not cells_arg:
        raise SystemExit("--stage confirm requires --cells key1,key2,...")
    by_key = {c["key"]: c for c in grid}
    out = []
    for key in cells_arg.split(","):
        for seed in (0, 1, 2):
            c = dict(by_key[key.strip()])
            c["sampler_seed"] = seed
            c["key"] = f"{key.strip()}_s{seed}"
            out.append(c)
    return out


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--stage", choices=["smoke", "screen", "confirm"], default="smoke")
    parser.add_argument("--cells", default=None, help="confirm stage: comma-separated cell keys")
    parser.add_argument("--seed", type=int, default=42, help="passage sampling seed (v1 = 42)")
    parser.add_argument("--out-repo", default=None)
    parser.add_argument("--workdir", type=Path, default=Path("data"))
    args = parser.parse_args()

    import torch
    import transformers
    from transformers import AutoProcessor, DiffusionGemmaForBlockDiffusion, TextDiffusionStreamer

    print(f"transformers {transformers.__version__}, torch {torch.__version__}, "
          f"cuda: {torch.cuda.get_device_name(0)}")

    n_passages = {"smoke": 3, "screen": 20, "confirm": 75}[args.stage]
    passages = sample_passages(download_bln600(args.workdir), n_passages, args.seed)
    conditions = make_conditions(args.stage, args.cells)
    print(f"stage={args.stage}: {len(passages)} passages x {len(conditions)} conditions")

    class StepCountingStreamer(TextDiffusionStreamer):
        def __init__(self, tokenizer):
            super().__init__(tokenizer=tokenizer)
            self.n_steps = 0

        def put_draft(self, value, **kwargs):
            self.n_steps += 1

        def put(self, value):
            pass

        def end(self):
            pass

    print(f"loading {DG_MODEL} ...")
    processor = AutoProcessor.from_pretrained(DG_MODEL)
    model = DiffusionGemmaForBlockDiffusion.from_pretrained(DG_MODEL, dtype="auto", device_map="auto")
    tokenizer = processor.tokenizer
    base_config = copy.deepcopy(model.generation_config)
    canvas_length = getattr(base_config, "canvas_length", None) or 256
    vocab = model.config.text_config.vocab_size

    def build_canvas(ocr_text: str, noise_p: float, rng: torch.Generator) -> torch.Tensor:
        ids = tokenizer(ocr_text, add_special_tokens=False)["input_ids"][:canvas_length]
        ids = torch.tensor(ids, dtype=torch.long)
        if noise_p > 0:
            corrupt = torch.rand(len(ids), generator=rng) < noise_p
            ids[corrupt] = torch.randint(vocab, (int(corrupt.sum()),), generator=rng)
        pad = torch.randint(vocab, (canvas_length - len(ids),), generator=rng)
        return torch.cat([ids, pad]).unsqueeze(0)

    def generate(ocr_text: str, cond: dict, canvas_rng: torch.Generator) -> dict:
        # mutate the live generation_config: guaranteed to be read, no reliance
        # on generate() kwarg plumbing for these fields
        gen_config = copy.deepcopy(base_config)
        gen_config.t_max = cond["t_max"]
        gen_config.sampler_config.entropy_bound = cond["entropy_bound"]
        model.generation_config = gen_config

        message = [{"role": "user", "content": PROMPT_TEMPLATE.format(ocr=ocr_text)}]
        inputs = processor.apply_chat_template(
            message, tokenize=True, add_generation_prompt=True, return_dict=True, return_tensors="pt"
        ).to(model.device)
        input_len = inputs["input_ids"].shape[-1]

        streamer = StepCountingStreamer(tokenizer)
        gen_kwargs: dict = {"max_new_tokens": 256, "streamer": streamer}
        if cond["canvas"] == "ocr":
            gen_kwargs["decoder_input_ids"] = build_canvas(
                ocr_text, cond["noise_p"], canvas_rng
            ).to(model.device)

        torch.manual_seed(cond["sampler_seed"])  # sampler stochasticity, reproducible
        torch.cuda.synchronize()
        t0 = time.perf_counter()
        output = model.generate(**inputs, **gen_kwargs)
        torch.cuda.synchronize()
        seconds = time.perf_counter() - t0

        seq = output.sequences if hasattr(output, "sequences") else output
        generated = seq[0][input_len:] if seq.shape[-1] > input_len else seq[0]
        raw = tokenizer.decode(generated, skip_special_tokens=False)
        answer, thought = extract_answer(raw)
        return {
            "text": clean_output(answer),
            "_raw": raw,
            "seconds": round(seconds, 3),
            "tokens_generated": count_generated_tokens(generated, tokenizer),
            "denoising_steps": streamer.n_steps,
            "thought_chars": len(thought),
        }

    # verify the overrides bite: print the effective config per condition
    print("warmup (uncounted) ...")
    warm_rng = torch.Generator().manual_seed(999)
    generate(passages[0]["ocr_input"], conditions[0], warm_rng)

    meta = {
        "date": time.strftime("%Y-%m-%d"),
        "stage": args.stage,
        "n_passages": len(passages),
        "passage_seed": args.seed,
        "prompt": PROMPT_TEMPLATE,
        "grid": {"noise_p": P_VALUES, "t_max": T_MAX_VALUES, "entropy_bound": ENTROPY_BOUNDS},
        "defaults": {"t_max": base_config.t_max,
                     "entropy_bound": base_config.sampler_config.entropy_bound,
                     "confidence_threshold": getattr(base_config, "confidence_threshold", None),
                     "max_denoising_steps": getattr(base_config, "max_denoising_steps", None)},
        "transformers": transformers.__version__,
        "torch": torch.__version__,
        "gpu": torch.cuda.get_device_name(0),
    }

    out_path = Path(f"raw_outputs_canvas_sweep_{args.stage}.jsonl")
    records = []
    t_sweep = time.perf_counter()
    for ci, cond in enumerate(conditions):
        print(f"[{ci + 1}/{len(conditions)}] {cond['key']}: t_max={cond['t_max']}, "
              f"entropy_bound={cond['entropy_bound']}, canvas={cond['canvas']}, p={cond['noise_p']}")
        # deterministic noise per (condition, passage)
        for pi, p in enumerate(passages):
            canvas_rng = torch.Generator().manual_seed(ci * 10_000 + pi)
            out = generate(p["ocr_input"], cond, canvas_rng)
            raw = out.pop("_raw")
            records.append({"id": p["id"], "ocr_input": p["ocr_input"], "gold": p["gold"],
                            "condition": {k: cond[k] for k in
                                          ("key", "canvas", "noise_p", "t_max", "entropy_bound",
                                           "sampler_seed")},
                            "output": out})
            print(f"    {p['id']}: {out['seconds']}s, {out['denoising_steps']} steps, "
                  f"{out['tokens_generated']} tok")
            if args.stage == "smoke":
                print(f"    OCR: {p['ocr_input'][:160]}")
                print(f"    OUT: {out['text'][:160]}")
                print(f"    RAW: {raw[:200]}")
    print(f"sweep wall-clock: {time.perf_counter() - t_sweep:.0f}s")

    with out_path.open("w") as f:
        for i, r in enumerate(records):
            if i == 0:
                r = {**r, "meta": meta}
            f.write(json.dumps(r) + "\n")
    print(f"wrote {out_path} ({len(records)} records)")

    if args.out_repo:
        from huggingface_hub import HfApi

        api = HfApi()
        api.create_repo(args.out_repo, repo_type="dataset", private=True, exist_ok=True)
        api.upload_file(path_or_fileobj=out_path, path_in_repo=out_path.name,
                        repo_id=args.out_repo, repo_type="dataset")
        print(f"uploaded to https://huggingface.co/datasets/{args.out_repo} (private)")

    model = None  # noqa: F841 — drop the closure-captured ref so the GPU frees
    gc.collect()
    torch.cuda.empty_cache()


if __name__ == "__main__":
    main()