275 lines
12 KiB
Python
275 lines
12 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
kronos_reference_capture.py
|
|
===========================
|
|
First build artifact for the "Kronos-small inference in native MQL5" project.
|
|
|
|
What it does
|
|
------------
|
|
1. Loads Kronos-small + Kronos-Tokenizer-base and PINS their exact resolved
|
|
architecture to JSON (so the tokenizer config gets captured automatically).
|
|
2. Freezes ONE fully reproducible input window and saves it as the contract the
|
|
MQL5 side will read verbatim -- MQL5 never re-synthesizes the input, it reads
|
|
these bytes, so the data generation here does not need to be portable.
|
|
3. Replicates KronosPredictor's preprocessing EXACTLY (per-window z-score with
|
|
eps=1e-5, population std, clip +/-5) and saves every intermediate.
|
|
4. Captures a DETERMINISTIC verification ladder -- no sampling anywhere -- at each
|
|
model boundary, so each MQL5 component can be checked against ground truth as
|
|
it is built:
|
|
|
|
x_norm -> (s1_ids, s2_ids) -> s1_logits -> s2_logits -> reconstruction
|
|
|
|
5. (Optional) one SEEDED end-to-end forecast as a coarse full-pipeline check.
|
|
This one is stochastic; it is documented as such and is not a bit-exact target.
|
|
|
|
Every array is written twice:
|
|
- <name>.npy for Python-side inspection
|
|
- <name>.bin + manifest.json little-endian, row-major, for MQL5 FileReadArray
|
|
|
|
Usage
|
|
-----
|
|
Run from inside the cloned Kronos repo (so `from model import ...` resolves), or
|
|
point --kronos_repo at it:
|
|
|
|
python kronos_reference_capture.py --kronos_repo /path/to/Kronos --out ./kronos_refs
|
|
|
|
Notes
|
|
-----
|
|
- CPU + float32 are forced for reproducibility. Do not run this on GPU.
|
|
- np.std defaults to population std (ddof=0); KronosPredictor relies on that, so
|
|
MQL5 must also use population std (divide by N, not N-1).
|
|
"""
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import sys
|
|
|
|
import numpy as np
|
|
|
|
|
|
def to_py(v):
|
|
"""Make a config value JSON-serializable."""
|
|
try:
|
|
import torch
|
|
if isinstance(v, torch.Tensor):
|
|
return v.tolist()
|
|
except Exception:
|
|
pass
|
|
if isinstance(v, (np.integer,)):
|
|
return int(v)
|
|
if isinstance(v, (np.floating,)):
|
|
return float(v)
|
|
if isinstance(v, (np.ndarray,)):
|
|
return v.tolist()
|
|
return v
|
|
|
|
|
|
def dump_config(obj, curated_keys):
|
|
"""Best-effort capture of a HF-mixin module's resolved config."""
|
|
cfg = {}
|
|
# 1) try the stored hub-mixin config dict
|
|
for attr in ("config", "_hub_mixin_config", "_config"):
|
|
c = getattr(obj, attr, None)
|
|
if isinstance(c, dict):
|
|
cfg.update({k: to_py(v) for k, v in c.items()})
|
|
# 2) scrape curated attributes directly off the instantiated module
|
|
for k in curated_keys:
|
|
if hasattr(obj, k):
|
|
cfg[k] = to_py(getattr(obj, k))
|
|
return cfg
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("--kronos_repo", default=".",
|
|
help="Path to cloned Kronos repo (the dir containing model/)")
|
|
ap.add_argument("--out", default="./kronos_refs")
|
|
ap.add_argument("--tokenizer_id", default="NeoQuasar/Kronos-Tokenizer-base")
|
|
ap.add_argument("--model_id", default="NeoQuasar/Kronos-small")
|
|
ap.add_argument("--lookback", type=int, default=256, help="context length, <=512")
|
|
ap.add_argument("--pred_len", type=int, default=16)
|
|
ap.add_argument("--clip", type=float, default=5.0)
|
|
ap.add_argument("--eps", type=float, default=1e-5)
|
|
ap.add_argument("--seed", type=int, default=0)
|
|
ap.add_argument("--skip_forecast", action="store_true",
|
|
help="Skip the seeded end-to-end forecast (stochastic) step")
|
|
args = ap.parse_args()
|
|
|
|
sys.path.insert(0, os.path.abspath(args.kronos_repo))
|
|
import torch
|
|
import pandas as pd
|
|
from model import Kronos, KronosTokenizer
|
|
|
|
torch.manual_seed(args.seed)
|
|
np.random.seed(args.seed)
|
|
torch.set_grad_enabled(False)
|
|
device = "cpu" # reproducible references only
|
|
|
|
os.makedirs(args.out, exist_ok=True)
|
|
manifest = {"arrays": {}, "scalars": {}, "notes": []}
|
|
|
|
def save(name, arr, dtype):
|
|
a = np.ascontiguousarray(np.asarray(arr).astype(dtype))
|
|
np.save(os.path.join(args.out, name + ".npy"), a)
|
|
a.tofile(os.path.join(args.out, name + ".bin"))
|
|
manifest["arrays"][name] = {
|
|
"file": name + ".bin",
|
|
"dtype": np.dtype(dtype).name,
|
|
"shape": list(a.shape),
|
|
"order": "little-endian, row-major (C)",
|
|
}
|
|
print(f" saved {name:24s} shape={tuple(a.shape)} dtype={np.dtype(dtype).name}")
|
|
|
|
# ------------------------------------------------------------------ load
|
|
print("Loading models (CPU, float32) ...")
|
|
tok = KronosTokenizer.from_pretrained(args.tokenizer_id).to(device).float().eval()
|
|
mdl = Kronos.from_pretrained(args.model_id).to(device).float().eval()
|
|
|
|
tok_cfg = dump_config(tok, [
|
|
"d_in", "d_model", "n_heads", "ff_dim", "enc_layers", "dec_layers",
|
|
"s1_bits", "s2_bits", "codebook_dim",
|
|
])
|
|
mdl_cfg = dump_config(mdl, [
|
|
"d_model", "n_heads", "ff_dim", "n_layers", "s1_bits", "s2_bits",
|
|
"s1_vocab_size", "s2_vocab_size", "learn_te",
|
|
])
|
|
with open(os.path.join(args.out, "config_tokenizer.json"), "w") as f:
|
|
json.dump(tok_cfg, f, indent=2)
|
|
with open(os.path.join(args.out, "config_predictor.json"), "w") as f:
|
|
json.dump(mdl_cfg, f, indent=2)
|
|
print("Pinned tokenizer config:", json.dumps(tok_cfg))
|
|
print("Pinned predictor config:", json.dumps(mdl_cfg))
|
|
|
|
# ------------------------------------------------- freeze an input window
|
|
L, P = args.lookback, args.pred_len
|
|
price_cols = ["open", "high", "low", "close"]
|
|
feat_cols = price_cols + ["volume", "amount"]
|
|
stamp_cols = ["minute", "hour", "weekday", "day", "month"]
|
|
|
|
idx = pd.date_range("2023-01-02 00:00:00", periods=L + P, freq="h")
|
|
rw = np.cumsum(np.random.randn(L + P).astype(np.float64)) * 0.5 + 100.0
|
|
close = rw
|
|
open_ = np.concatenate([[close[0]], close[:-1]])
|
|
high = np.maximum(open_, close) + np.abs(np.random.randn(L + P)) * 0.2
|
|
low = np.minimum(open_, close) - np.abs(np.random.randn(L + P)) * 0.2
|
|
volume = np.abs(np.random.randn(L + P)) * 1000.0 + 500.0
|
|
amount = volume * (open_ + high + low + close) / 4.0
|
|
df_all = pd.DataFrame(
|
|
{"open": open_, "high": high, "low": low, "close": close,
|
|
"volume": volume, "amount": amount}, index=idx)
|
|
|
|
hist_df = df_all.iloc[:L].copy()
|
|
x_ts = pd.Series(df_all.index[:L])
|
|
y_ts = pd.Series(df_all.index[L:L + P])
|
|
|
|
def stamps(ts):
|
|
out = np.zeros((len(ts), 5), dtype=np.float32)
|
|
out[:, 0] = ts.dt.minute
|
|
out[:, 1] = ts.dt.hour
|
|
out[:, 2] = ts.dt.weekday
|
|
out[:, 3] = ts.dt.day
|
|
out[:, 4] = ts.dt.month
|
|
return out
|
|
|
|
x_raw = hist_df[feat_cols].values.astype(np.float32) # (L, 6)
|
|
x_stamp = stamps(x_ts) # (L, 5)
|
|
y_stamp = stamps(y_ts) # (P, 5)
|
|
|
|
# --------------------------------------------- preprocessing (exact copy)
|
|
x_mean = x_raw.mean(axis=0) # (6,)
|
|
x_std = x_raw.std(axis=0) # (6,) population std, ddof=0
|
|
x_norm = (x_raw - x_mean) / (x_std + args.eps)
|
|
x_norm = np.clip(x_norm, -args.clip, args.clip).astype(np.float32)
|
|
|
|
print("\nStage 0 preprocessing")
|
|
save("x_raw", x_raw, np.float32)
|
|
save("x_norm", x_norm, np.float32)
|
|
save("x_mean", x_mean, np.float32)
|
|
save("x_std", x_std, np.float32)
|
|
save("x_stamp", x_stamp, np.float32)
|
|
save("y_stamp", y_stamp, np.float32)
|
|
|
|
xt = torch.from_numpy(x_norm)[None].to(device) # (1, L, 6)
|
|
xs = torch.from_numpy(x_stamp)[None].to(device) # (1, L, 5)
|
|
|
|
# ----------------------------------------------- Stage 1: tokenizer.encode
|
|
print("\nStage 1 tokenizer.encode (half=True -> [s1_ids, s2_ids])")
|
|
z_idx = tok.encode(xt, half=True)
|
|
s1_ids = z_idx[0] if isinstance(z_idx, (list, tuple)) else z_idx
|
|
s2_ids = z_idx[1] if isinstance(z_idx, (list, tuple)) else None
|
|
save("s1_ids", s1_ids.cpu().numpy(), np.int32)
|
|
if s2_ids is not None:
|
|
save("s2_ids", s2_ids.cpu().numpy(), np.int32)
|
|
|
|
# --------------------------------- Stage 2: predictor decode_s1 / decode_s2
|
|
print("\nStage 2 predictor logits (deterministic, teacher-forced)")
|
|
try:
|
|
s1_logits, context = mdl.decode_s1(s1_ids, s2_ids, xs)
|
|
save("s1_logits_full", s1_logits.cpu().numpy(), np.float32)
|
|
save("s1_logits_last", s1_logits[0, -1].cpu().numpy(), np.float32)
|
|
# deterministic condition for s2: argmax of the last-step s1 logits
|
|
s1_pick = s1_logits[:, -1:].argmax(dim=-1) # (1, 1)
|
|
save("s1_pick_last", s1_pick.cpu().numpy(), np.int32)
|
|
try:
|
|
s2_logits = mdl.decode_s2(context, s1_pick)
|
|
save("s2_logits_last", s2_logits.reshape(-1).cpu().numpy(), np.float32)
|
|
except Exception as e:
|
|
manifest["notes"].append(f"decode_s2 capture failed: {e!r} "
|
|
f"(check arg convention vs auto_regressive_inference)")
|
|
print(" [warn] decode_s2 failed:", repr(e))
|
|
except Exception as e:
|
|
manifest["notes"].append(f"decode_s1 capture failed: {e!r}")
|
|
print(" [warn] decode_s1 failed:", repr(e))
|
|
|
|
# ------------------------------------ Stage 3: tokenizer.decode round-trip
|
|
print("\nStage 3 tokenizer.decode (reconstruction of the context window)")
|
|
try:
|
|
recon = tok.decode([s1_ids, s2_ids], half=True) # (1, L, 6) normalized
|
|
recon_np = recon[0].cpu().numpy()
|
|
save("recon_norm", recon_np, np.float32)
|
|
save("recon_denorm", recon_np * (x_std + args.eps) + x_mean, np.float32)
|
|
except Exception as e:
|
|
manifest["notes"].append(f"tokenizer.decode capture failed: {e!r}")
|
|
print(" [warn] tokenizer.decode failed:", repr(e))
|
|
|
|
# ------------------------------- Stage 4: seeded end-to-end forecast (opt.)
|
|
if not args.skip_forecast:
|
|
print("\nStage 4 seeded end-to-end forecast (STOCHASTIC, seed-pinned)")
|
|
try:
|
|
from model import KronosPredictor
|
|
torch.manual_seed(args.seed)
|
|
predictor = KronosPredictor(mdl, tok, device=device,
|
|
max_context=512, clip=args.clip)
|
|
pred_df = predictor.predict(
|
|
df=hist_df, x_timestamp=x_ts, y_timestamp=y_ts,
|
|
pred_len=P, T=1.0, top_k=0, top_p=0.9,
|
|
sample_count=1, verbose=False)
|
|
save("forecast_seeded", pred_df[feat_cols].values, np.float32)
|
|
manifest["notes"].append(
|
|
"forecast_seeded uses torch.manual_seed(seed), sample_count=1, "
|
|
"T=1.0, top_p=0.9 -- reproducible in PyTorch but NOT a bit-exact "
|
|
"MQL5 target; use the logits-level refs for exact verification.")
|
|
except Exception as e:
|
|
manifest["notes"].append(f"seeded forecast failed: {e!r}")
|
|
print(" [warn] forecast failed:", repr(e))
|
|
|
|
# ----------------------------------------------------------- write manifest
|
|
manifest["scalars"] = {
|
|
"lookback": L, "pred_len": P, "clip": args.clip, "eps": args.eps,
|
|
"seed": args.seed, "feature_order": feat_cols, "stamp_order": stamp_cols,
|
|
"std": "population (ddof=0)",
|
|
"tokenizer_config": tok_cfg, "predictor_config": mdl_cfg,
|
|
}
|
|
with open(os.path.join(args.out, "manifest.json"), "w") as f:
|
|
json.dump(manifest, f, indent=2)
|
|
|
|
print(f"\nDone. {len(manifest['arrays'])} arrays written to {args.out}")
|
|
if manifest["notes"]:
|
|
print("Notes / warnings:")
|
|
for n in manifest["notes"]:
|
|
print(" -", n)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|