81 lines
2.9 KiB
Python
81 lines
2.9 KiB
Python
|
#!/usr/bin/env python3
|
||
|
|
"""
|
||
|
|
export_kronos_weights.py
|
||
|
|
========================
|
||
|
|
Companion to kronos_reference_capture.py. Dumps EVERY parameter and buffer of
|
||
|
|
Kronos-small and Kronos-Tokenizer-base to flat little-endian float32 .bin files,
|
||
|
|
plus a manifest (tensor name -> filename, shape), for loading in MQL5 via
|
||
|
|
FileReadArray. One-time, offline; no Python at inference.
|
||
|
|
|
||
|
|
python export_kronos_weights.py --kronos_repo /path/to/Kronos --out ./kronos_weights
|
||
|
|
|
||
|
|
It also prints the full name+shape listing to stdout. Paste that back and the
|
||
|
|
MQL5 weight loader + stacks can be written against the exact layout.
|
||
|
|
"""
|
||
|
|
|
||
|
|
import argparse
|
||
|
|
import json
|
||
|
|
import os
|
||
|
|
import re
|
||
|
|
import sys
|
||
|
|
|
||
|
|
import numpy as np
|
||
|
|
|
||
|
|
|
||
|
|
def safe_name(name):
|
||
|
|
"""PyTorch names use dots; make a filesystem- and MQL5-friendly filename."""
|
||
|
|
return re.sub(r"[^A-Za-z0-9]", "_", name)
|
||
|
|
|
||
|
|
|
||
|
|
def dump_module(mod, outdir, label):
|
||
|
|
os.makedirs(outdir, exist_ok=True)
|
||
|
|
manifest = {"label": label, "tensors": {}}
|
||
|
|
tensors = list(mod.named_parameters()) + list(mod.named_buffers())
|
||
|
|
total = 0
|
||
|
|
print(f"\n=== {label} : {len(tensors)} tensors ===")
|
||
|
|
for name, t in tensors:
|
||
|
|
arr = np.ascontiguousarray(t.detach().cpu().float().numpy().astype(np.float32))
|
||
|
|
fn = safe_name(name) + ".bin"
|
||
|
|
arr.tofile(os.path.join(outdir, fn))
|
||
|
|
manifest["tensors"][name] = {
|
||
|
|
"file": fn,
|
||
|
|
"shape": list(arr.shape),
|
||
|
|
"dtype": "float32",
|
||
|
|
"count": int(arr.size),
|
||
|
|
}
|
||
|
|
total += int(arr.size)
|
||
|
|
print(f" {name:55s} shape={tuple(arr.shape)}")
|
||
|
|
manifest["total_params"] = total
|
||
|
|
with open(os.path.join(outdir, "manifest.json"), "w") as f:
|
||
|
|
json.dump(manifest, f, indent=2)
|
||
|
|
print(f" -> {total:,} values, manifest at {os.path.join(outdir, 'manifest.json')}")
|
||
|
|
return manifest
|
||
|
|
|
||
|
|
|
||
|
|
def main():
|
||
|
|
ap = argparse.ArgumentParser()
|
||
|
|
ap.add_argument("--kronos_repo", default=".")
|
||
|
|
ap.add_argument("--out", default="./kronos_weights")
|
||
|
|
ap.add_argument("--tokenizer_id", default="NeoQuasar/Kronos-Tokenizer-base")
|
||
|
|
ap.add_argument("--model_id", default="NeoQuasar/Kronos-small")
|
||
|
|
args = ap.parse_args()
|
||
|
|
|
||
|
|
sys.path.insert(0, os.path.abspath(args.kronos_repo))
|
||
|
|
import torch
|
||
|
|
from model import Kronos, KronosTokenizer
|
||
|
|
|
||
|
|
torch.set_grad_enabled(False)
|
||
|
|
tok = KronosTokenizer.from_pretrained(args.tokenizer_id).float().eval()
|
||
|
|
mdl = Kronos.from_pretrained(args.model_id).float().eval()
|
||
|
|
|
||
|
|
os.makedirs(args.out, exist_ok=True)
|
||
|
|
dump_module(tok, os.path.join(args.out, "tokenizer"), "Kronos-Tokenizer-base")
|
||
|
|
dump_module(mdl, os.path.join(args.out, "predictor"), "Kronos-small")
|
||
|
|
|
||
|
|
print("\nAll weights exported (float32, row-major, little-endian).")
|
||
|
|
print("Send back: config_tokenizer.json (from the capture script) + the two")
|
||
|
|
print("manifest.json name/shape listings, and the MQL5 stacks can be written.")
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|