Finnish-NLP/Ahma-3B

1# This script converts model checkpoint trained by EsayLM to a standard2# mspack checkpoint that can be loaded by huggingface transformers or3# flax.serialization.msgpack_restore. Such conversion allows models to be4# used by other frameworks that integrate with huggingface transformers.5 6import pprint7from functools import partial8import os9import numpy as np10import mlxu11import jax.numpy as jnp12import flax.serialization13from EasyLM.checkpoint import StreamingCheckpointer14from EasyLM.jax_utils import float_to_dtype15 16 17FLAGS, FLAGS_DEF = mlxu.define_flags_with_default(18    load_checkpoint='',19    output_file='',20    streaming=False,21    float_dtype='bf16',22)23 24 25def main(argv):26    assert FLAGS.load_checkpoint != '' and FLAGS.output_file != '', 'input and output must be specified'27    params = StreamingCheckpointer.load_trainstate_checkpoint(28        FLAGS.load_checkpoint, disallow_trainstate=True29    )[1]['params']30 31    if FLAGS.streaming:32        StreamingCheckpointer.save_train_state_to_file(33            params, FLAGS.output_file, float_dtype=FLAGS.float_dtype34        )35    else:36        params = float_to_dtype(params, FLAGS.float_dtype)37        with mlxu.open_file(FLAGS.output, 'wb') as fout:38            fout.write(flax.serialization.msgpack_serialize(params, in_place=True))39 40 41if __name__ == "__main__":42    mlxu.run(main)43