scripts/preflight_model.py

"""Read-only model package and optional native-processor checks. Never GPU inference."""
import argparse
import json
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))
from grounding_demo.native_runtime import NativeRuntime, inspect_contract


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("model_path", type=Path)
    parser.add_argument("--metadata-only", action="store_true", help="Do not require weight shards")
    parser.add_argument("--processor", action="store_true", help="Execute reviewed custom processor on all configured cases (CPU)")
    parser.add_argument("--report", type=Path, help="Write a new JSON report (refuse overwrite)")
    args = parser.parse_args()
    report = inspect_contract(args.model_path.resolve(), require_weights=not args.metadata_only)
    if args.processor:
        from transformers import AutoProcessor
        from grounding_demo.catalog import load_cases

        processor = AutoProcessor.from_pretrained(str(args.model_path.resolve()), trust_remote_code=True)
        runtime = NativeRuntime(None, processor, report["eos_token_ids"])
        report["cases"] = []
        for row in load_cases():
            inputs, count = runtime.prepare(ROOT / row["image"], row["prompt"], 1024)
            image_tokens = int((inputs["input_ids"] == processor.image_token_id).sum())
            grid_tokens = int(inputs["image_grid_thw"].prod()) // processor.image_processor.merge_size ** 2
            expected_ids = processor.tokenizer.encode(row["expected"], add_special_tokens=False)
            roundtrip = processor.decode(expected_ids, skip_special_tokens=False,
                                         spaces_between_special_tokens=False) == row["expected"]
            if image_tokens != grid_tokens or not roundtrip:
                raise RuntimeError(f"Processor/tokenizer contract failed: {row['id']}")
            report["cases"].append({
                "id": row["id"], "route": row["route"], "prompt_tokens": count,
                "image_tokens": image_tokens, "reference_tokens": len(expected_ids),
                "reference_roundtrip": roundtrip,
            })
        report["scope"] = "native_processor_cpu_only_not_model_accuracy_or_gpu_inference"
    rendered = json.dumps(report, indent=2, ensure_ascii=False)
    if args.report:
        with args.report.open("x", encoding="utf-8") as output:
            output.write(rendered + "\n")
    print(rendered)


if __name__ == "__main__":
    main()