convert_flax_checkpoint.py

from huggingface_hub import hf_hub_download
from flax.serialization import msgpack_restore
from safetensors.flax import save_file
import numpy as np

filename = hf_hub_download("gpt2", filename="flax_model.msgpack")
with open(filename, "rb") as f:
    data = f.read()
flax_weights = msgpack_restore(data)


def flatten(weights, prefix=""):
    values = {}
    for k, v in weights.items():
        newprefix = f"{prefix}.{k}" if prefix else f"{k}"
        print(newprefix)
        if isinstance(v, dict):
            values.update(flatten(v, prefix=newprefix))
        elif isinstance(v, np.ndarray):
            values[newprefix] = v
    return values


weights = flatten(flax_weights)
save_file(weights, "model.safetensors")