Cccccz/HY

1#!/usr/bin/env python32"""Validate BF16 prefeature files, targets, task cardinality, and videos."""3 4from __future__ import annotations5 6import argparse7import json8from collections import Counter9from pathlib import Path10 11from safetensors.torch import load_file12 13from predictor_data.prefeature_schema import (14    CONTEXT_BLOCK_IDS,15    SCHEMA_VERSION_PREFEATURE,16    validate_case_tensors,17    validate_context_prefeature,18    validate_step_tensors,19)20 21 22def main() -> None:23    parser = argparse.ArgumentParser(description=__doc__)24    parser.add_argument(25        "--manifest",26        type=Path,27        default=Path("/mnt/local_nvme/zoubin/cz/hyworldplay_predictor_prefeature_v4/manifest.jsonl"),28    )29    parser.add_argument("--expected_records", type=int, default=3200)30    parser.add_argument("--check_videos", action=argparse.BooleanOptionalAction, default=True)31    args = parser.parse_args()32    manifest = args.manifest.resolve()33    root = manifest.parent34    with manifest.open("r", encoding="utf-8") as handle:35        records = [json.loads(line) for line in handle if line.strip()]36    if len(records) != args.expected_records:37        raise ValueError(f"Found {len(records)} records, expected {args.expected_records}")38    keys = Counter()39    cases_checked = set()40    frame_histogram = Counter()41    for record in records:42        if record["schema_version"] != SCHEMA_VERSION_PREFEATURE:43            raise ValueError("Schema mismatch")44        key = (int(record["case_id"]), int(record["action_id"]), int(record["chunk_id"]))45        keys[key] += 146        if key[0] not in cases_checked:47            validate_case_tensors(load_file(str(root / record["case_tensor_file"]), device="cpu"))48            cases_checked.add(key[0])49        validate_step_tensors(load_file(str(root / record["step_tensor_file"]), device="cpu"))50        block_frames = []51        selected = None52        for block_id in CONTEXT_BLOCK_IDS:53            value = load_file(str(root / record["context_tensor_files"][str(block_id)]), device="cpu")54            block_frames.append(validate_context_prefeature(block_id, value))55            current = value["selected_frame_indices"]56            if selected is not None and not current.equal(selected):57                raise ValueError(f"Selected frame mismatch across blocks for {key}")58            selected = current59        if len(set(block_frames)) != 1 or block_frames[0] != int(record["context_frames"]):60            raise ValueError(f"Context frame mismatch for {key}")61        frame_histogram[block_frames[0]] += 162    duplicates = [key for key, count in keys.items() if count != 1]63    if duplicates:64        raise ValueError(f"Duplicate keys: {duplicates[:10]}")65    if args.check_videos:66        missing_videos = [67            (case_id, action_id)68            for case_id in range(100) for action_id in range(4)69            if not (root / "videos" / f"case_{case_id:04d}_action_{action_id:02d}.mp4").is_file()70        ]71        if missing_videos:72            raise ValueError(f"Missing videos: {missing_videos[:10]}")73    print(json.dumps({74        "status": "ok", "records": len(records), "cases": len(cases_checked),75        "frame_histogram": dict(sorted(frame_histogram.items())),76    }, indent=2))77 78 79if __name__ == "__main__":80    main()81