import json
from safetensors.torch import load_file

path = "model.safetensors"
tensors = load_file(path)

weight_map = {}
total_size = 0

for name, tensor in tensors.items():
    weight_map[name] = "model.safetensors"
    total_size += tensor.element_size() * tensor.numel()

index = {
    "metadata": {"total_size": str(total_size)},
    "weight_map": weight_map,
}

with open("model.safetensors.index.json", "w") as f:
    json.dump(index, f, indent=2)