#!/usr/bin/env python
"""Merge sharded lambda/low-rank results into one calibration artifact.

Shards are disjoint by layer, so the merge is a dict update; the check is that all 312 target
layers are present, since a silently short artifact would convert without complaint and only
show up as a bad video.
"""
import argparse, glob
from pathlib import Path
import torch

ap = argparse.ArgumentParser()
ap.add_argument("--shards", required=True)
ap.add_argument("--out", required=True)
ap.add_argument("--span", required=True)
ap.add_argument("--expect", type=int, default=312)
a = ap.parse_args()

layers = {}
files = sorted(Path(a.shards).glob("slr_*.pt"))
for f in files:
    layers.update(torch.load(f, map_location="cpu", weights_only=False)["layers"])
assert len(layers) == a.expect, f"got {len(layers)} layers from {len(files)} shards, want {a.expect}"
torch.save({"rank": 32, "num_grids": 20, "span": a.span,
            "base_model": "MiniMaxAI/MiniMax-H3@bfc8ed0353f5a9733be73e6b2c98ec0948195b86",
            "layers": layers}, a.out)
print(f"merged {len(files)} shards, {len(layers)} layers -> {a.out}")