21 lines
456 B
Python
21 lines
456 B
Python
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)
|