infer_nested_model.py

#!/usr/bin/env python3
"""Cached byte generation for two- and three-level nested Mamba checkpoints."""
from __future__ import annotations

import argparse
import base64
import gc
import io
import json
import time
from pathlib import Path
from typing import Dict, List, Optional

import torch

from nested_inference_tools import generate_bytes, load_nested_checkpoint


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--checkpoint",
        default=".",
        help="Hugging Face model directory (default: current directory) or legacy last.pt.",
    )
    prompt = parser.add_mutually_exclusive_group()
    prompt.add_argument("--prompt", default="", help="UTF-8 prompt text.")
    prompt.add_argument("--prompt-file", help="Read prompt bytes from this file.")
    parser.add_argument("--max-new-bytes", type=int, default=256)
    parser.add_argument("--temperature", type=float, default=0.8)
    parser.add_argument("--top-p", type=float, default=0.9)
    parser.add_argument("--top-k", type=int, default=0)
    parser.add_argument("--repeat-penalty", type=float, default=1.05)
    parser.add_argument("--repeat-window", type=int, default=256)
    parser.add_argument("--greedy", action="store_true")
    parser.add_argument("--seed", type=int, default=1234)
    parser.add_argument(
        "--precision",
        choices=["fp16", "bf16", "fp32"],
        default="bf16",
        help=(
            "Model weight/activation precision. BF16 is the safe default for "
            "Mamba-2's long cached rollouts; use FP16 only for checkpoints and "
            "GPUs verified to remain finite."
        ),
    )
    parser.add_argument("--device", default="cuda:0", help="Single-device inference target.")
    parser.add_argument("--fine-device", default=None)
    parser.add_argument("--nested-devices", default=None)
    parser.add_argument("--tertiary-device", default=None)
    parser.add_argument("--use-saved-placement", action="store_true")
    parser.add_argument("--output", help="Optional raw-byte output file.")
    parser.add_argument("--stats-json", help="Optional JSON statistics output.")
    parser.add_argument(
        "--image-output-dir",
        help=(
            "Extract complete generated P6 PPM images, convert them to PNG, and "
            "write them to this directory. Images are also embedded in --html-output."
        ),
    )
    parser.add_argument("--html-output", "--output-html", dest="html_output", help="Write a self-contained prompt/response report with toggleable L2/L3 influence coloring.")
    return parser.parse_args()


def _ppm_token(data: bytes, position: int) -> tuple[bytes, int]:
    """Read one whitespace/comment-delimited PPM header token."""
    size = len(data)
    while position < size:
        if data[position] in b" \t\r\n":
            position += 1
            continue
        if data[position] == ord("#"):
            newline = data.find(b"\n", position)
            if newline < 0:
                raise ValueError("unterminated PPM header comment")
            position = newline + 1
            continue
        break
    start = position
    while position < size and data[position] not in b" \t\r\n#":
        position += 1
    if position == start:
        raise ValueError("missing PPM header token")
    return data[start:position], position


def extract_ppm_images(data: bytes) -> tuple[List[Dict[str, object]], List[str]]:
    """Extract complete P6 RGB images from an arbitrary generated byte stream."""
    try:
        from PIL import Image
    except ImportError as error:
        return [], [f"Pillow is required to render generated images: {error}"]

    images: List[Dict[str, object]] = []
    warnings: List[str] = []
    search_from = 0
    while len(images) < 32:
        start = data.find(b"P6", search_from)
        if start < 0:
            break
        search_from = start + 2
        if start > 0 and data[start - 1] not in b" \t\r\n>":
            continue
        if start + 2 >= len(data) or data[start + 2] not in b" \t\r\n":
            continue
        try:
            width_token, position = _ppm_token(data, start + 2)
            height_token, position = _ppm_token(data, position)
            maximum_token, position = _ppm_token(data, position)
            width, height, maximum = int(width_token), int(height_token), int(maximum_token)
            if width < 1 or height < 1 or width * height > 16_777_216:
                raise ValueError(f"unsafe dimensions {width}x{height}")
            if maximum != 255:
                raise ValueError(f"unsupported maximum channel value {maximum}; expected 255")
            if position >= len(data) or data[position] not in b" \t\r\n":
                raise ValueError("missing whitespace before PPM pixel payload")
            # The delimiter is one whitespace unit. Treat CRLF as one unit so
            # the first pixel is never shifted on Windows-produced headers.
            pixel_start = position + 1
            if data[position] == ord("\r") and pixel_start < len(data) and data[pixel_start] == ord("\n"):
                pixel_start += 1
            pixel_bytes = width * height * 3
            pixel_end = pixel_start + pixel_bytes
            if pixel_end > len(data):
                warnings.append(
                    f"Incomplete PPM at byte {start}: {width}x{height} needs "
                    f"{pixel_bytes:,} RGB bytes, but only {max(0, len(data)-pixel_start):,} remain."
                )
                continue
            image = Image.frombytes("RGB", (width, height), data[pixel_start:pixel_end])
            encoded = io.BytesIO()
            image.save(encoded, format="PNG", optimize=True)
            images.append(
                {
                    "width": width,
                    "height": height,
                    "start": start,
                    "end": pixel_end,
                    "png": encoded.getvalue(),
                }
            )
            search_from = pixel_end
        except (TypeError, ValueError) as error:
            warnings.append(f"Invalid PPM candidate at byte {start}: {error}.")
    return images, warnings


def image_report_payload(images: List[Dict[str, object]], warnings: List[str]) -> Dict[str, object]:
    return {
        "images": [
            {
                "width": int(item["width"]),
                "height": int(item["height"]),
                "start": int(item["start"]),
                "end": int(item["end"]),
                "data_uri": "data:image/png;base64," + base64.b64encode(item["png"]).decode("ascii"),
            }
            for item in images
        ],
        "warnings": list(warnings),
    }


def write_extracted_images(
    directory: Path, images: List[Dict[str, object]], *, prefix: str
) -> List[str]:
    directory.mkdir(parents=True, exist_ok=True)
    paths: List[str] = []
    for index, item in enumerate(images, 1):
        path = directory / f"{prefix}_{index:03d}_{item['width']}x{item['height']}.png"
        path.write_bytes(item["png"])
        paths.append(str(path))
    return paths


def attributed_utf8_segments(
    data: bytes, attributions: List[Dict[str, object]]
) -> List[Dict[str, object]]:
    """Group byte attribution into valid UTF-8 characters without losing data."""
    segments: List[Dict[str, object]] = []
    index = 0
    while index < len(data):
        first = data[index]
        width = (
            1
            if first < 0x80
            else 2
            if 0xC2 <= first <= 0xDF
            else 3
            if 0xE0 <= first <= 0xEF
            else 4
            if 0xF0 <= first <= 0xF4
            else 1
        )
        chunk = data[index : index + width]
        try:
            text = chunk.decode("utf-8", "strict")
        except UnicodeDecodeError:
            width = 1
            chunk = data[index : index + 1]
            text = chunk.decode("utf-8", "replace")
        values = attributions[index : index + width]
        count = max(1, len(values))
        segments.append(
            {
                "text": text,
                "bytes": list(chunk),
                "level2_active": any(bool(item.get("level2_active")) for item in values),
                "level3_active": any(bool(item.get("level3_active")) for item in values),
                "level2_delta_logp": sum(
                    float(item.get("level2_delta_logp", 0.0)) for item in values
                )
                / count,
                "level3_delta_logp": sum(
                    float(item.get("level3_delta_logp", 0.0)) for item in values
                )
                / count,
                "hierarchy_delta_logp": sum(
                    float(item.get("hierarchy_delta_logp", 0.0)) for item in values
                )
                / count,
            }
        )
        index += width
    return segments


def make_generation_html(
    prompt: bytes,
    generated: bytes,
    attributions: List[Dict[str, object]],
    stats: Dict[str, object],
    comparisons: Optional[List[Dict[str, object]]] = None,
    image_report: Optional[Dict[str, object]] = None,
) -> str:
    segments = attributed_utf8_segments(generated, attributions)
    payload = json.dumps(
        {
            "prompt": prompt.decode("utf-8", "replace"),
            "segments": segments,
            "stats": stats,
            "comparisons": comparisons or [],
            "image_report": image_report or {"images": [], "warnings": []},
        },
        separators=(",", ":"),
    ).replace("</", "<\\/")
    return f"""<!doctype html>
<html lang="en"><head><meta charset="utf-8"><meta name="viewport" content="width=device-width,initial-scale=1">
<title>Nested Mamba generation influence</title>
<style>
:root{{--bg:#090d17;--panel:#111a2b;--line:#273652;--text:#e7edf9;--muted:#91a0b8;--prompt:#d6deec;--response:#63a4ff;--fine:#51d6ca;--l2:#ffbe55;--l3:#a78bfa;--negative:#fb7185}}
*{{box-sizing:border-box}}body{{margin:0;background:linear-gradient(145deg,#080c15,#11192b);color:var(--text);font:14px/1.5 system-ui,sans-serif}}main{{max-width:1200px;margin:auto;padding:28px}}
h1{{margin:0 0 6px;font-size:27px}}p{{color:var(--muted)}}.toolbar,.panel{{background:#111a2bf2;border:1px solid var(--line);border-radius:12px}}
.toolbar{{display:flex;align-items:center;gap:18px;flex-wrap:wrap;padding:13px 16px;margin:18px 0}}label{{display:flex;align-items:center;gap:8px;font-weight:650}}input{{accent-color:var(--l3)}}
.legend{{display:flex;gap:14px;flex-wrap:wrap;color:var(--muted);font-size:12px}}.dot{{width:10px;height:10px;border-radius:50%;display:inline-block;margin-right:5px}}
.panel{{padding:20px}}.text{{white-space:pre-wrap;overflow-wrap:anywhere;font:16px/1.65 ui-monospace,SFMono-Regular,Menlo,Consolas,monospace}}.responses{{display:grid;grid-template-columns:repeat(3,minmax(0,1fr));gap:12px;margin-top:16px}}.response-card{{min-width:0;background:#0c1321;border:1px solid var(--line);border-radius:10px;padding:14px}}.response-card h2{{font:700 14px system-ui,sans-serif;margin:0 0 4px}}.response-meta{{color:var(--muted);font:11px system-ui,sans-serif;margin-bottom:11px}}
.prompt{{color:var(--prompt)}}.generated{{color:var(--response);transition:color .15s,opacity .15s}}.divider{{display:block;color:var(--muted);font:12px system-ui,sans-serif;margin:18px 0 7px;border-top:1px solid var(--line);padding-top:10px}}
.cards{{display:grid;grid-template-columns:repeat(auto-fit,minmax(170px,1fr));gap:10px;margin-top:14px}}.card{{background:#0c1321;border:1px solid var(--line);border-radius:9px;padding:11px}}.name{{color:var(--muted);font-size:11px;text-transform:uppercase}}.value{{font-size:18px;font-weight:700}}
.note{{border-left:3px solid var(--l2);padding:9px 12px;background:#ffbe550d}}
.image-panel{{margin-top:16px}}.image-galleries{{display:grid;grid-template-columns:repeat(3,minmax(0,1fr));gap:12px}}.image-gallery{{background:#0c1321;border:1px solid var(--line);border-radius:10px;padding:12px;min-width:0}}.image-gallery h2{{font-size:14px;margin:0 0 9px}}.image-item{{margin:0 0 12px}}.image-item img{{display:block;max-width:100%;height:auto;image-rendering:auto;border:1px solid #354766;background:#05070b}}.image-item figcaption,.image-status{{color:var(--muted);font-size:11px;margin-top:5px;white-space:pre-wrap}}
@media(max-width:900px){{.responses{{grid-template-columns:1fr}}}}
</style></head><body><main>
<h1>Nested Mamba generation influence</h1><p id="subtitle"></p>
<div class="toolbar"><label><input id="influence" type="checkbox" checked> Color generated text by hierarchy influence</label><div class="legend"><span><i class="dot" style="background:var(--fine)"></i>fine only</span><span><i class="dot" style="background:var(--l2)"></i>L2 dominant</span><span><i class="dot" style="background:var(--l3)"></i>L3 dominant</span><span><i class="dot" style="background:var(--negative)"></i>dominant branch reduced selected-byte probability</span></div></div>
<p class="note">Color is a decoder counterfactual, not merely an activation marker. Positive delta means the cached dynamic level increased the selected byte's log probability relative to replacing that level with its beginning-of-document state. Hover text for exact values.</p>
<section class="panel"><div class="text"><span class="divider">PROMPT SHARED BY ALL ROLLOUTS</span><span class="prompt" id="prompt"></span></div><div class="responses"><article class="response-card"><h2>Level 1 only</h2><div class="response-meta" id="level1Meta"></div><div class="text generated" id="level1Response"></div></article><article class="response-card"><h2>Levels 1 + 2</h2><div class="response-meta" id="level12Meta"></div><div class="text generated" id="level12Response"></div></article><article class="response-card"><h2>Levels 1 + 2 + 3</h2><div class="response-meta" id="fullMeta"></div><div class="text" id="response"></div></article></div><div class="cards" id="cards"></div></section>
<section class="panel image-panel"><h1>Generated image rendering</h1><p>Complete P6 RGB byte sequences found in each prompt-plus-rollout stream are validated and embedded here as PNG.</p><div class="image-galleries"><div class="image-gallery"><h2>Level 1 only</h2><div id="level1Images"></div></div><div class="image-gallery"><h2>Levels 1 + 2</h2><div id="level12Images"></div></div><div class="image-gallery"><h2>Levels 1 + 2 + 3</h2><div id="fullImages"></div></div></div></section>
<script>
const R={payload},prompt=document.getElementById("prompt"),response=document.getElementById("response"),toggle=document.getElementById("influence");
prompt.textContent=R.prompt;document.getElementById("subtitle").textContent=`${{R.stats.checkpoint}} • step ${{Number(R.stats.checkpoint_step).toLocaleString()}}`;
const byMode=Object.fromEntries(R.comparisons.map(x=>[x.mode,x]));for(const [mode,id,color] of [["level1","level1Response","#51d6ca"],["level12","level12Response","#ffbe55"]]){{const item=byMode[mode]||{{text:"",generated_bytes:0,seconds:0}};document.getElementById(id).textContent=item.text;document.getElementById(id).dataset.color=color;document.getElementById(mode+"Meta").textContent=`${{Number(item.generated_bytes).toLocaleString()}} bytes • ${{Number(item.seconds).toFixed(2)}} s`;}}document.getElementById("fullMeta").textContent=`${{Number(R.stats.generated_bytes).toLocaleString()}} bytes • ${{Number(R.stats.generation_seconds_including_prefill).toFixed(2)}} s`;
const magnitudes=R.segments.flatMap(x=>[Math.abs(x.level2_delta_logp),Math.abs(x.level3_delta_logp)]).filter(Number.isFinite).sort((a,b)=>a-b),scale=magnitudes[Math.floor(magnitudes.length*.9)]||.01;
function classification(x){{const candidates=[];if(x.level2_active)candidates.push([Math.abs(x.level2_delta_logp),x.level2_delta_logp,"#ffbe55","L2"]);if(x.level3_active)candidates.push([Math.abs(x.level3_delta_logp),x.level3_delta_logp,"#a78bfa","L3"]);if(!candidates.length)return ["#51d6ca",.72,"fine only"];candidates.sort((a,b)=>b[0]-a[0]);const best=candidates[0],color=best[1]<0?"#fb7185":best[2],opacity=.48+.52*Math.min(1,best[0]/scale);return [color,opacity,best[1]<0?best[3]+" negative":best[3]+" dominant"]}}
R.segments.forEach(x=>{{const span=document.createElement("span"),style=classification(x);span.className="generated";span.textContent=x.text;span.dataset.color=style[0];span.dataset.opacity=style[1];span.title=`${{style[2]}} • bytes ${{x.bytes.join(",")}} • L2 Δlogp ${{x.level2_delta_logp.toFixed(5)}} • L3 Δlogp ${{x.level3_delta_logp.toFixed(5)}} • combined Δlogp ${{x.hierarchy_delta_logp.toFixed(5)}}`;response.appendChild(span)}});
function recolor(){{response.querySelectorAll(".generated").forEach(span=>{{span.style.color=toggle.checked?span.dataset.color:"#63a4ff";span.style.opacity=toggle.checked?span.dataset.opacity:"1"}});for(const id of ["level1Response","level12Response"]){{const node=document.getElementById(id);node.style.color=toggle.checked?node.dataset.color:"#63a4ff"}}}}toggle.addEventListener("change",recolor);recolor();
const attrs=R.segments,mean=key=>attrs.length?attrs.reduce((s,x)=>s+Number(x[key]||0),0)/attrs.length:0,cards=[["Prompt bytes",R.stats.prompt_bytes],["Generated bytes",R.stats.generated_bytes],["Generation tok/s",Number(R.stats.generated_bytes_per_second_including_prefill).toFixed(1)],["Mean L2 Δlogp",mean("level2_delta_logp").toFixed(5)],["Mean L3 Δlogp",mean("level3_delta_logp").toFixed(5)],["Mean combined Δlogp",mean("hierarchy_delta_logp").toFixed(5)]];
document.getElementById("cards").innerHTML=cards.map(x=>`<div class="card"><div class="name">${{x[0]}}</div><div class="value">${{x[1]}}</div></div>`).join("");
function renderImages(target,report){{const root=document.getElementById(target),images=report?.images||[],warnings=report?.warnings||[];if(!images.length&&!warnings.length){{root.innerHTML='<div class="image-status">No complete P6 image detected.</div>';return}}images.forEach((item,index)=>{{const figure=document.createElement("figure");figure.className="image-item";const image=document.createElement("img");image.src=item.data_uri;image.alt=`Generated image ${{index+1}}, ${{item.width}} by ${{item.height}} pixels`;const caption=document.createElement("figcaption");caption.textContent=`Image ${{index+1}} • ${{item.width}}×${{item.height}} • byte range ${{item.start.toLocaleString()}}–${{(item.end-1).toLocaleString()}}`;figure.append(image,caption);root.appendChild(figure)}});if(warnings.length){{const status=document.createElement("div");status.className="image-status";status.textContent=warnings.join("\\n");root.appendChild(status)}}}}
renderImages("fullImages",R.image_report);renderImages("level1Images",byMode.level1?.image_report);renderImages("level12Images",byMode.level12?.image_report);
</script></main></body></html>"""


def main() -> None:
    args = parse_args()
    if args.max_new_bytes < 0 or args.temperature <= 0 or not 0 < args.top_p <= 1:
        raise ValueError("max-new-bytes must be non-negative, temperature positive, and top-p in (0, 1]")
    if args.prompt_file:
        prompt = Path(args.prompt_file).expanduser().read_bytes()
    else:
        prompt = args.prompt.encode("utf-8")

    load_started = time.perf_counter()
    loaded = load_nested_checkpoint(
        args.checkpoint,
        precision=args.precision,
        device=args.device,
        fine_device=args.fine_device,
        nested_devices=args.nested_devices,
        tertiary_device=args.tertiary_device,
        use_saved_placement=args.use_saved_placement,
    )
    torch.cuda.synchronize()
    load_seconds = time.perf_counter() - load_started
    generation_started = time.perf_counter()
    generated, stream = generate_bytes(
        loaded,
        prompt,
        max_new_bytes=args.max_new_bytes,
        temperature=args.temperature,
        top_p=args.top_p,
        top_k=args.top_k,
        repeat_penalty=args.repeat_penalty,
        repeat_window=args.repeat_window,
        greedy=args.greedy,
        seed=args.seed,
        collect_hierarchy_attribution=bool(args.html_output),
    )
    torch.cuda.synchronize()
    generation_seconds = time.perf_counter() - generation_started
    combined = prompt + generated
    full_images, full_image_warnings = extract_ppm_images(combined)
    full_image_report = image_report_payload(full_images, full_image_warnings)
    print(combined.decode("utf-8", "replace"))
    attributions = list(stream.get("generated_attribution", []))
    stats = {
        "checkpoint": str(loaded.checkpoint_path),
        "checkpoint_step": loaded.checkpoint_step,
        "trained_tokens": loaded.trained_tokens,
        "precision": loaded.precision,
        "model_parallel": loaded.model_parallel,
        "prompt_bytes": len(prompt),
        "generated_bytes": len(generated),
        "load_seconds": load_seconds,
        "generation_seconds_including_prefill": generation_seconds,
        "generated_bytes_per_second_including_prefill": (
            len(generated) / generation_seconds if generation_seconds else 0.0
        ),
        "fine_patches": int(stream["completed_patches"]),
        "level2_pools": int(stream["completed_nested_patches"]),
        "level3_pools": int(stream.get("completed_tertiary_patches", 0)),
        "hierarchy_attribution": bool(args.html_output),
        "rendered_images": len(full_images),
        "image_render_warnings": full_image_warnings,
    }
    if attributions:
        stats.update(
            {
                "attributed_bytes": len(attributions),
                "level2_active_generated_bytes": sum(
                    bool(item.get("level2_active")) for item in attributions
                ),
                "level3_active_generated_bytes": sum(
                    bool(item.get("level3_active")) for item in attributions
                ),
                "mean_level2_delta_logp": sum(
                    float(item.get("level2_delta_logp", 0.0))
                    for item in attributions
                )
                / len(attributions),
                "mean_level3_delta_logp": sum(
                    float(item.get("level3_delta_logp", 0.0))
                    for item in attributions
                )
                / len(attributions),
                "mean_hierarchy_delta_logp": sum(
                    float(item.get("hierarchy_delta_logp", 0.0))
                    for item in attributions
                )
                / len(attributions),
            }
        )
    comparisons: List[Dict[str, object]] = []
    del stream
    gc.collect()
    if args.html_output:
        for mode, label in (("level1", "Level 1 only"), ("level12", "Levels 1 + 2")):
            comparison_started = time.perf_counter()
            comparison_bytes, comparison_stream = generate_bytes(
                loaded,
                prompt,
                max_new_bytes=args.max_new_bytes,
                temperature=args.temperature,
                top_p=args.top_p,
                top_k=args.top_k,
                repeat_penalty=args.repeat_penalty,
                repeat_window=args.repeat_window,
                greedy=args.greedy,
                seed=args.seed,
                hierarchy_mode=mode,
            )
            torch.cuda.synchronize()
            comparison_seconds = time.perf_counter() - comparison_started
            comparisons.append(
                {
                    "mode": mode,
                    "label": label,
                    "text": comparison_bytes.decode("utf-8", "replace"),
                    "generated_bytes": len(comparison_bytes),
                    "seconds": comparison_seconds,
                    "fine_patches": int(comparison_stream["completed_patches"]),
                    "level2_pools": int(
                        comparison_stream["completed_nested_patches"]
                    ),
                    "level3_pools": int(
                        comparison_stream.get("completed_tertiary_patches", 0)
                    ),
                    "image_report": image_report_payload(
                        *extract_ppm_images(prompt + comparison_bytes)
                    ),
                }
            )
            del comparison_stream
            gc.collect()
        # Keep base64 PNG payloads in the HTML only; terminal/stats JSON should
        # stay compact even when a rollout contains a large rendered image.
        stats["comparison_rollouts"] = [
            {key: value for key, value in item.items() if key != "image_report"}
            for item in comparisons
        ]
    if args.output:
        Path(args.output).expanduser().write_bytes(combined)
    if args.image_output_dir:
        image_paths = write_extracted_images(
            Path(args.image_output_dir).expanduser(), full_images, prefix="full"
        )
        stats["image_output_files"] = image_paths
        for item in comparisons:
            report = item.get("image_report") or {}
            comparison_images = []
            for encoded in report.get("images", []):
                raw = base64.b64decode(str(encoded["data_uri"]).split(",", 1)[1])
                comparison_images.append({**encoded, "png": raw})
            stats.setdefault("comparison_image_output_files", {})[item["mode"]] = (
                write_extracted_images(
                    Path(args.image_output_dir).expanduser(),
                    comparison_images,
                    prefix=str(item["mode"]),
                )
            )
    print(json.dumps(stats, indent=2))
    if args.stats_json:
        Path(args.stats_json).expanduser().write_text(
            json.dumps(stats, indent=2), encoding="utf-8"
        )
    if args.html_output:
        html_path = Path(args.html_output).expanduser()
        html_path.parent.mkdir(parents=True, exist_ok=True)
        html_path.write_text(
            make_generation_html(
                prompt,
                generated,
                attributions,
                stats,
                comparisons=comparisons,
                image_report=full_image_report,
            ),
            encoding="utf-8",
        )
        print(f"wrote hierarchy influence report: {html_path}")


if __name__ == "__main__":
    main()