scripts/verify_assets.py

"""Validate local assets needed for an Orienter smoke or reproduction run."""

import argparse
import hashlib
import json
import sys
from pathlib import Path


if __package__ in {None, ""}:
    sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from approach.app_metadata import load_app_metadata_cache
from approach.pipeline_utils import parse_orienter_image_name, resolve_image_path


CHECKPOINT_RELATIVE_PATH = Path("approach/ovod/APE/ape_d_model_final.pth")
CHECKPOINT_SHA256 = "3548f41a3238148180e08fd4b16c71f4abc3ac3caf9c8434444462d1bdb7f965"
CHECKPOINT_SIZE = 5_956_547_279


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


def load_jsonl(path: Path):
    with path.open(encoding="utf-8") as file:
        return [json.loads(line) for line in file if line.strip()]


def check_checkpoint(repo_root: Path, hash_checkpoint: bool):
    checkpoint = repo_root / CHECKPOINT_RELATIVE_PATH
    if not checkpoint.is_file():
        raise FileNotFoundError(f"Missing APE checkpoint: {checkpoint}")
    size = checkpoint.stat().st_size
    if size != CHECKPOINT_SIZE:
        raise ValueError(f"Unexpected checkpoint size: {size}")
    result = {"path": str(CHECKPOINT_RELATIVE_PATH), "size_bytes": size}
    if hash_checkpoint:
        digest = sha256_file(checkpoint)
        if digest != CHECKPOINT_SHA256:
            raise ValueError(f"Unexpected checkpoint sha256: {digest}")
        result["sha256"] = digest
    return result


def check_questions(questions_path: Path, images_dir: Path):
    records = load_jsonl(questions_path)
    if not records:
        raise ValueError(f"Question manifest is empty: {questions_path}")
    seen_question_ids = set()
    seen_image_ids = set()
    app_ids = set()
    image_ids = []
    for record in records:
        question_id = record["question_id"]
        if question_id in seen_question_ids:
            raise ValueError(f"Duplicate question_id: {question_id}")
        seen_question_ids.add(question_id)
        image_name = record["image"]
        image_path = resolve_image_path(images_dir, image_name)
        if not image_path.is_file():
            raise FileNotFoundError(f"Question references a missing image: {image_path}")
        app_id, _, image_id = parse_orienter_image_name(Path(image_name).name)
        if image_id in seen_image_ids:
            raise ValueError(f"Duplicate image_id in question manifest: {image_id}")
        seen_image_ids.add(image_id)
        if "image_id" in record and int(record["image_id"]) != image_id:
            raise ValueError(f"image_id mismatch for {image_name}: {record['image_id']} != {image_id}")
        app_ids.add(app_id)
        image_ids.append(image_id)
    return {
        "questions": len(records),
        "unique_question_ids": len(seen_question_ids),
        "unique_images": len(set(image_ids)),
        "app_ids": sorted(app_ids),
    }


def check_metadata_cache(cache_path: Path, app_ids):
    cache = load_app_metadata_cache(cache_path)
    missing = [app_id for app_id in app_ids if str(app_id) not in cache]
    if missing:
        raise KeyError(f"Metadata cache is missing app IDs: {', '.join(missing)}")
    return {"metadata_cache": str(cache_path), "covered_app_ids": len(app_ids)}


def check_embedding_cache(cache_path: Path, manifest_path: Path):
    manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
    expected_size = int(manifest["size_bytes"])
    expected_sha256 = manifest["sha256"]
    actual_size = cache_path.stat().st_size
    if actual_size != expected_size:
        raise ValueError(
            f"Unexpected embedding cache size: {actual_size}; expected {expected_size}"
        )
    actual_sha256 = sha256_file(cache_path)
    if actual_sha256 != expected_sha256:
        raise ValueError(
            f"Unexpected embedding cache sha256: {actual_sha256}; expected {expected_sha256}"
        )
    return {
        "artifact": manifest.get("artifact", cache_path.name),
        "size_bytes": actual_size,
        "sha256": actual_sha256,
        "entry_count": manifest.get("entry_count"),
        "embedding_dimension": manifest.get("embedding_dimension"),
    }


def build_parser():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--repo-root", type=Path, default=Path.cwd())
    parser.add_argument("--questions", type=Path)
    parser.add_argument("--images-dir", type=Path)
    parser.add_argument("--app-metadata-cache", type=Path)
    parser.add_argument(
        "--embedding-cache",
        type=Path,
        help="external frozen semantic cache to verify against evaluation/cache_manifest.json",
    )
    parser.add_argument(
        "--embedding-cache-manifest",
        type=Path,
        help="override the semantic cache manifest (defaults to the repository manifest)",
    )
    parser.add_argument("--hash-checkpoint", action="store_true")
    parser.add_argument("--skip-checkpoint", action="store_true")
    return parser


def main(argv=None):
    args = build_parser().parse_args(argv)
    report = {}
    if not args.skip_checkpoint:
        report["checkpoint"] = check_checkpoint(args.repo_root, args.hash_checkpoint)
    if args.embedding_cache:
        manifest_path = args.embedding_cache_manifest or (
            args.repo_root / "evaluation" / "cache_manifest.json"
        )
        report["embedding_cache"] = check_embedding_cache(
            args.embedding_cache,
            manifest_path,
        )
    elif args.embedding_cache_manifest:
        raise SystemExit("--embedding-cache-manifest requires --embedding-cache")
    if args.questions or args.images_dir:
        if not args.questions or not args.images_dir:
            raise SystemExit("--questions and --images-dir must be supplied together")
        report["questions"] = check_questions(args.questions, args.images_dir)
        if args.app_metadata_cache:
            report["metadata"] = check_metadata_cache(
                args.app_metadata_cache,
                report["questions"]["app_ids"],
            )
    elif args.app_metadata_cache:
        raise SystemExit("--app-metadata-cache requires --questions and --images-dir")
    print(json.dumps(report, indent=2, sort_keys=True))


if __name__ == "__main__":
    main()