beam-boot / scripts /export_checkpoint.py
StarAtNyte1's picture
Add runnable checkpoint and source
bcb308c verified
Raw History Blame Contribute Delete
566 Bytes
"""Compile the bundled ROM into a standalone, stock Transformers checkpoint."""
import argparse
from pathlib import Path
from beamboot.causal_lm import CausalMachine
from beamboot.os import firmware
def main():
parser = argparse.ArgumentParser()
parser.add_argument('output', type=Path, help='output directory for config and safetensors')
args = parser.parse_args()
model = CausalMachine(firmware()[0])
model.save_pretrained(args.output)
print(f'Saved constructed Beam Boot checkpoint to {args.output}')
if __name__ == '__main__': main()