kronos-mql5/Kronos_Python/kronos_slide_capture.py

#!/usr/bin/env python3
"""
kronos_slide_capture.py
=======================
Golden reference for the AR loop's WINDOW-SLIDE branch (buffer > max_context).

The main verification used L=256, pred_len=16, so the 512-token ring buffer never
filled and the roll-left/slide path never executed. Here we force it: a long
context (L0) plus a horizon (pred_len) whose sum exceeds max_context, run with
GREEDY decoding (deterministic), and capture the exact generated token sequence.
MQL5 greedy AR must reproduce gen_s1/gen_s2 token-for-token through the slide.

We replicate the source auto_regressive_inference loop EXACTLY (same windowing,
same buffer roll), but force greedy (sample_logits=False) so it is reproducible.

    python kronos_slide_capture.py --kronos_repo ./kronos-git --out ./kronos_refs_slide

Outputs (little-endian, row-major, for MQL5 FileReadArray):
  x_norm_slide.bin    [L0, 6]        float32   normalized long context
  x_stamp_slide.bin   [L0, 5]        float32
  y_stamp_slide.bin   [pred_len, 5]  float32
  gen_s1_slide.bin    [pred_len]     int32     greedy s1 tokens
  gen_s2_slide.bin    [pred_len]     int32     greedy s2 tokens
  manifest_slide.json
"""
import argparse, json, os, sys
import numpy as np


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--kronos_repo", default="./kronos-git")
    ap.add_argument("--out", default="./kronos_refs_slide")
    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=500, help="long context, <=512")
    ap.add_argument("--pred_len", type=int, default=30, help="so L0+pred_len > 512")
    ap.add_argument("--max_context", type=int, default=512)
    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)
    args = ap.parse_args()

    sys.path.insert(0, os.path.abspath(args.kronos_repo))
    import torch
    import pandas as pd
    import torch.nn.functional as F
    from model import Kronos, KronosTokenizer

    torch.manual_seed(args.seed); np.random.seed(args.seed)
    torch.set_grad_enabled(False)
    device = "cpu"
    os.makedirs(args.out, exist_ok=True)

    L, P, MC = args.lookback, args.pred_len, args.max_context
    assert L + P > MC, f"need L0+pred_len > max_context to exercise the slide ({L}+{P} vs {MC})"

    tok = KronosTokenizer.from_pretrained(args.tokenizer_id).to(device).float().eval()
    mdl = Kronos.from_pretrained(args.model_id).to(device).float().eval()

    # ---- synth a reproducible long window (same recipe as the main capture) ----
    feat_cols = ["open", "high", "low", "close", "volume", "amount"]
    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 = pd.DataFrame({"open": open_, "high": high, "low": low, "close": close,
                       "volume": volume, "amount": amount}, index=idx)

    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 = df[feat_cols].values[:L].astype(np.float32)
    x_stamp = stamps(pd.Series(df.index[:L]))
    y_stamp = stamps(pd.Series(df.index[L:L + P]))

    x_mean = x_raw.mean(axis=0); x_std = x_raw.std(axis=0)
    x_norm = np.clip((x_raw - x_mean) / (x_std + args.eps), -args.clip, args.clip).astype(np.float32)

    xt = torch.from_numpy(x_norm)[None]            # (1,L,6)
    xs = torch.from_numpy(x_stamp)[None]           # (1,L,5)
    ys = torch.from_numpy(y_stamp)[None]           # (1,P,5)

    # ---- replicate auto_regressive_inference loop, GREEDY (sample_logits=False) ----
    x_token = tok.encode(xt, half=True)
    initial_seq_len = L
    full_stamp = torch.cat([xs, ys], dim=1)        # (1, L+P, 5)

    pre_buffer = x_token[0].new_zeros(1, MC)
    post_buffer = x_token[1].new_zeros(1, MC)
    buffer_len = min(initial_seq_len, MC)
    start_idx = max(0, initial_seq_len - MC)
    pre_buffer[:, :buffer_len] = x_token[0][:, start_idx:start_idx + buffer_len]
    post_buffer[:, :buffer_len] = x_token[1][:, start_idx:start_idx + buffer_len]

    gen_s1 = np.zeros(P, dtype=np.int32)
    gen_s2 = np.zeros(P, dtype=np.int32)

    def greedy(logits):
        return torch.topk(F.softmax(logits, dim=-1), k=1, dim=-1)[1]   # argmax id, (B,1)

    crossed = -1
    for i in range(P):
        current_seq_len = initial_seq_len + i
        window_len = min(current_seq_len, MC)
        if current_seq_len <= MC:
            input_tokens = [pre_buffer[:, :window_len], post_buffer[:, :window_len]]
        else:
            input_tokens = [pre_buffer, post_buffer]
            if crossed < 0:
                crossed = i

        ctx_end = current_seq_len
        ctx_start = max(0, ctx_end - MC)
        current_stamp = full_stamp[:, ctx_start:ctx_end, :].contiguous()

        s1_logits, context = mdl.decode_s1(input_tokens[0], input_tokens[1], current_stamp)
        sample_pre = greedy(s1_logits[:, -1, :])
        s2_logits = mdl.decode_s2(context, sample_pre)
        sample_post = greedy(s2_logits[:, -1, :])

        gen_s1[i] = int(sample_pre.item())
        gen_s2[i] = int(sample_post.item())

        if current_seq_len < MC:
            pre_buffer[:, current_seq_len] = sample_pre.squeeze(-1)
            post_buffer[:, current_seq_len] = sample_post.squeeze(-1)
        else:
            pre_buffer.copy_(torch.roll(pre_buffer, shifts=-1, dims=1))
            post_buffer.copy_(torch.roll(post_buffer, shifts=-1, dims=1))
            pre_buffer[:, -1] = sample_pre.squeeze(-1)
            post_buffer[:, -1] = sample_post.squeeze(-1)

    print(f"slide branch first fires at generation step i={crossed} "
          f"(current_seq_len crosses {MC})")
    print("gen_s1:", gen_s1.tolist())
    print("gen_s2:", gen_s2.tolist())

    def save(name, arr, dtype):
        a = np.ascontiguousarray(np.asarray(arr).astype(dtype)); a.tofile(os.path.join(args.out, name + ".bin"))
        return {"file": name + ".bin", "dtype": np.dtype(dtype).name, "shape": list(a.shape)}

    man = {"arrays": {}, "scalars": {"lookback": L, "pred_len": P, "max_context": MC,
                                     "slide_first_step": int(crossed), "decoding": "greedy"}}
    man["arrays"]["x_norm_slide"]  = save("x_norm_slide",  x_norm,  np.float32)
    man["arrays"]["x_stamp_slide"] = save("x_stamp_slide", x_stamp, np.float32)
    man["arrays"]["y_stamp_slide"] = save("y_stamp_slide", y_stamp, np.float32)
    man["arrays"]["gen_s1_slide"]  = save("gen_s1_slide",  gen_s1,  np.int32)
    man["arrays"]["gen_s2_slide"]  = save("gen_s2_slide",  gen_s2,  np.int32)
    with open(os.path.join(args.out, "manifest_slide.json"), "w") as f:
        json.dump(man, f, indent=2)
    print(f"\nDone. Wrote slide refs to {args.out}")


if __name__ == "__main__":
    main()