#!/usr/bin/env python3
import json
import sys
from pathlib import Path

def clean_index(infile, outfile=None):
    infile = Path(infile)
    outfile = Path(outfile) if outfile else infile

    with infile.open("r", encoding="utf-8") as f:
        data = json.load(f)

    new_weight_map = {}
    for k, v in data["weight_map"].items():
        new_k = k.replace("._checkpoint_wrapped_module", "")
        new_weight_map[new_k] = v

    data["weight_map"] = new_weight_map

    with outfile.open("w", encoding="utf-8") as f:
        json.dump(data, f, indent=2)

if __name__ == "__main__":
    if len(sys.argv) < 2:
        print(f"Usage: {sys.argv[0]} <index_file> [output_file]")
        sys.exit(1)

    in_file = sys.argv[1]
    out_file = sys.argv[2] if len(sys.argv) > 2 else None
    clean_index(in_file, out_file)