inference-samples/generate.py

#!/usr/bin/env python3
"""Generate or verify Pollock's fixed cross-revision inference samples."""

from __future__ import annotations

import argparse
import gc
import hashlib
import json
import platform
from importlib.metadata import version
from pathlib import Path
from typing import Any

import torch
from huggingface_hub import hf_hub_download
from transformers import AutoModelForCausalLM, AutoTokenizer


HERE = Path(__file__).resolve().parent
PACKAGE_NAMES = {
    "torch": "torch",
    "transformers": "transformers",
    "huggingface_hub": "huggingface-hub",
}


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--config", type=Path, default=HERE / "config.json")
    parser.add_argument("--output", type=Path, default=HERE / "results.json")
    parser.add_argument("--markdown", type=Path, default=HERE / "README.md")
    parser.add_argument(
        "--check",
        action="store_true",
        help="Regenerate in memory and require an exact match with --output",
    )
    parser.add_argument(
        "--allow-version-mismatch",
        action="store_true",
        help="Run even if the Python or package versions differ from config.json",
    )
    parser.add_argument(
        "--render-only",
        action="store_true",
        help="Render --markdown from the existing --output without inference",
    )
    return parser.parse_args()


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(8 * 1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def sha256_text(text: str) -> str:
    return hashlib.sha256(text.encode("utf-8")).hexdigest()


def runtime_environment() -> dict[str, str]:
    operating_system = platform.system()
    if operating_system == "Darwin":
        operating_system = f"macOS {platform.mac_ver()[0]}"
    return {
        "operating_system": operating_system,
        "architecture": platform.machine(),
        "python": platform.python_version(),
        **{
            config_name: version(package_name)
            for config_name, package_name in PACKAGE_NAMES.items()
        },
    }


def validate_environment(config: dict[str, Any], allow_mismatch: bool) -> dict[str, str]:
    actual = runtime_environment()
    expected = config["environment"]
    mismatches = {
        name: {"expected": expected[name], "actual": actual[name]}
        for name in expected
        if actual.get(name) != expected[name]
    }
    if mismatches and not allow_mismatch:
        raise RuntimeError(f"Environment mismatch: {json.dumps(mismatches)}")
    return actual


def generate_results(config: dict[str, Any], environment: dict[str, str]) -> dict[str, Any]:
    protocol = config["protocol"]
    if protocol["device"] != "cpu" or protocol["dtype"] != "float32":
        raise ValueError("The fixed suite requires CPU float32 inference")

    torch.set_num_threads(protocol["cpu_threads"])
    torch.set_num_interop_threads(1)
    torch.use_deterministic_algorithms(True)

    revision_results = []
    for revision in config["revisions"]:
        local_path = revision.get("local_path")
        if local_path:
            model_source = Path(local_path)
            if not model_source.is_absolute():
                model_source = (HERE / model_source).resolve()
            weights_path = model_source / "model.safetensors"
            load_kwargs: dict[str, Any] = {}
            source_label = str(model_source)
        else:
            model_source = config["model_id"]
            weights_path = Path(
                hf_hub_download(
                    repo_id=config["model_id"],
                    filename="model.safetensors",
                    revision=revision["commit"],
                    token=False,
                )
            )
            load_kwargs = {"revision": revision["commit"], "token": False}
            source_label = revision["commit"]
        print(
            f"Loading {revision['revision_id']} from {source_label}...",
            flush=True,
        )
        actual_model_sha = sha256_file(weights_path)
        if actual_model_sha != revision["model_sha256"]:
            raise RuntimeError(
                f"Weight hash mismatch for {revision['revision_id']}: "
                f"{actual_model_sha} != {revision['model_sha256']}"
            )

        tokenizer = AutoTokenizer.from_pretrained(
            model_source,
            trust_remote_code=False,
            **load_kwargs,
        )
        model = AutoModelForCausalLM.from_pretrained(
            model_source,
            dtype=torch.float32,
            trust_remote_code=False,
            **load_kwargs,
        )
        model.to("cpu")
        model.eval()

        prompt_results = []
        for prompt in config["prompts"]:
            torch.manual_seed(protocol["seed"])
            inputs = tokenizer(prompt["text"], return_tensors="pt")
            input_length = inputs["input_ids"].shape[1]
            with torch.inference_mode():
                output = model.generate(
                    **inputs,
                    do_sample=protocol["do_sample"],
                    temperature=protocol["temperature"],
                    top_k=protocol["top_k"],
                    top_p=protocol["top_p"],
                    max_new_tokens=protocol["max_new_tokens"],
                    num_return_sequences=protocol["num_return_sequences"],
                    eos_token_id=tokenizer.eos_token_id,
                    pad_token_id=tokenizer.eos_token_id,
                    use_cache=True,
                )

            full_token_ids = output[0].tolist()
            completion_token_ids = full_token_ids[input_length:]
            full_text = tokenizer.decode(
                full_token_ids,
                skip_special_tokens=protocol["skip_special_tokens"],
            )
            completion_text = tokenizer.decode(
                completion_token_ids,
                skip_special_tokens=protocol["skip_special_tokens"],
            )
            prompt_results.append(
                {
                    "prompt_id": prompt["prompt_id"],
                    "prompt": prompt["text"],
                    "prompt_token_ids": full_token_ids[:input_length],
                    "completion_token_ids": completion_token_ids,
                    "completion_tokens": len(completion_token_ids),
                    "stopped_on_eos": bool(
                        completion_token_ids
                        and completion_token_ids[-1] == tokenizer.eos_token_id
                    ),
                    "completion": completion_text,
                    "full_text": full_text,
                    "full_text_sha256": sha256_text(full_text),
                }
            )
            print(
                f"  {prompt['prompt_id']}: {len(completion_token_ids)} tokens",
                flush=True,
            )

        revision_results.append(
            {
                **{key: value for key, value in revision.items() if key != "local_path"},
                "loaded_model_version": getattr(model.config, "model_version", None),
                "results": prompt_results,
            }
        )
        del model, tokenizer
        gc.collect()

    return {
        "schema_version": 1,
        "suite_id": config["suite_id"],
        "model_id": config["model_id"],
        "environment": environment,
        "protocol": protocol,
        "revisions": revision_results,
    }


def render_markdown(config: dict[str, Any], results: dict[str, Any]) -> str:
    protocol = results["protocol"]
    environment = results["environment"]
    revisions = {item["revision_id"]: item for item in results["revisions"]}
    lines = [
        "# Fixed cross-revision inference samples",
        "",
        "These samples compare Pollock revisions under one fixed inference "
        "protocol. They are behavioral examples, not benchmark scores or factuality claims. "
        "Outputs are shown without cherry-picking, including errors and repetition.",
        "",
        "## Protocol",
        "",
        f"- Suite: `{results['suite_id']}`",
        f"- Runtime: {environment['operating_system']} "
        f"{environment['architecture']}, Python {environment['python']}, "
        f"PyTorch {environment['torch']}, "
        f"Transformers {environment['transformers']}, huggingface_hub "
        f"{environment['huggingface_hub']}",
        f"- Execution: {protocol['device'].upper()}, `{protocol['dtype']}`, "
        f"{protocol['cpu_threads']} PyTorch CPU thread",
        f"- Sampling: seed {protocol['seed']} reset before every prompt, temperature "
        f"{protocol['temperature']}, top-k {protocol['top_k']}, top-p "
        f"{protocol['top_p']}",
        f"- Limit: {protocol['max_new_tokens']} new tokens, with the prompt included in "
        "each displayed output",
        "",
        "The older samples embedded in `training-history/r001.md` remain part of the "
        "historical record. They differ from this suite because their complete execution "
        "environment was not pinned. Only the results below are used for direct "
        "cross-revision comparison.",
        "",
        "Exact token replay was verified in the recorded environment. Sampling may "
        "diverge on another operating system or CPU architecture even when package "
        "versions and seeds match.",
        "",
        "## Revisions",
        "",
        "| Revision | Release | Immutable weights | `model.safetensors` SHA-256 |",
        "|---|---|---|---|",
    ]
    for revision in results["revisions"]:
        commit = revision["commit"]
        commit_url = f"https://huggingface.co/{results['model_id']}/commit/{commit}"
        weights = f"[`{commit[:7]}`]({commit_url})"
        lines.append(
            f"| {revision['revision_id']} | {revision['release']} | "
            f"{weights} | `{revision['model_sha256']}` |"
        )

    lines.extend(
        [
            "",
            "## Results",
            "",
        ]
    )
    for prompt in config["prompts"]:
        lines.extend(
            [
                f"### {prompt['label']}",
                "",
                f"Prompt: `{prompt['text']}`",
                "",
            ]
        )
        for revision_config in config["revisions"]:
            revision_id = revision_config["revision_id"]
            result = next(
                item
                for item in revisions[revision_id]["results"]
                if item["prompt_id"] == prompt["prompt_id"]
            )
            lines.extend(
                [
                    f"#### {revision_id} - {revision_config['release']}",
                    "",
                    "````text",
                    result["full_text"],
                    "````",
                    "",
                    f"Output SHA-256: `{result['full_text_sha256']}`",
                    "",
                ]
            )

    lines.extend(
        [
            "## Reproduce or verify",
            "",
            "Run from this directory. The check downloads approximately 2.4 GB of "
            "published model artifacts if they are not already cached.",
            "",
            "```bash",
            "python -m venv .venv",
            "source .venv/bin/activate",
            "pip install -r requirements.txt",
            "python generate.py --check",
            "```",
            "",
            "`generate.py` rejects package-version mismatches by default, verifies each "
            "weight file against `config.json`, and compares both token IDs and rendered "
            "text through `results.json`.",
            "",
            "Machine-readable protocol and outputs: [`config.json`](./config.json) and "
            "[`results.json`](./results.json).",
            "",
        ]
    )
    return "\n".join(lines)


def main() -> None:
    args = parse_args()
    config = json.loads(args.config.read_text(encoding="utf-8"))
    if args.render_only:
        if args.check:
            raise ValueError("--render-only and --check cannot be combined")
        results = json.loads(args.output.read_text(encoding="utf-8"))
        args.markdown.write_text(
            render_markdown(config, results),
            encoding="utf-8",
        )
        print(f"Wrote {args.markdown}")
        return

    environment = validate_environment(config, args.allow_version_mismatch)
    results = generate_results(config, environment)
    markdown = render_markdown(config, results)

    if args.check:
        expected = json.loads(args.output.read_text(encoding="utf-8"))
        if results != expected:
            raise RuntimeError(f"Generated results differ from {args.output}")
        if markdown != args.markdown.read_text(encoding="utf-8"):
            raise RuntimeError(f"Rendered Markdown differs from {args.markdown}")
        print(f"Exact match: {args.output} and {args.markdown}")
        return

    args.output.write_text(
        json.dumps(results, indent=2, ensure_ascii=False) + "\n",
        encoding="utf-8",
    )
    args.markdown.write_text(markdown, encoding="utf-8")
    print(f"Wrote {args.output} and {args.markdown}")


if __name__ == "__main__":
    main()