#!/usr/bin/env python3
"""Average several checkpoints of the same run into one ("model soup" / SWA).

All inputs must be my own checkpoints from this run's training tree; averaging along a
training trajectory is the classic stochastic-weight-averaging trick and usually buys a
little robustness for free. Copies the tokenizer/config/processor files from the first input.

Usage: soup.py OUT CKPT1 CKPT2 [CKPT3 ...]
"""

import json
import os
import shutil
import sys

import torch
from safetensors.torch import load_file, save_file


def main():
    out, ckpts = sys.argv[1], sys.argv[2:]
    if len(ckpts) < 2:
        sys.exit(__doc__)
    os.makedirs(out, exist_ok=True)

    index_path = os.path.join(ckpts[0], "model.safetensors.index.json")
    index = json.load(open(index_path))
    shards = sorted(set(index["weight_map"].values()))
    print(f"{len(ckpts)} checkpoints, {len(shards)} shards")

    for shard in shards:
        acc = None
        for i, c in enumerate(ckpts):
            sd = load_file(os.path.join(c, shard))
            if acc is None:
                acc = {k: v.to(torch.float32) for k, v in sd.items()}
            else:
                for k in acc:
                    acc[k] += sd[k].to(torch.float32)
            del sd
        for k in acc:
            acc[k] = (acc[k] / len(ckpts)).to(torch.bfloat16)
        save_file(acc, os.path.join(out, shard), metadata={"format": "pt"})
        print("wrote", shard, flush=True)
        del acc

    for f in os.listdir(ckpts[0]):
        if f.endswith(".safetensors") or f == "STABLE":
            continue
        shutil.copy2(os.path.join(ckpts[0], f), os.path.join(out, f))
    open(os.path.join(out, "STABLE"), "w").close()
    print("soup ->", out)


if __name__ == "__main__":
    main()