Blame for code/tests/data/make_qwen3_bytelevel_oracle.py

"""Regenerate tests/data/qwen3_bytelevel_oracle.json.gz -- the frozen HF reference for
tests/test_tokenizer_bytelevel.py.

RUN ONLY in the isolated rollout venv (it has transformers; the align venv must never):

    CUDA_VISIBLE_DEVICES= .venv-rollout/bin/python tests/data/make_qwen3_bytelevel_oracle.py

transformers is imported HERE, as a separate process, to DUMP reference ids/decodes.
align itself never imports it (tests/test_dependency_firewall.py scans align/ and
scripts/ only; this file lives under tests/data and is data tooling, not runtime).

Corpus: all 244 minif2f + 360 MOBench formal statements, the P2-B probe completions
(first 2000 chars each), Unicode stress strings, the exact chat-templated Goedel prompts
for the first 50 minif2f problems + correction-round prompts, chat-template cases, and
id sequences for decode. Ids are stored as zlib(uint32 LE) base64 to keep the file small.
"""
import array
import base64
import gzip
import json
import os
import random
import sys
import zlib

REPO = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
SNAP = os.environ.get(
    "GOEDEL_SNAPSHOT",
    "/apps/development/rlhf/hf_cache/hub/models--Goedel-LM--Goedel-Prover-V2-8B/snapshots/"
    "dfd02e6271a58375dfbf3ece0175277cf6b6a89a",
)
DATASETS = "/apps/development/rlhf/reference-repos/Goedel-Prover-V2/dataset"
PROBE = os.path.join(REPO, ".tmp/p2/P2-B/probe_16k_seqs256.jsonl")
OUT = os.path.join(REPO, "tests/data/qwen3_bytelevel_oracle.json.gz")

import unicodedata  # noqa: E402

from transformers import AutoTokenizer  # noqa: E402

sys.path.insert(0, REPO)
from extras.vllm_rollout import lean_prompts  # noqa: E402  (stdlib-only at import)

hf = AutoTokenizer.from_pretrained(SNAP)


def pack(ids):
    return base64.b64encode(zlib.compress(array.array("I", ids).tobytes(), 9)).decode()


def jl(path):
    with open(path) as f:
        return [json.loads(x) for x in f if x.strip()]


STRESS = [
    "", " ", "\n", "a", "Hello, world!", "hello  world", "  leading", "trailing  ", "a\n\nb\n\n\nc",
    "line1\r\nline2\rline3\n", " \n ", "\t\tx\t", "x \t\n y", "a\u00a0b\u2003c\u3000d\u2028e\u2029f\x0bg\x0ch\x85i",
    "a\x1cb\x1dc\x1ed\x1fe", "zero\u200bwidth\u200djoiner\ufeffbom",
    "2024", "x=1234567890", "3.14159", "١٢٣٤", "²³ ½ ⅷ Ⅷ ①", "1,000,000", "a1b22c333", "0x1F",
    "don't can't I'm we're they've she'll he'd it's", "DON'T I'M WE'RE THEY'VE SHE'LL HE'D IT'S", "'s 't 're 've 'm 'll 'd",
    "'ſ", "'ſt", "x'ſ", "'S'T'RE", "rock'n'roll", "'quoted'", "o'clock", "''s", "'re's",
    "∀ x ∈ ℝ, ∃ y, x < y → ⟨x, y⟩", "theorem foo (f : ℝ → ℝ) (h₀ : ∀ x, f x = 5 * x + 4) : f 1 = 9 := by",
    "ℕ ℤ ℚ ℝ ℂ ∑ ∏ ∫ √ ≤ ≥ ≠ ≈ ∈ ∉ ⊂ ⊆ ∪ ∩ ¬ ∧ ∨ ↔ ⇒ λ π", "x₁ x₂ x₃ a⁻¹ x²",
    "你好，世界！这是一个测试。", "日本語のテキスト、ひらがな、カタカナ", "한국어 텍스트", "Привет, мир! Ёж", "مرحبا بالعالم", "שלום עולם",
    "😀 😃 👨‍👩‍👧‍👦 🏳️‍🌈 👍🏽 🇮🇳", "\U0001f600\U0001f601", "emoji😀attached",
    "e\u0301", "a\u0300\u0301\u0302", "\u1100\u1161\u11a8", "\u0041\u030a vs \u00c5", "\u212b \u2126 \ufb01",
    "caf\u00e9 cafe\u0301", "\u0958 \u0915\u093c", "fi fl ㎏ ㍿",
    "<|im_start|>user\nhi<|im_end|>\n<|im_start|>assistant\n", "<think>\nreason\n</think>\n\nanswer", "<|endoftext|>",
    "<|im_end|>\u0301", "<|im_end|>e\u0301<|im_start|>", "e<|im_end|>\u0301x", "<|im_start|", "<|IM_START|>", "<think", "</think>>",
    "<think><think></think></think>", "<|im_start|><|im_end|>", "<tool_call>{}</tool_call>", "<|fim_prefix|>a<|fim_suffix|>b<|fim_middle|>",
    "<|vision_pad|><|image_pad|><|video_pad|>", "<|box_start|>x<|box_end|>", "<|file_sep|><|repo_name|>",
    "\x00\x01\x02\x7f", "a\x00b", "\ue000 private use", "\ufffe \uffff", "\U0010ffff", "\ufffd",
    "\U00001c89 new in Unicode 16", "\U00010d4a garay digit", "\U0001e5d0 ol onal", "\U0000a7ce \U0000a7cf", "\U00031350 CJK ext H",
    "\u0378 unassigned", "\U000e0001 tag", "\u180e mongolian vowel sep",
    "x" * 5000, " " * 5000, "-" * 3000, "\n" * 2000, "1" * 3000, "ab " * 2000, "=" * 100 + "\n" + "=" * 100,
    "a" + " " * 40 + "b", "a" + "\n" * 40 + " b", "\t" * 50 + "x",
    "!!!???...", "...", "a,b;c:d", "(a)[b]{c}", "a_b_c", "__init__", "foo.bar(baz)", "$x^2$", "\\frac{1}{3}Bh",
    "/-- doc -/\ntheorem t : True := trivial", "```lean4\nexample : 1 + 1 = 2 := by norm_num\n```", "  have h : x = 1 := by\n    simp\n    <;> linarith\n",
    "https://example.com/a?b=c&d=e#f", "user@example.com", "C:\\Users\\x", "a\u2028b", "x\u00a0\u00a0y", " \u00a0 ",
    "it'sx", "'tis", "I'dnt", "she'lled", "'rest", "'very", "'Sx", "'REx", "we'veto", "they'rent", "i'mnot",
    "'The 'This 'Time 'Maybe 'Mean 'Does 'REAL 'Still 'Total", "He said 'The End' and 'THIS IS IT'",
    "word\u0301s", "1\u0301", "!\u0301", " \u0301", "\u0301", "\u0301\u0301a",
]


def main():
    out = {"snapshot": os.path.basename(SNAP), "hf_transformers": __import__("transformers").__version__,
           "hf_tokenizers": __import__("tokenizers").__version__, "texts": [], "id_seqs": [], "chat": []}

    texts = []  # (group, text)
    for r in jl(os.path.join(DATASETS, "minif2f.jsonl")):
        texts.append(("minif2f", r["formal_statement"]))
    for r in jl(os.path.join(DATASETS, "MOBench.jsonl")):
        texts.append(("mobench", r["formal_statement"]))
    for r in jl(PROBE):
        texts.append(("completion", r["completion_text"][:2000]))
    for s in STRESS:
        texts.append(("stress", s))

    minif = jl(os.path.join(DATASETS, "minif2f.jsonl"))
    for r in minif[:50]:
        prompt, msgs = lean_prompts.initial_prompt(r["formal_statement"], hf)
        texts.append(("chat_prompt", prompt))
        out["chat"].append({"messages": msgs, "kwargs": {"add_generation_prompt": True}, "expected": prompt})
    # correction rounds (assistant turn replayed, template rewrites it)
    probe = jl(PROBE)
    n_corr = 0
    for r, p in zip(minif[:20], probe[:20]):
        try:
            _, m0 = lean_prompts.initial_prompt(r["formal_statement"], hf)
            prompt, msgs = lean_prompts.correction_prompt(
                r["formal_statement"], m0, "<think>\nplan\n</think>\n\n" + p["completion_text"][:1500],
                "import Mathlib\ntheorem x : True := by\n  sorry", [], hf, 1)
        except Exception as e:  # noqa: BLE001
            print("correction prompt skipped:", repr(e)[:200])
            continue
        texts.append(("chat_prompt", prompt))
        out["chat"].append({"messages": msgs, "kwargs": {"add_generation_prompt": True}, "expected": prompt})
        n_corr += 1
    print("correction prompts:", n_corr)

    for g, t in texts:
        ids = hf.encode(t, add_special_tokens=False)
        ids2 = hf(t)["input_ids"]
        assert ids == ids2, ("add_special_tokens changes ids", g, t[:40])
        dec = hf.decode(ids)
        dec_skip = hf.decode(ids, skip_special_tokens=True)
        e = {"g": g, "t": t, "ids": pack(ids), "n": len(ids)}
        nfc = None
        # decode of the ids is the NFC text segment-wise; store explicit decode only when it is not t itself
        if dec != t:
            e["dec"] = dec
        if dec_skip != dec:
            e["dec_skip"] = dec_skip
        out["texts"].append(e)

    # chat-template cases
    U = lambda c: {"role": "user", "content": c}  # noqa: E731
    A = lambda c, **k: {"role": "assistant", "content": c, **k}  # noqa: E731
    S = lambda c: {"role": "system", "content": c}  # noqa: E731
    cases = [
        [U("hi")], [S("be brief"), U("hi")], [U("hi"), A("hello")], [S("s"), U("q1"), A("a1"), U("q2")],
        [U("q1"), A("<think>\nr1\n</think>\n\na1"), U("q2")],
        [U("q1"), A("<think>\nr1\n</think>\n\na1"), U("q2"), A("<think>\nr2\n</think>\n\na2")],
        [U("q1"), A("<think>\nr1\n</think>\n\na1")],
        [U("q1"), A("a1 no think")],
        [U("q1"), A("x"), A("<think>\nr\n</think>\n\ny")],
        [U("q1"), A("<think>\nr\n</think>\n\ny"), A("<think>\nr2\n</think>\n\nz")],
        [U("q1"), A("ans", reasoning_content="RC")],
        [U("q1"), A("ans", reasoning_content="\n\nRC\n\n")],
        [U("q1"), A("ans", reasoning_content="\n\nRC\n\n"), U("q2"), A("a2", reasoning_content="R2")],
        [U("q1"), A("</think>after")], [U("q1"), A("<think>a</think>b</think>c")],
        [U("q1"), A("<think>\n\n\nspaced\n\n\n</think>\n\n\n\nbody\n\n")],
        [U("q1"), A("\n\nleading newlines")], [U("q1"), A("")], [U(""), A("")],
        [S("sys1"), S("sys2"), U("u")], [A("first is assistant"), U("u")],
        [U("<tool_response>x</tool_response>"), A("a")], [U("u1"), U("u2")],
        [S("only system")], [U("é e\u0301 😀 <|im_end|> literal")],
    ]
    for msgs in cases:
        for kw in ({}, {"add_generation_prompt": True}, {"add_generation_prompt": True, "enable_thinking": False},
                   {"add_generation_prompt": True, "enable_thinking": True}, {"enable_thinking": False}):
            exp = hf.apply_chat_template(msgs, tokenize=False, **kw)
            out["chat"].append({"messages": msgs, "kwargs": kw, "expected": exp})
            texts_ids = hf.encode(exp, add_special_tokens=False)
            out["chat"][-1]["ids"] = pack(texts_ids)

    # id sequences for decode
    rng = random.Random(1234)
    V = len(hf)
    seqs = [[i] for i in sorted(hf.get_added_vocab().values())]
    seqs += [[rng.randrange(V) for _ in range(rng.randrange(1, 48))] for _ in range(150)]
    for s in ("😀", "你好", "👨‍👩‍👧‍👦", "é", "\U0001f1ee\U0001f1f3"):
        ids = hf.encode(s, add_special_tokens=False)
        seqs += [ids, ids[:1], ids[1:], ids[:-1], ids[-1:], ids + ids[:1]]
        seqs.append(ids[:1] + [151667] + ids[1:])  # added token splitting a character
    seqs += [[151643, 9707, 151645], [9707, 151667, 15, 151668, 220, 151644, 872, 198],
             [V - 1, V, V + 5, 151935, 151936, 200000, 9707], [], [9707] * 5]
    for s in seqs:
        rec = {"ids": s}
        for skip in (False, True):
            try:
                rec["dec_skip" if skip else "dec"] = hf.decode(s, skip_special_tokens=skip)
            except Exception as e:  # noqa: BLE001
                rec["dec_skip" if skip else "dec"] = None
                rec["err"] = type(e).__name__
        out["id_seqs"].append(rec)

    out["added_vocab"] = hf.get_added_vocab()
    out["special_ids"] = sorted(hf.all_special_ids)
    out["meta"] = {"vocab_size": hf.vocab_size, "len": len(hf), "eos": hf.eos_token_id, "pad": hf.pad_token_id,
                   "bos": hf.bos_token_id}
    with gzip.open(OUT, "wt", encoding="utf-8", compresslevel=9) as f:
        json.dump(out, f, ensure_ascii=False, separators=(",", ":"))
    print("wrote", OUT, os.path.getsize(OUT), "bytes;", len(out["texts"]), "texts,", len(out["chat"]), "chat,",
          len(out["id_seqs"]), "id seqs; tokens:", sum(e["n"] for e in out["texts"]))


if __name__ == "__main__":
    main()