src/adapters/kronos.py

"""Kronos -- a foundation model trained on candlesticks rather than on numbers.

Kronos is the only family here that is native to OHLCV: it tokenises whole
candles and generates whole candles, so it is the family that can emit sampled
price *paths* rather than just a band on the close. That is what drives the
ghost-path view and the sampled-path dispersion widget.

Two facts about it shape this adapter.

**It has no inference package.** Weights ship on the Hub, the code ships only
on GitHub, and the `kronos` name on PyPI is an unrelated Django library. The
source is vendored under `vendor/kronos` at a pinned commit; see the
PROVENANCE.md there.

**It has no seed argument.** Sampling runs through `torch.multinomial` against
the global RNG. Determinism is therefore something this adapter imposes, by
seeding immediately before inference, rather than something Kronos offers.
"""

from __future__ import annotations

import logging
import sys
from pathlib import Path

import numpy as np
import pandas as pd

from .. import config
from .base import (OHLCV_COLUMNS, OUTPUT_OHLCV_PATHS, AdapterError,
                   Capabilities, ForecastAdapter, ForecastResult,
                   check_context, seed_everything)

log = logging.getLogger("arena.adapters.kronos")

_VENDOR = Path(__file__).resolve().parents[2] / "vendor"
if str(_VENDOR) not in sys.path:
    sys.path.insert(0, str(_VENDOR))

# Each Kronos checkpoint is trained against one tokenizer and is meaningless
# with any other, so the pairing is fixed here rather than left to the caller.
# Context lengths are the model card's, not guesses.
KRONOS_MODELS = {
    "NeoQuasar/Kronos-mini": {
        "tokenizer": "NeoQuasar/Kronos-Tokenizer-2k",
        "max_context": 2048,
        "params": "4.1M",
        # Measured 6.7-7.1s warm at h=24, n=32 on an M-series laptop. That is
        # inside the 8s budget here but cpu-basic is slower, so the tier is
        # confirmed by `scripts/benchmark.py` running on the Space itself and
        # demoted in the registry if it misses.
        "hardware": "cpu",
    },
    "NeoQuasar/Kronos-small": {
        "tokenizer": "NeoQuasar/Kronos-Tokenizer-base",
        "max_context": 512,
        "params": "24.7M",
        # Measured 23.5-24.3s warm at h=24, n=32 on an M-series laptop, which
        # is faster than cpu-basic. Nowhere near the 8s warm budget.
        "hardware": "gpu",
    },
    "NeoQuasar/Kronos-base": {
        "tokenizer": "NeoQuasar/Kronos-Tokenizer-base",
        "max_context": 512,
        "params": "102.3M",
        # 102M params generating autoregressively is not a cpu-basic workload.
        "hardware": "gpu",
    },
}

# Sampling defaults. These are the model card's own recommended values; they
# are recorded in `component_versions` because changing them changes every
# number the model produces.
DEFAULT_T = 1.0
DEFAULT_TOP_P = 0.9
DEFAULT_TOP_K = 0
CLIP = 5


class KronosAdapter(ForecastAdapter):
    """`NeoQuasar/Kronos-*` candlestick generators."""

    family = "kronos"
    adapter_version = "1"

    def __init__(self, model_id: str, revision: str | None = None,
                 device: str | None = None, tokenizer_id: str | None = None,
                 tokenizer_revision: str | None = None,
                 hardware: str | None = None):
        super().__init__(model_id, revision=revision, device=device)
        spec = KRONOS_MODELS.get(model_id, {})
        self.tokenizer_id = tokenizer_id or spec.get("tokenizer")
        self.tokenizer_revision = tokenizer_revision
        self._max_context = int(spec.get("max_context", 512))
        # The declared default, which the registry may override downward after
        # a measured latency run. Declaring a tier a model cannot hold is how
        # a user ends up watching a spinner for ninety seconds.
        self._hardware = hardware or spec.get("hardware", "gpu")
        self._tokenizer = None
        self._resolved_tokenizer_revision: str | None = None

    # -- capabilities -----------------------------------------------------

    def capabilities(self) -> Capabilities:
        return Capabilities(
            output=OUTPUT_OHLCV_PATHS,
            # Kronos generates autoregressively -- one forward pass per
            # forecast step. The tier here is the declared default; a measured
            # miss against the CPU budget demotes it in the registry.
            hardware=self._hardware,
            max_context=self._max_context,
            asset_generality="financial",
            seedable_natively=False,
        )

    def component_versions(self) -> dict[str, str]:
        versions = {
            "model": f"{self.model_id}@{self.resolved_revision}",
            "tokenizer": f"{self.tokenizer_id}@{self._resolved_tokenizer_revision or 'unpinned'}",
            "sampling": f"T={DEFAULT_T},top_p={DEFAULT_TOP_P},top_k={DEFAULT_TOP_K},clip={CLIP}",
            "vendor": "kronos@67b630e6",
        }
        try:
            import torch

            versions["torch"] = torch.__version__
        except ImportError:
            pass
        return versions

    # -- load -------------------------------------------------------------

    def load(self, model_id: str | None = None, revision: str | None = None):
        if model_id and model_id != self.model_id:
            self.model_id = model_id
            spec = KRONOS_MODELS.get(model_id, {})
            self.tokenizer_id = spec.get("tokenizer", self.tokenizer_id)
            self._max_context = int(spec.get("max_context", self._max_context))
            self._hardware = spec.get("hardware", self._hardware)
            self._model = None
        if revision:
            self.revision = revision
        if self._model is not None:
            return self

        if not self.tokenizer_id:
            raise AdapterError(
                f"no tokenizer known for {self.model_id}; Kronos checkpoints are "
                f"only valid with the tokenizer they were trained against"
            )

        try:
            from kronos import Kronos, KronosTokenizer
        except ImportError as e:  # pragma: no cover
            raise AdapterError(
                "the vendored Kronos source is not importable; check vendor/kronos"
            ) from e

        self.resolve_revision()
        self._resolved_tokenizer_revision = _resolve(self.tokenizer_id,
                                                     self.tokenizer_revision)

        self._tokenizer = KronosTokenizer.from_pretrained(
            self.tokenizer_id, revision=self._resolved_tokenizer_revision)
        self._model = Kronos.from_pretrained(
            self.model_id, revision=self._resolved_revision)

        self._tokenizer = self._tokenizer.to(self.device).eval()
        self._model = self._model.to(self.device).eval()
        return self

    # -- predict ----------------------------------------------------------

    def predict(self, context_ohlcv: pd.DataFrame, horizon: int,
                n_samples: int = config.DEFAULT_N_SAMPLES, seed: int = 0,
                issued_ts: pd.Timestamp | None = None) -> ForecastResult:
        check_context(context_ohlcv, issued_ts=issued_ts)
        if horizon < 1:
            raise AdapterError("horizon must be at least 1")
        if n_samples < 1:
            raise AdapterError("n_samples must be at least 1")

        if self._model is None:
            self.load()

        import torch

        from kronos.kronos import auto_regressive_inference, calc_time_stamps

        ctx = self._trim(context_ohlcv).reset_index(drop=True)
        ts = pd.to_datetime(ctx["ts"], utc=True)

        # Kronos consumes `amount` (quote volume) alongside OHLCV. Where the
        # cache has no amount column, the upstream predictor's own fallback is
        # volume times the mean price, and that is reproduced here so the
        # inputs match what `KronosPredictor` would have built.
        frame = ctx[list(OHLCV_COLUMNS)].astype("float64").copy()
        if "amount" in ctx.columns:
            frame["amount"] = ctx["amount"].astype("float64")
        else:
            frame["amount"] = frame["volume"] * frame[["open", "high", "low", "close"]].mean(axis=1)

        x = frame.to_numpy(dtype="float32")
        future_ts = _future_timestamps(ts, horizon)

        # Normalisation, replicated from `KronosPredictor.predict`. It is
        # replicated rather than called because the upstream method averages
        # the sampled paths away before returning, and the paths are the point.
        # `test_kronos_vendor.py` asserts this reproduces the upstream result.
        x_mean = x.mean(axis=0)
        x_std = x.std(axis=0)
        x_norm = (x - x_mean) / (x_std + 1e-5)
        x_norm = np.clip(x_norm, -CLIP, CLIP)

        x_stamp = calc_time_stamps(ts.reset_index(drop=True)).to_numpy(dtype="float32")
        y_stamp = calc_time_stamps(pd.Series(future_ts)).to_numpy(dtype="float32")

        x_t = torch.from_numpy(x_norm[np.newaxis, :].astype("float32")).to(self.device)
        xs_t = torch.from_numpy(x_stamp[np.newaxis, :]).to(self.device)
        ys_t = torch.from_numpy(y_stamp[np.newaxis, :]).to(self.device)

        # The seed goes in here and nowhere else: Kronos draws through the
        # global torch RNG, so this call is what makes the result reproducible.
        seed_everything(seed)

        with torch.inference_mode():
            raw = auto_regressive_inference(
                self._tokenizer, self._model, x_t, xs_t, ys_t,
                max_context=self._max_context, pred_len=horizon, clip=CLIP,
                T=DEFAULT_T, top_k=DEFAULT_TOP_K, top_p=DEFAULT_TOP_P,
                sample_count=int(n_samples), verbose=False,
                return_paths=True,
            )

        # (batch=1, n_samples, seq, features) -> (n_samples, horizon, features)
        paths = np.asarray(raw)[0][:, -horizon:, :]
        paths = paths * (x_std + 1e-5) + x_mean

        # A sampled candle can come back internally inconsistent -- a high
        # below the close, say -- because each field is decoded from its own
        # token. Repairing it is more honest than rendering an impossible
        # candle, and it only ever widens the bar to contain what it must.
        paths = _repair_candles(paths)

        close = paths[:, :, OHLCV_COLUMNS.index("close")]
        quantiles = self._quantiles_from_paths(close)

        return ForecastResult(
            quantiles=quantiles,
            levels=config.QUANTILE_LEVELS,
            horizon=horizon,
            context_len=len(ctx),
            inference_version=self.inference_version(),
            seed=int(seed),
            n_samples=int(n_samples),
            paths=paths[:, :, :len(OHLCV_COLUMNS)],
        )


# --------------------------------------------------------------------------
# Helpers
# --------------------------------------------------------------------------


def _resolve(repo_id: str, revision: str | None) -> str:
    from huggingface_hub import HfApi

    return HfApi().model_info(repo_id, revision=revision).sha


def _future_timestamps(ts: pd.Series, horizon: int) -> pd.DatetimeIndex:
    """Continue the context's own cadence forward.

    Kronos conditions on calendar features, so the future stamps have to be
    plausible rather than arbitrary. The modal spacing of the context is used
    so that this works for both 1h and 1d without being told which it is.
    """
    if len(ts) < 2:
        raise AdapterError("cannot infer cadence from fewer than two bars")
    deltas = ts.diff().dropna()
    step = deltas.mode().iloc[0] if len(deltas.mode()) else deltas.median()
    last = ts.iloc[-1]
    return pd.DatetimeIndex([last + step * (i + 1) for i in range(horizon)])


def _repair_candles(paths: np.ndarray) -> np.ndarray:
    """Force high >= max(o,c) and low <= min(o,c) on every sampled candle."""
    o, h, l, c = (OHLCV_COLUMNS.index(k) for k in ("open", "high", "low", "close"))
    body_hi = np.maximum(paths[:, :, o], paths[:, :, c])
    body_lo = np.minimum(paths[:, :, o], paths[:, :, c])
    paths[:, :, h] = np.maximum(paths[:, :, h], body_hi)
    paths[:, :, l] = np.minimum(paths[:, :, l], body_lo)
    # Volume is a count; a negative one is a decode artefact, not information.
    v = OHLCV_COLUMNS.index("volume")
    paths[:, :, v] = np.maximum(paths[:, :, v], 0.0)
    return paths