File size: 2,598 Bytes
53c10a4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 | #!/usr/bin/env python3
"""Merge, validate, and officially summarize ReVSI vLLM shards."""
import argparse
import json
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parent))
from eval_revsi_vllm import print_summary, summarise # noqa: E402
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--output-dir", type=Path, required=True)
parser.add_argument("--num-shards", type=int, required=True)
parser.add_argument("--expected-samples", type=int, default=0)
args = parser.parse_args()
missing = [
shard
for shard in range(args.num_shards)
if not (args.output_dir / f"results_shard{shard}.jsonl").is_file()
]
if missing:
raise RuntimeError(
f"Missing {len(missing)}/{args.num_shards} ReVSI shards: {missing}"
)
records_by_id = {}
for shard in range(args.num_shards):
path = args.output_dir / f"results_shard{shard}.jsonl"
with path.open(encoding="utf-8") as handle:
for line_number, line in enumerate(handle, 1):
if not line.strip():
continue
record = json.loads(line)
sample_id = record.get("id")
if sample_id is None:
raise RuntimeError(f"{path}:{line_number}: record has no id")
records_by_id[str(sample_id)] = record
records = list(records_by_id.values())
records.sort(
key=lambda row: (
0,
int(row["id"]),
)
if str(row["id"]).isdigit()
else (1, str(row["id"]))
)
if args.expected_samples > 0 and len(records) != args.expected_samples:
raise RuntimeError(
f"Expected {args.expected_samples} unique ReVSI samples, "
f"got {len(records)}"
)
merged_path = args.output_dir / "merged_results.jsonl"
with merged_path.open("w", encoding="utf-8") as handle:
for record in records:
handle.write(json.dumps(record, ensure_ascii=False) + "\n")
summary = summarise(records)
summary["num_shards"] = args.num_shards
summary["frame_budgets"] = sorted(
{str(record.get("num_frames") or "") for record in records}
)
summary_path = args.output_dir / "summary.json"
summary_path.write_text(
json.dumps(summary, ensure_ascii=False, indent=2) + "\n",
encoding="utf-8",
)
print_summary(summary)
print(f"Merged: {merged_path}")
print(f"Summary: {summary_path}")
if __name__ == "__main__":
main()
|