Initial import: grid-bot — grid trading bot for BTC-USDT on Cifra Markets

kronos/__init__.py

@@ -0,0 +1,21 @@
+"""Kronos advisor for grid-bot: OHLCV → step/center/pause advice.
+
+Public surface:
+  - GridAdvice
+  - KronosAdvisor
+  - TradernetHlocSource
+  - BinancePublicSource
+  - run_backtest (in backtest.py)
+"""
+
+from .advisor import GridAdvice, KronosAdvisor, MIN_STEP_PCT, MAX_STEP_PCT
+from .data_source import TradernetHlocSource, BinancePublicSource
+
+__all__ = [
+    "GridAdvice",
+    "KronosAdvisor",
+    "MIN_STEP_PCT",
+    "MAX_STEP_PCT",
+    "TradernetHlocSource",
+    "BinancePublicSource",
+]

kronos/advisor.py

@@ -0,0 +1,278 @@
+"""
+KronosAdvisor: turn OHLCV history into grid-trading advice.
+
+Input:  DataFrame with columns ['open','high','low','close','volume'] + timestamps
+Output: GridAdvice(step_percent, center_offset_pct, pause_grid, confidence, ...)
+
+Design notes:
+  - step_percent is clamped to [MIN_STEP_PCT, MAX_STEP_PCT] to prevent silly values
+  - pause_grid = True if expected_range_pct is huge (>2*baseline) or bias is strong
+  - confidence = a heuristic 0..1 — how much we trust this forecast
+  - This module does NOT touch money. It only computes numbers.
+"""
+
+from __future__ import annotations
+
+import math
+import os
+import sys
+import logging
+from dataclasses import dataclass, asdict
+from datetime import datetime, timezone
+from typing import Optional
+
+import numpy as np
+import pandas as pd
+
+logger = logging.getLogger(__name__)
+
+# Safety clamps for step_percent. 0.1% — слишком мелко (съест комиссия),
+# 5% — слишком грубо для grid.
+MIN_STEP_PCT = 0.001   # 0.1%
+MAX_STEP_PCT = 0.05    # 5.0%
+
+# Thresholds
+PAUSE_RANGE_PCT = 0.05      # expected_range > 5% → pause (тренд)
+PAUSE_BIAS_THRESHOLD = 0.6  # strong directional bias → pause
+
+
+@dataclass
+class GridAdvice:
+    step_percent: float        # адаптивный шаг сетки (доля)
+    center_offset_pct: float   # сдвиг центра от текущей цены (доля, + вверх, - вниз)
+    pause_grid: bool           # True = стоп сетки (тренд / пробой)
+    confidence: float          # 0..1
+    forecast_horizon_min: int  # таймфрейм прогноза (минуты)
+    expected_range_pct: float  # (high-low)/close прогноза
+    bias: str                  # "up" | "down" | "flat"
+    source: str                # "kronos-mini" | "kronos-small" | "fallback"
+    generated_at: str          # ISO timestamp
+
+    def to_dict(self) -> dict:
+        return asdict(self)
+
+    def to_json_safe(self) -> str:
+        import json
+        return json.dumps(self.to_dict(), indent=2, ensure_ascii=False)
+
+
+class KronosAdvisor:
+    """
+    Загружает Kronos (mini или small) и считает advice.
+
+    Пример:
+        advisor = KronosAdvisor(model_name="kronos-mini", device="cpu")
+        advice = advisor.advise(ohlcv_df, lookback=400, pred_len=24, tf_min=60)
+        print(advice)
+    """
+
+    def __init__(
+        self,
+        model_name: str = "kronos-mini",
+        device: str = "cpu",
+        max_context: int = 512,
+        cache_dir: Optional[str] = None,
+    ):
+        self.model_name = model_name
+        self.device = device
+        self.max_context = max_context
+        self.cache_dir = cache_dir
+        self._predictor = None
+        self._tokenizer = None
+        self._loaded = False
+        self._load_error: Optional[str] = None
+
+    def _ensure_loaded(self):
+        if self._loaded:
+            return
+        if self._load_error:
+            raise RuntimeError(f"Kronos not loaded: {self._load_error}")
+
+        # Попытка 1: уже в sys.path (например, вызвали из ~/projects/Kronos)
+        # Попытка 2: лежит в соседнем каталоге — ищем Kronos/ вверх по дереву
+        import pathlib
+        try:
+            from model import Kronos, KronosTokenizer, KronosPredictor
+        except ImportError:
+            kronos_dir = None
+            cur = pathlib.Path(__file__).resolve().parent
+            for _ in range(6):  # до 6 уровней вверх
+                candidate = cur / "Kronos"
+                if candidate.is_dir() and (candidate / "model").is_dir():
+                    kronos_dir = candidate
+                    break
+                cur = cur.parent
+            if kronos_dir is not None:
+                sys.path.insert(0, str(kronos_dir))
+            try:
+                from model import Kronos, KronosTokenizer, KronosPredictor
+            except ImportError as e:
+                self._load_error = f"cannot import Kronos model: {e}"
+                raise RuntimeError(self._load_error)
+
+        # Выбор модели. tokenizer для mini и small разный.
+        if self.model_name == "kronos-mini":
+            tokenizer_name = os.path.abspath("models/Kronos-Tokenizer-2k")
+            model_hf = os.path.abspath("models/Kronos-mini")
+        elif self.model_name == "kronos-small":
+            tokenizer_name = os.path.abspath("models/Kronos-Tokenizer-base")
+            model_hf = os.path.abspath("models/Kronos-small")
+            if not pathlib.Path(tokenizer_name).exists():
+                tokenizer_name = os.path.abspath("models/Kronos-Tokenizer-2k")
+                model_hf = os.path.abspath("models/Kronos-mini")
+        else:
+            raise ValueError(f"Unknown model_name: {self.model_name!r}")
+
+        logger.info("Loading Kronos tokenizer %s ...", tokenizer_name)
+        self._tokenizer = KronosTokenizer.from_pretrained(tokenizer_name)
+        logger.info("Loading Kronos model %s ...", model_hf)
+        self._model = Kronos.from_pretrained(model_hf)
+
+        # max_context=512 для small/base. mini поддерживает 2048.
+        ctx = 2048 if self.model_name == "kronos-mini" else min(512, self.max_context)
+        self._predictor = KronosPredictor(self._model, self._tokenizer, max_context=ctx, device=self.device)
+        self._loaded = True
+        logger.info("Kronos %s loaded on %s (max_context=%d)", self.model_name, self.device, ctx)
+
+    def _safe_predict(self, df: pd.DataFrame, lookback: int, pred_len: int, tf_min: int):
+        """Прогноз с несколькими sample_count, усреднение."""
+        self._ensure_loaded()
+
+        # Обрезаем lookback до того, что у нас есть
+        actual_lookback = min(lookback, len(df) - 1)
+        if actual_lookback < 32:
+            raise ValueError(f"need at least 32 candles, got {len(df)}")
+
+        x_df = df.iloc[-actual_lookback:][["open", "high", "low", "close", "volume"]].copy()
+        # Если нет volume — заполняем нулями (модель это умеет)
+        if "volume" not in x_df.columns:
+            x_df["volume"] = 0.0
+
+        # timestamps — DatetimeIndex или Series. Kronos ожидает pandas Series (нужен .dt).
+        if isinstance(df.index, pd.DatetimeIndex):
+            x_ts = pd.Series(df.index[-actual_lookback:], name="timestamps")
+            last_ts = df.index[-1]
+        elif "timestamps" in df.columns:
+            ts = pd.to_datetime(df["timestamps"].iloc[-actual_lookback:])
+            x_ts = pd.Series(ts.values, name="timestamps")
+            last_ts = pd.to_datetime(df["timestamps"].iloc[-1])
+        else:
+            raise ValueError("df must have DatetimeIndex or 'timestamps' column")
+
+        y_ts = pd.Series(pd.date_range(start=last_ts, periods=pred_len + 1, freq=f"{tf_min}min")[1:], name="timestamps")
+
+        pred_df = self._predictor.predict(
+            df=x_df,
+            x_timestamp=x_ts,
+            y_timestamp=y_ts,
+            pred_len=pred_len,
+            T=1.0,
+            top_p=0.9,
+            sample_count=1,   # одна траектория — быстрее (для бэктеста)
+            verbose=False,
+        )
+        return pred_df, x_df
+
+    def advise(
+        self,
+        df: pd.DataFrame,
+        lookback: int = 400,
+        pred_len: int = 24,
+        tf_min: int = 60,
+    ) -> GridAdvice:
+        """
+        Главный метод. Возвращает GridAdvice.
+        При любой ошибке — fallback: дефолтный шаг, без паузы, confidence=0.
+        """
+        try:
+            pred_df, x_df = self._safe_predict(df, lookback, pred_len, tf_min)
+            advice = self._postprocess(pred_df, x_df, pred_len, tf_min)
+            return advice
+        except Exception as e:
+            logger.exception("Kronos advise failed: %s", e)
+            return self._fallback_advice(tf_min, reason=str(e))
+
+    # -----------------------------------------------------------------------
+    # Внутренние штуки
+    # -----------------------------------------------------------------------
+
+    def _postprocess(self, pred_df: pd.DataFrame, x_df: pd.DataFrame, pred_len: int, tf_min: int) -> GridAdvice:
+        """Превращаем прогноз Kronos в GridAdvice."""
+        # Базовая волатильность по обучающей выборке: среднее |log return|
+        rets = np.log(x_df["close"] / x_df["close"].shift(1)).dropna()
+        baseline_vol = float(rets.abs().mean() or 1e-4)
+
+        # Метрики прогноза
+        f_open = float(pred_df["open"].iloc[0])
+        f_close = float(pred_df["close"].iloc[-1])
+        f_high = float(pred_df["high"].max())
+        f_low = float(pred_df["low"].min())
+        f_range_pct = (f_high - f_low) / f_close if f_close > 0 else 0.0
+        bias = (f_close - f_open) / f_open if f_open > 0 else 0.0
+
+        # 1) Адаптивный шаг: step ≈ 0.5 * expected_range / sqrt(pred_len)
+        #    Чтобы один «квант» сетки укладывался в ~половину прогнозного диапазона
+        #    за pred_len свечей.
+        raw_step = (f_range_pct * 0.5) / max(1.0, math.sqrt(pred_len))
+        # Также учитываем baseline vol (не дать шагу схлопнуться на низковолатильных участках)
+        raw_step = max(raw_step, baseline_vol * 1.5)
+        step = float(np.clip(raw_step, MIN_STEP_PCT, MAX_STEP_PCT))
+
+        # 2) Center offset: лёгкий сдвиг в сторону bias, но не больше половины диапазона
+        if f_range_pct > 0:
+            center_offset = float(np.clip(bias * 0.5, -f_range_pct * 0.25, f_range_pct * 0.25))
+        else:
+            center_offset = 0.0
+
+        # 3) Pause detection
+        bias_strength = abs(bias) / max(f_range_pct, 1e-4)
+        pause = (
+            (f_range_pct > PAUSE_RANGE_PCT) or
+            (bias_strength > PAUSE_BIAS_THRESHOLD)
+        )
+
+        # 4) Confidence — евристика
+        #    Чем уже ожидаемый диапазон относительно baseline → тем менее уверены
+        #    Чем ближе bias к нулю → тем менее уверены
+        conf = 1.0
+        if baseline_vol > 0:
+            ratio = f_range_pct / (baseline_vol * math.sqrt(pred_len))
+            # ratio < 1 → прогноз «теснее» исторического → понижаем уверенность
+            conf *= float(np.clip(ratio, 0.3, 1.2)) / 1.2
+        conf *= 1.0 - min(0.5, bias_strength * 0.5)
+        conf = float(np.clip(conf, 0.05, 0.95))
+
+        # 5) bias label
+        if abs(bias) < 0.001:
+            bias_label = "flat"
+        elif bias > 0:
+            bias_label = "up"
+        else:
+            bias_label = "down"
+
+        return GridAdvice(
+            step_percent=step,
+            center_offset_pct=center_offset,
+            pause_grid=bool(pause),
+            confidence=conf,
+            forecast_horizon_min=tf_min * pred_len,
+            expected_range_pct=f_range_pct,
+            bias=bias_label,
+            source=self.model_name,
+            generated_at=datetime.now(timezone.utc).isoformat(),
+        )
+
+    def _fallback_advice(self, tf_min: int, reason: str = "") -> GridAdvice:
+        """Если Kronos не справился — отдаём безопасный fallback."""
+        logger.warning("Kronos fallback: %s", reason)
+        return GridAdvice(
+            step_percent=0.005,            # 0.5% — как в дефолте .env
+            center_offset_pct=0.0,
+            pause_grid=False,
+            confidence=0.0,
+            forecast_horizon_min=tf_min * 24,
+            expected_range_pct=0.0,
+            bias="flat",
+            source="fallback",
+            generated_at=datetime.now(timezone.utc).isoformat(),
+        )

kronos/backtest.py

@@ -0,0 +1,298 @@
+"""
+Backtest: static grid vs adaptive grid (with Kronos advisor).
+
+Использование:
+    python -m kronos.backtest --csv ./data/BTCUSDT-1h.csv --tf 60 --step 0.005 --levels 10
+
+Что делает:
+  1) Загружает историю.
+  2) На каждом шаге t (от lookback до конца истории):
+       - запускает Kronos на окне [t-lookback .. t] с прогнозом на pred_len свечей
+       - вычисляет adaptive step
+       - симулирует сделку на следующих M свечей:
+            * static  — фиксированный шаг, центр = close[t]
+            * adaptive — шаг от Kronos, центр = close[t] + center_offset
+  3) Сравнивает метрики: PnL, max drawdown, число сделок, % прибыльных.
+"""
+
+from __future__ import annotations
+
+import argparse
+import logging
+import math
+import sys
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Optional
+
+import numpy as np
+import pandas as pd
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class BTResult:
+    name: str
+    total_pnl_pct: float
+    n_trades: int
+    winrate: float
+    max_dd_pct: float
+    sharpe: float
+
+
+def simulate_grid(
+    df: pd.DataFrame,
+    start: int,
+    horizon: int,
+    step_pct: float,
+    levels: int,
+    qty: float = 0.001,
+    center_offset_pct: float = 0.0,
+    fee_pct: float = 0.001,   # 0.1% taker fee per order (типично для Tradernet)
+) -> tuple[float, int, int]:
+    """
+    Простая сеточная симуляция с реалистичной логикой матчинга.
+
+    Правила:
+      - Одна сделка на свечу (BUY или SELL, не обе).
+      - BUY разрешён, если pos_open is None.
+      - SELL разрешён, если pos_open is not None.
+      - При исполнении снимаем комиссию (fee_pct × qty × price).
+      - Возвращаем (pnl_pct, n_roundtrips, n_winning_roundtrips).
+    """
+    if start >= len(df) - 1:
+        return 0.0, 0, 0
+
+    center = float(df["close"].iloc[start]) * (1.0 + center_offset_pct)
+    step = center * step_pct
+    if step <= 0:
+        return 0.0, 0, 0
+
+    buy_levels = sorted([center - n * step for n in range(1, levels + 1) if center - n * step > 0])
+    sell_levels = sorted([center + n * step for n in range(1, levels + 1)])
+
+    end = min(start + horizon, len(df))
+    pnl = 0.0
+    n_roundtrips = 0
+    n_wins = 0
+    pos_open: Optional[float] = None
+
+    for i in range(start + 1, end):
+        hi = float(df["high"].iloc[i])
+        lo = float(df["low"].iloc[i])
+
+        if pos_open is None:
+            # Ищем ближайший BUY-уровень, который достигнут
+            for bl in buy_levels:
+                if lo <= bl:
+                    # fee на вход
+                    pnl -= bl * qty * fee_pct
+                    pos_open = bl
+                    break
+        else:
+            # Ищем ближайший SELL-уровень, который достигнут
+            for sl in sell_levels:
+                if hi >= sl:
+                    gross = (sl - pos_open) * qty
+                    fee = sl * qty * fee_pct + pos_open * qty * fee_pct
+                    net = gross - fee
+                    pnl += net
+                    if net > 0:
+                        n_wins += 1
+                    n_roundtrips += 1
+                    pos_open = None
+                    break
+
+    # PnL в долях от notional (center * qty)
+    pnl_pct = pnl / (center * qty) if center > 0 else 0.0
+    return pnl_pct, n_roundtrips, n_wins
+
+
+def run_backtest(
+    df: pd.DataFrame,
+    tf_min: int,
+    lookback: int,
+    pred_len: int,
+    static_step_pct: float,
+    levels: int,
+    advisor=None,
+    stride: int = 24,
+    regime_window_hours: int = 24 * 30,  # окно для классификации режима (тренд/боковик)
+) -> list[BTResult]:
+    """
+    Прогоняет статическую и (если задан advisor) адаптивную сетки.
+    stride — через сколько свечей делаем новый «тик» сетки.
+
+    Дополнительно возвращает _raw_arrays в results[0].__dict__:
+      - regimes_at_t: список ('trending'|'sideways') для каждого тика
+      - static_pnls_per_t, adapt_pnls_per_t
+    """
+    n = len(df)
+    if n < lookback + pred_len + 1:
+        raise ValueError(f"need at least {lookback + pred_len + 1} candles, got {n}")
+
+    static_pnls, static_n, static_w = [], [], []
+    adapt_pnls, adapt_n, adapt_w = [], [], []
+    regimes = []
+    pause_hits = 0
+
+    peak_equity_s, peak_equity_a = 1.0, 1.0
+    eq_s, eq_a = 1.0, 1.0
+    max_dd_s, max_dd_a = 0.0, 0.0
+
+    for t in range(lookback, n - pred_len, stride):
+        horizon = min(pred_len, stride * 2)
+        center = float(df["close"].iloc[t])
+
+        # Режим рынка: в окне [t-regime_window, t] считаем |return|
+        win_start = max(0, t - regime_window_hours)
+        win = df["close"].iloc[win_start:t + 1]
+        ret = (win.iloc[-1] - win.iloc[0]) / win.iloc[0] if len(win) > 1 and win.iloc[0] > 0 else 0.0
+        regime = "trending" if abs(ret) > 0.05 else "sideways"   # ±5% за месяц → тренд
+        regimes.append(regime)
+
+        # STATIC
+        s_pnl, s_nt, s_w = simulate_grid(
+            df, t, horizon,
+            step_pct=static_step_pct,
+            levels=levels,
+            center_offset_pct=0.0,
+        )
+        static_pnls.append(s_pnl)
+        static_n.append(s_nt)
+        static_w.append(s_w)
+        eq_s *= (1.0 + s_pnl)
+        peak_equity_s = max(peak_equity_s, eq_s)
+        dd = (eq_s - peak_equity_s) / peak_equity_s
+        max_dd_s = min(max_dd_s, dd)
+
+        # ADAPTIVE
+        if advisor is not None:
+            try:
+                advice = advisor.advise(df.iloc[:t + 1], lookback=lookback, pred_len=pred_len, tf_min=tf_min)
+                if advice.pause_grid and advice.confidence > 0.5:
+                    pause_hits += 1
+                    adapt_pnls.append(0.0)
+                    adapt_n.append(0)
+                    adapt_w.append(0)
+                else:
+                    a_pnl, a_nt, a_w = simulate_grid(
+                        df, t, horizon,
+                        step_pct=max(static_step_pct * 0.5, advice.step_percent),
+                        levels=levels,
+                        center_offset_pct=advice.center_offset_pct,
+                    )
+                    adapt_pnls.append(a_pnl)
+                    adapt_n.append(a_nt)
+                    adapt_w.append(a_w)
+                    eq_a *= (1.0 + a_pnl)
+                    peak_equity_a = max(peak_equity_a, eq_a)
+                    dd = (eq_a - peak_equity_a) / peak_equity_a
+                    max_dd_a = min(max_dd_a, dd)
+            except Exception as e:
+                logger.warning("Advisor failed at t=%d: %s", t, e)
+                continue
+
+    def _agg(name, pnls, ns, ws, dd):
+        if not pnls:
+            return BTResult(name, 0.0, 0, 0.0, 0.0, 0.0)
+        total = float(np.prod(1.0 + np.array(pnls)) - 1.0)
+        n_total = sum(ns)
+        n_w = sum(ws)
+        wr = n_w / n_total if n_total else 0.0
+        if len(pnls) > 1:
+            sharpe = float(np.mean(pnls) / (np.std(pnls) + 1e-9) * math.sqrt(252 / max(1, len(pnls) // 24)))
+        else:
+            sharpe = 0.0
+        return BTResult(name, total, n_total, wr, dd, sharpe)
+
+    results = [_agg("static", static_pnls, static_n, static_w, max_dd_s)]
+    if advisor is not None:
+        results.append(_agg("adaptive", adapt_pnls, adapt_n, adapt_w, max_dd_a))
+        # Метрики по режимам
+        for regime in ("trending", "sideways"):
+            mask = [i for i, r in enumerate(regimes[:len(static_pnls)]) if r == regime]
+            if not mask:
+                continue
+            s_pnls_r = [static_pnls[i] for i in mask]
+            a_pnls_r = [adapt_pnls[i] for i in mask]
+            s_total = float(np.prod(1.0 + np.array(s_pnls_r)) - 1.0) if s_pnls_r else 0.0
+            a_total = float(np.prod(1.0 + np.array(a_pnls_r)) - 1.0) if a_pnls_r else 0.0
+            logger.info(
+                "regime=%-9s n_ticks=%-4d static_pnl=%+7.2f%%  adaptive_pnl=%+7.2f%%  delta=%+.2f%%  pause_hits_total=%d",
+                regime, len(mask), s_total * 100, a_total * 100,
+                (a_total - s_total) * 100, pause_hits,
+            )
+    return results
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--csv", required=True, help="CSV с колонками open,high,low,close,volume")
+    ap.add_argument("--tf", type=int, default=60)
+    ap.add_argument("--lookback", type=int, default=400)
+    ap.add_argument("--pred-len", type=int, default=24)
+    ap.add_argument("--step", type=float, default=0.005)
+    ap.add_argument("--levels", type=int, default=10)
+    ap.add_argument("--stride", type=int, default=24)
+    ap.add_argument("--with-advisor", action="store_true")
+    ap.add_argument("--model", default="kronos-mini")
+    ap.add_argument("--device", default="cpu")
+    ap.add_argument("--out", default=None)
+    args = ap.parse_args()
+
+    logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
+    log = logging.getLogger("backtest")
+
+    df = pd.read_csv(args.csv)
+    # нормализуем имена колонок
+    df.columns = [c.strip().lower() for c in df.columns]
+    if "timestamps" in df.columns:
+        df["timestamps"] = pd.to_datetime(df["timestamps"])
+        df = df.set_index("timestamps")
+    elif "timestamp" in df.columns:
+        df["timestamps"] = pd.to_datetime(df["timestamp"])
+        df = df.set_index("timestamps")
+    elif "open_time" in df.columns:
+        df.index = pd.to_datetime(df["open_time"], unit="ms", utc=True).tz_convert(None)
+    elif "date" in df.columns:
+        df["timestamps"] = pd.to_datetime(df["date"])
+        df = df.set_index("timestamps")
+    df = df[["open", "high", "low", "close", "volume"]].astype(float).sort_index()
+    log.info("Loaded %d candles, last close=%.2f", len(df), float(df["close"].iloc[-1]))
+
+    advisor = None
+    if args.with_advisor:
+        from kronos import KronosAdvisor
+        advisor = KronosAdvisor(model_name=args.model, device=args.device)
+        log.info("Loaded advisor %s", args.model)
+
+    results = run_backtest(
+        df,
+        tf_min=args.tf,
+        lookback=args.lookback,
+        pred_len=args.pred_len,
+        static_step_pct=args.step,
+        levels=args.levels,
+        advisor=advisor,
+        stride=args.stride,
+    )
+
+    print()
+    print("=" * 72)
+    print(f"{'MODE':<12} {'PnL %':>10} {'#Trades':>10} {'WinRate':>10} {'MaxDD %':>10} {'Sharpe':>10}")
+    print("-" * 72)
+    for r in results:
+        print(f"{r.name:<12} {r.total_pnl_pct*100:>9.2f}% {r.n_trades:>10d} "
+              f"{r.winrate*100:>9.1f}% {r.max_dd_pct*100:>9.2f}% {r.sharpe:>9.2f}")
+    print("=" * 72)
+
+    if args.out:
+        import json
+        Path(args.out).write_text(json.dumps([r.__dict__ for r in results], indent=2))
+        log.info("Wrote %s", args.out)
+
+
+if __name__ == "__main__":
+    main()

kronos/compare_backtest.py

@@ -0,0 +1,205 @@
+"""
+Compare two Kronos advisor configurations via backtest.
+
+Config A: "conservative" — текущие параметры
+  - pause при range > 5%
+  - step = max(base*0.7, min(base*1.5, advice))
+  - center_offset clamp ±0.2%
+
+Config B: "adjusted" — мягкие
+  - pause при range > 8% (или bias>0.7)
+  - step = max(base*0.85, min(base*1.3, advice))  # меньше снижение
+  - center_offset clamp ±0.15%
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import logging
+import math
+import sys
+import time
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Optional
+
+import numpy as np
+import pandas as pd
+
+logger = logging.getLogger(__name__)
+
+
+def simulate_grid(df, start, horizon, step_pct, levels,
+                  qty=0.001, center_offset_pct=0.0, fee_pct=0.001):
+    if start >= len(df) - 1:
+        return 0.0, 0, 0
+    center = float(df["close"].iloc[start]) * (1.0 + center_offset_pct)
+    step = center * step_pct
+    if step <= 0:
+        return 0.0, 0, 0
+    buy_levels = sorted([center - n * step for n in range(1, levels + 1) if center - n * step > 0])
+    sell_levels = sorted([center + n * step for n in range(1, levels + 1)])
+    end = min(start + horizon, len(df))
+    pnl, n_roundtrips, n_wins = 0.0, 0, 0
+    pos_open = None
+    for i in range(start + 1, end):
+        hi, lo = float(df["high"].iloc[i]), float(df["low"].iloc[i])
+        if pos_open is None:
+            for bl in buy_levels:
+                if lo <= bl:
+                    pnl -= bl * qty * fee_pct
+                    pos_open = bl
+                    break
+        else:
+            for sl in sell_levels:
+                if hi >= sl:
+                    gross = (sl - pos_open) * qty
+                    fee = sl * qty * fee_pct + pos_open * qty * fee_pct
+                    net = gross - fee
+                    pnl += net
+                    if net > 0:
+                        n_wins += 1
+                    n_roundtrips += 1
+                    pos_open = None
+                    break
+    pnl_pct = pnl / (center * qty) if center > 0 else 0.0
+    return pnl_pct, n_roundtrips, n_wins
+
+
+@dataclass
+class BTResult:
+    name: str
+    total_pnl_pct: float
+    n_trades: int
+    winrate: float
+    max_dd_pct: float
+    sharpe: float
+
+
+def run_one(df, tf_min, lookback, pred_len, static_step_pct, levels,
+            advisor, cfg: dict, stride=12, regime_window_hours=720):
+    n = len(df)
+    pnls, ns, ws, regimes = [], [], [], []
+    pause_hits = 0
+    peak, eq, max_dd = 1.0, 1.0, 0.0
+    for t in range(lookback, n - pred_len, stride):
+        horizon = min(pred_len, stride * 2)
+        # Regime
+        win_start = max(0, t - regime_window_hours)
+        win = df["close"].iloc[win_start:t + 1]
+        ret = (win.iloc[-1] - win.iloc[0]) / win.iloc[0] if len(win) > 1 and win.iloc[0] > 0 else 0.0
+        regime = "trending" if abs(ret) > 0.05 else "sideways"
+        regimes.append(regime)
+        # PAUSE check
+        advice = advisor.advise(df.iloc[:t + 1], lookback=lookback, pred_len=pred_len, tf_min=tf_min)
+        if advice.pause_grid and advice.confidence >= cfg["pause_min_conf"]:
+            if advice.expected_range_pct > cfg["pause_range_thr"] or (
+                abs(advice.center_offset_pct) / max(advice.expected_range_pct, 1e-4) > cfg["pause_bias_thr"]
+            ):
+                pause_hits += 1
+                pnls.append(0.0); ns.append(0); ws.append(0)
+                continue
+        # Step
+        kronos_step = float(advice.step_percent)
+        eff_step = max(static_step_pct * cfg["step_floor"],
+                       min(static_step_pct * cfg["step_ceil"], kronos_step))
+        # Center
+        eff_off = max(-cfg["center_clamp"], min(cfg["center_clamp"], advice.center_offset_pct))
+        pnl_pct, nt, nw = simulate_grid(df, t, horizon, eff_step, levels,
+                                        center_offset_pct=eff_off)
+        pnls.append(pnl_pct); ns.append(nt); ws.append(nw)
+        eq *= (1.0 + pnl_pct)
+        peak = max(peak, eq)
+        max_dd = min(max_dd, (eq - peak) / peak)
+    if not pnls:
+        return BTResult("?", 0.0, 0, 0.0, 0.0, 0.0), []
+    total = float(np.prod(1.0 + np.array(pnls)) - 1.0)
+    nt, nw = sum(ns), sum(ws)
+    wr = nw / nt if nt else 0.0
+    sharpe = float(np.mean(pnls) / (np.std(pnls) + 1e-9) * math.sqrt(252 / max(1, len(pnls) // 24))) if len(pnls) > 1 else 0.0
+    return BTResult(cfg["name"], total, nt, wr, max_dd, sharpe), regimes
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--csv", required=True)
+    ap.add_argument("--tf", type=int, default=60)
+    ap.add_argument("--lookback", type=int, default=400)
+    ap.add_argument("--pred-len", type=int, default=24)
+    ap.add_argument("--step", type=float, default=0.005)
+    ap.add_argument("--levels", type=int, default=10)
+    ap.add_argument("--stride", type=int, default=12)
+    ap.add_argument("--model", default="kronos-mini")
+    ap.add_argument("--device", default="cpu")
+    args = ap.parse_args()
+
+    logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
+    log = logging.getLogger("compare")
+
+    df = pd.read_csv(args.csv)
+    df.columns = [c.strip().lower() for c in df.columns]
+    if "timestamps" in df.columns:
+        df["timestamps"] = pd.to_datetime(df["timestamps"])
+        df = df.set_index("timestamps")
+    elif "timestamp" in df.columns:
+        df["timestamps"] = pd.to_datetime(df["timestamp"])
+        df = df.set_index("timestamps")
+    df = df[["open", "high", "low", "close", "volume"]].astype(float).sort_index()
+    log.info(f"Loaded {len(df)} candles")
+
+    # Static baseline (no advisor)
+    log.info("=== STATIC baseline (no advisor) ===")
+    static_pnls, static_n, static_w = [], [], []
+    peak, eq, max_dd = 1.0, 1.0, 0.0
+    n = len(df)
+    for t in range(args.lookback, n - args.pred_len, args.stride):
+        horizon = min(args.pred_len, args.stride * 2)
+        pnl, nt, nw = simulate_grid(df, t, horizon, args.step, args.levels)
+        static_pnls.append(pnl); static_n.append(nt); static_w.append(nw)
+        eq *= (1.0 + pnl)
+        peak = max(peak, eq)
+        max_dd = min(max_dd, (eq - peak) / peak)
+    s_total = float(np.prod(1.0 + np.array(static_pnls)) - 1.0)
+    s_sharpe = float(np.mean(static_pnls) / (np.std(static_pnls) + 1e-9) * math.sqrt(252 / max(1, len(static_pnls) // 24)))
+    static_res = BTResult("static", s_total, sum(static_n),
+                          sum(static_w) / sum(static_n) if sum(static_n) else 0.0,
+                          max_dd, s_sharpe)
+
+    from kronos import KronosAdvisor
+    advisor = KronosAdvisor(model_name=args.model, device=args.device)
+    log.info("Loaded advisor")
+
+    cfgs = [
+        {"name": "conservative",
+         "pause_min_conf": 0.6, "pause_range_thr": 0.05, "pause_bias_thr": 0.6,
+         "step_floor": 0.7, "step_ceil": 1.5, "center_clamp": 0.002},
+        {"name": "adjusted",
+         "pause_min_conf": 0.7, "pause_range_thr": 0.08, "pause_bias_thr": 0.75,
+         "step_floor": 0.85, "step_ceil": 1.3, "center_clamp": 0.0015},
+        {"name": "minimal",
+         "pause_min_conf": 0.85, "pause_range_thr": 0.10, "pause_bias_thr": 0.9,
+         "step_floor": 0.9, "step_ceil": 1.15, "center_clamp": 0.001},
+    ]
+    results = [static_res]
+    for cfg in cfgs:
+        log.info(f"=== {cfg['name']} ===")
+        t0 = time.time()
+        r, regimes = run_one(df, args.tf, args.lookback, args.pred_len,
+                             args.step, args.levels, advisor, cfg, args.stride)
+        log.info(f"  {cfg['name']} done in {time.time()-t0:.1f}s")
+        # Per-regime breakdown
+        results.append(r)
+
+    print()
+    print("=" * 88)
+    print(f"{'MODE':<15} {'PnL %':>10} {'#Trades':>10} {'WinRate':>10} {'MaxDD %':>10} {'Sharpe':>10}")
+    print("-" * 88)
+    for r in results:
+        print(f"{r.name:<15} {r.total_pnl_pct*100:>9.2f}% {r.n_trades:>10d} "
+              f"{r.winrate*100:>9.1f}% {r.max_dd_pct*100:>9.2f}% {r.sharpe:>9.2f}")
+    print("=" * 88)
+
+
+if __name__ == "__main__":
+    main()

kronos/data_source.py

@@ -0,0 +1,130 @@
+"""
+Data sources for Kronos advisor.
+
+Two flavours:
+  - TradernetHlocSource: production, uses Tradernet API (needs PRIVATE_KEY)
+  - BinancePublicSource: public, no auth, for offline tests & fallback
+
+Both expose the same interface:
+  fetch_ohlcv(lookback: int) -> pd.DataFrame
+    with columns ['open','high','low','close','volume'] and DatetimeIndex
+"""
+
+from __future__ import annotations
+
+import logging
+import time
+from typing import Optional
+
+import pandas as pd
+import requests
+
+logger = logging.getLogger(__name__)
+
+
+# Map our SYMBOL (Tradernet format) → Binance symbol
+SYMBOL_MAP_TRADERNET_TO_BINANCE = {
+    "BTC-USDT.IMEX": "BTCUSDT",
+    "ETH-USDT.IMEX": "ETHUSDT",
+    "SOL-USDT.IMEX": "SOLUSDT",
+    "TON-USDT.IMEX": "TONUSDT",
+}
+
+
+class TradernetHlocSource:
+    """
+    Тянет OHLCV с Tradernet (продакшн).
+    Требует приватный ключ → не используй в офлайн-тестах.
+    """
+
+    def __init__(self, api, symbol: str, tf_min: int = 60):
+        self.api = api           # экземпляр TradernetAPI
+        self.symbol = symbol
+        self.tf_min = tf_min
+
+    def fetch_ohlcv(self, lookback: int = 500, timeout: float = 20.0) -> pd.DataFrame:
+        # Берём запас побольше — на случай пропусков
+        # ВАЖНО: Tradernet по (date_from="", date_to="", count=N) возвращает
+        # САМЫЕ СТАРЫЕ N свечей, а не свежие. Подставляем явный date_to=NOW.
+        from datetime import datetime, timedelta
+        now = datetime.now()
+        # lookback свечей × tf_min минут = сколько часов назад начинать
+        hours_back = max(1, int(lookback * self.tf_min / 60) + 1)
+        date_from = (now - timedelta(hours=hours_back)).strftime("%d.%m.%Y %H:%M")
+        date_to = now.strftime("%d.%m.%Y %H:%M")
+        resp = self.api.get_hloc_sync(
+            ticker=self.symbol,
+            timeframe_min=self.tf_min,
+            date_from=date_from,
+            date_to=date_to,
+            count=0,  # count=0 — берём только между датами
+            timeout=timeout,
+        )
+        if not resp or "hloc" not in resp:
+            raise RuntimeError(f"Tradernet getHloc returned empty for {self.symbol}")
+
+        hloc_map = resp["hloc"]
+        x_map = resp.get("xSeries", {})
+        vl_map = resp.get("vl", {})
+
+        rows = hloc_map.get(self.symbol) or next(iter(hloc_map.values()))
+        ts_list = x_map.get(self.symbol) or next(iter(x_map.values()), [])
+        vols = vl_map.get(self.symbol) or next(iter(vl_map.values()), []) if vl_map else []
+
+        if not rows or not ts_list:
+            raise RuntimeError(f"Tradernet getHloc: empty series for {self.symbol}")
+
+        df = pd.DataFrame(rows, columns=["open", "high", "low", "close"])
+        df["volume"] = vols if len(vols) == len(df) else 0.0
+        # xSeries — unix-секунды
+        df.index = pd.to_datetime(ts_list, unit="s", utc=True).tz_convert(None)
+        df = df.sort_index()
+        return df.tail(lookback)
+
+
+class BinancePublicSource:
+    """
+    Публичный API Binance (https://api.binance.com). Без ключей.
+    Используется для offline-тестов Kronos и как fallback, если Tradernet недоступен.
+    """
+
+    BASE_URL = "https://api.binance.com"
+
+    def __init__(self, symbol: str, tf_min: int = 60):
+        # symbol в формате Binance: BTCUSDT
+        self.symbol = symbol
+        self.tf_min = tf_min
+
+    @classmethod
+    def from_tradernet(cls, tradernet_symbol: str, tf_min: int = 60) -> "BinancePublicSource":
+        """Конвертирует Tradernet-символ в Binance-символ."""
+        bsym = SYMBOL_MAP_TRADERNET_TO_BINANCE.get(tradernet_symbol)
+        if not bsym:
+            raise ValueError(f"No Binance mapping for {tradernet_symbol}")
+        return cls(bsym, tf_min)
+
+    def fetch_ohlcv(self, lookback: int = 500, timeout: float = 20.0) -> pd.DataFrame:
+        # Binance: 1m/3m/5m/15m/30m/1h/2h/4h/... (не "60m", а "1h")
+        interval_map = {1: "1m", 3: "3m", 5: "5m", 15: "15m", 30: "30m",
+                        60: "1h", 120: "2h", 240: "4h", 360: "6h",
+                        720: "12h", 1440: "1d"}
+        interval = interval_map.get(self.tf_min, f"{self.tf_min}m")
+        limit = min(1000, lookback)
+        url = f"{self.BASE_URL}/api/v3/klines"
+        params = {"symbol": self.symbol, "interval": interval, "limit": limit}
+
+        r = requests.get(url, params=params, timeout=timeout)
+        r.raise_for_status()
+        data = r.json()
+        if not data:
+            raise RuntimeError(f"Binance returned empty for {self.symbol}")
+
+        cols = ["open_time", "open", "high", "low", "close", "volume",
+                "close_time", "quote_vol", "trades", "taker_buy_base", "taker_buy_quote", "_"]
+        df = pd.DataFrame(data, columns=cols)
+        for c in ("open", "high", "low", "close", "volume"):
+            df[c] = df[c].astype(float)
+        df.index = pd.to_datetime(df["open_time"], unit="ms", utc=True)
+        df.index = df.index.tz_convert(None)
+        df = df[["open", "high", "low", "close", "volume"]].sort_index()
+        return df.tail(lookback)

kronos/data_source.py.bak.fresh_dates.20260614_042159

@@ -0,0 +1,122 @@
+"""
+Data sources for Kronos advisor.
+
+Two flavours:
+  - TradernetHlocSource: production, uses Tradernet API (needs PRIVATE_KEY)
+  - BinancePublicSource: public, no auth, for offline tests & fallback
+
+Both expose the same interface:
+  fetch_ohlcv(lookback: int) -> pd.DataFrame
+    with columns ['open','high','low','close','volume'] and DatetimeIndex
+"""
+
+from __future__ import annotations
+
+import logging
+import time
+from typing import Optional
+
+import pandas as pd
+import requests
+
+logger = logging.getLogger(__name__)
+
+
+# Map our SYMBOL (Tradernet format) → Binance symbol
+SYMBOL_MAP_TRADERNET_TO_BINANCE = {
+    "BTC-USDT.IMEX": "BTCUSDT",
+    "ETH-USDT.IMEX": "ETHUSDT",
+    "SOL-USDT.IMEX": "SOLUSDT",
+    "TON-USDT.IMEX": "TONUSDT",
+}
+
+
+class TradernetHlocSource:
+    """
+    Тянет OHLCV с Tradernet (продакшн).
+    Требует приватный ключ → не используй в офлайн-тестах.
+    """
+
+    def __init__(self, api, symbol: str, tf_min: int = 60):
+        self.api = api           # экземпляр TradernetAPI
+        self.symbol = symbol
+        self.tf_min = tf_min
+
+    def fetch_ohlcv(self, lookback: int = 500, timeout: float = 20.0) -> pd.DataFrame:
+        # Берём запас побольше — на случай пропусков
+        resp = self.api.get_hloc_sync(
+            ticker=self.symbol,
+            timeframe_min=self.tf_min,
+            date_from="",
+            date_to="",
+            count=lookback,
+            timeout=timeout,
+        )
+        if not resp or "hloc" not in resp:
+            raise RuntimeError(f"Tradernet getHloc returned empty for {self.symbol}")
+
+        hloc_map = resp["hloc"]
+        x_map = resp.get("xSeries", {})
+        vl_map = resp.get("vl", {})
+
+        rows = hloc_map.get(self.symbol) or next(iter(hloc_map.values()))
+        ts_list = x_map.get(self.symbol) or next(iter(x_map.values()), [])
+        vols = vl_map.get(self.symbol) or next(iter(vl_map.values()), []) if vl_map else []
+
+        if not rows or not ts_list:
+            raise RuntimeError(f"Tradernet getHloc: empty series for {self.symbol}")
+
+        df = pd.DataFrame(rows, columns=["open", "high", "low", "close"])
+        df["volume"] = vols if len(vols) == len(df) else 0.0
+        # xSeries — unix-секунды
+        df.index = pd.to_datetime(ts_list, unit="s", utc=True).tz_convert(None)
+        df = df.sort_index()
+        return df.tail(lookback)
+
+
+class BinancePublicSource:
+    """
+    Публичный API Binance (https://api.binance.com). Без ключей.
+    Используется для offline-тестов Kronos и как fallback, если Tradernet недоступен.
+    """
+
+    BASE_URL = "https://api.binance.com"
+
+    def __init__(self, symbol: str, tf_min: int = 60):
+        # symbol в формате Binance: BTCUSDT
+        self.symbol = symbol
+        self.tf_min = tf_min
+
+    @classmethod
+    def from_tradernet(cls, tradernet_symbol: str, tf_min: int = 60) -> "BinancePublicSource":
+        """Конвертирует Tradernet-символ в Binance-символ."""
+        bsym = SYMBOL_MAP_TRADERNET_TO_BINANCE.get(tradernet_symbol)
+        if not bsym:
+            raise ValueError(f"No Binance mapping for {tradernet_symbol}")
+        return cls(bsym, tf_min)
+
+    def fetch_ohlcv(self, lookback: int = 500, timeout: float = 20.0) -> pd.DataFrame:
+        # Binance: 1m/3m/5m/15m/30m/1h/2h/4h/... (не "60m", а "1h")
+        interval_map = {1: "1m", 3: "3m", 5: "5m", 15: "15m", 30: "30m",
+                        60: "1h", 120: "2h", 240: "4h", 360: "6h",
+                        720: "12h", 1440: "1d"}
+        interval = interval_map.get(self.tf_min, f"{self.tf_min}m")
+        limit = min(1000, lookback)
+        url = f"{self.BASE_URL}/api/v3/klines"
+        params = {"symbol": self.symbol, "interval": interval, "limit": limit}
+
+        r = requests.get(url, params=params, timeout=timeout)
+        r.raise_for_status()
+        data = r.json()
+        if not data:
+            raise RuntimeError(f"Binance returned empty for {self.symbol}")
+
+        cols = ["open_time", "open", "high", "low", "close", "volume",
+                "close_time", "quote_vol", "trades", "taker_buy_base", "taker_buy_quote", "_"]
+        df = pd.DataFrame(data, columns=cols)
+        for c in ("open", "high", "low", "close", "volume"):
+            df[c] = df[c].astype(float)
+        df.index = pd.to_datetime(df["open_time"], unit="ms", utc=True)
+        df.index = df.index.tz_convert(None)
+        df = df[["open", "high", "low", "close", "volume"]].sort_index()
+        return df.tail(lookback)

kronos/model/__init__.py

@@ -0,0 +1,17 @@
+from .kronos import KronosTokenizer, Kronos, KronosPredictor
+
+model_dict = {
+    'kronos_tokenizer': KronosTokenizer,
+    'kronos': Kronos,
+    'kronos_predictor': KronosPredictor
+}
+
+
+def get_model_class(model_name):
+    if model_name in model_dict:
+        return model_dict[model_name]
+    else:
+        print(f"Model {model_name} not found in model_dict")
+        raise NotImplementedError
+
+

kronos/model/kronos.py

@@ -0,0 +1,662 @@
+import numpy as np
+import pandas as pd
+import torch
+from huggingface_hub import PyTorchModelHubMixin
+import sys
+
+from tqdm import trange
+
+sys.path.append("../")
+from model.module import *
+
+
+class KronosTokenizer(nn.Module, PyTorchModelHubMixin):
+    """
+    KronosTokenizer module for tokenizing input data using a hybrid quantization approach.
+
+    This tokenizer utilizes a combination of encoder and decoder Transformer blocks
+    along with the Binary Spherical Quantization (BSQuantizer) to compress and decompress input data.
+
+    Args:
+           d_in (int): Input dimension.
+           d_model (int): Model dimension.
+           n_heads (int): Number of attention heads.
+           ff_dim (int): Feed-forward dimension.
+           n_enc_layers (int): Number of encoder layers.
+           n_dec_layers (int): Number of decoder layers.
+           ffn_dropout_p (float): Dropout probability for feed-forward networks.
+           attn_dropout_p (float): Dropout probability for attention mechanisms.
+           resid_dropout_p (float): Dropout probability for residual connections.
+           s1_bits (int): Number of bits for the pre token in BSQuantizer.
+           s2_bits (int): Number of bits for the post token in BSQuantizer.
+           beta (float): Beta parameter for BSQuantizer.
+           gamma0 (float): Gamma0 parameter for BSQuantizer.
+           gamma (float): Gamma parameter for BSQuantizer.
+           zeta (float): Zeta parameter for BSQuantizer.
+           group_size (int): Group size parameter for BSQuantizer.
+
+    """
+
+    def __init__(self, d_in, d_model, n_heads, ff_dim, n_enc_layers, n_dec_layers, ffn_dropout_p, attn_dropout_p, resid_dropout_p, s1_bits, s2_bits, beta, gamma0, gamma, zeta, group_size):
+
+        super().__init__()
+        self.d_in = d_in
+        self.d_model = d_model
+        self.n_heads = n_heads
+        self.ff_dim = ff_dim
+        self.enc_layers = n_enc_layers
+        self.dec_layers = n_dec_layers
+        self.ffn_dropout_p = ffn_dropout_p
+        self.attn_dropout_p = attn_dropout_p
+        self.resid_dropout_p = resid_dropout_p
+
+        self.s1_bits = s1_bits
+        self.s2_bits = s2_bits
+        self.codebook_dim = s1_bits + s2_bits # Total dimension of the codebook after quantization
+        self.embed = nn.Linear(self.d_in, self.d_model)
+        self.head = nn.Linear(self.d_model, self.d_in)
+
+        # Encoder Transformer Blocks
+        self.encoder = nn.ModuleList([
+            TransformerBlock(self.d_model, self.n_heads, self.ff_dim, self.ffn_dropout_p, self.attn_dropout_p, self.resid_dropout_p)
+            for _ in range(self.enc_layers - 1)
+        ])
+        # Decoder Transformer Blocks
+        self.decoder = nn.ModuleList([
+            TransformerBlock(self.d_model, self.n_heads, self.ff_dim, self.ffn_dropout_p, self.attn_dropout_p, self.resid_dropout_p)
+            for _ in range(self.dec_layers - 1)
+        ])
+        self.quant_embed = nn.Linear(in_features=self.d_model, out_features=self.codebook_dim) # Linear layer before quantization
+        self.post_quant_embed_pre = nn.Linear(in_features=self.s1_bits, out_features=self.d_model) # Linear layer after quantization (pre part - s1 bits)
+        self.post_quant_embed = nn.Linear(in_features=self.codebook_dim, out_features=self.d_model) # Linear layer after quantization (full codebook)
+        self.tokenizer = BSQuantizer(self.s1_bits, self.s2_bits, beta, gamma0, gamma, zeta, group_size) # BSQuantizer module
+
+    def forward(self, x):
+        """
+        Forward pass of the KronosTokenizer.
+
+        Args:
+            x (torch.Tensor): Input tensor of shape (batch_size, seq_len, d_in).
+
+        Returns:
+            tuple: A tuple containing:
+                - tuple: (z_pre, z) - Reconstructed outputs from decoder with s1_bits and full codebook respectively,
+                         both of shape (batch_size, seq_len, d_in).
+                - torch.Tensor: bsq_loss - Loss from the BSQuantizer.
+                - torch.Tensor: quantized - Quantized representation from BSQuantizer.
+                - torch.Tensor: z_indices - Indices from the BSQuantizer.
+        """
+        z = self.embed(x)
+
+        for layer in self.encoder:
+            z = layer(z)
+
+        z = self.quant_embed(z) # (B, T, codebook)
+
+        bsq_loss, quantized, z_indices = self.tokenizer(z)
+
+        quantized_pre = quantized[:, :, :self.s1_bits] # Extract the first part of quantized representation (s1_bits)
+        z_pre = self.post_quant_embed_pre(quantized_pre)
+
+        z = self.post_quant_embed(quantized)
+
+        # Decoder layers (for pre part - s1 bits)
+        for layer in self.decoder:
+            z_pre = layer(z_pre)
+        z_pre = self.head(z_pre)
+
+        # Decoder layers (for full codebook)
+        for layer in self.decoder:
+            z = layer(z)
+        z = self.head(z)
+
+        return (z_pre, z), bsq_loss, quantized, z_indices
+
+    def indices_to_bits(self, x, half=False):
+        """
+        Converts indices to bit representations and scales them.
+
+        Args:
+            x (torch.Tensor): Indices tensor.
+            half (bool, optional): Whether to process only half of the codebook dimension. Defaults to False.
+
+        Returns:
+            torch.Tensor: Bit representation tensor.
+        """
+        if half:
+            x1 = x[0] # Assuming x is a tuple of indices if half is True
+            x2 = x[1]
+            mask = 2 ** torch.arange(self.codebook_dim//2, device=x1.device, dtype=torch.long) # Create a mask for bit extraction
+            x1 = (x1.unsqueeze(-1) & mask) != 0 # Extract bits for the first half
+            x2 = (x2.unsqueeze(-1) & mask) != 0 # Extract bits for the second half
+            x = torch.cat([x1, x2], dim=-1) # Concatenate the bit representations
+        else:
+            mask = 2 ** torch.arange(self.codebook_dim, device=x.device, dtype=torch.long) # Create a mask for bit extraction
+            x = (x.unsqueeze(-1) & mask) != 0 # Extract bits
+
+        x = x.float() * 2 - 1 # Convert boolean to bipolar (-1, 1)
+        q_scale = 1. / (self.codebook_dim ** 0.5) # Scaling factor
+        x = x * q_scale
+        return x
+
+    def encode(self, x, half=False):
+        """
+        Encodes the input data into quantized indices.
+
+        Args:
+            x (torch.Tensor): Input tensor of shape (batch_size, seq_len, d_in).
+            half (bool, optional): Whether to use half quantization in BSQuantizer. Defaults to False.
+
+        Returns:
+            torch.Tensor: Quantized indices from BSQuantizer.
+        """
+        z = self.embed(x)
+        for layer in self.encoder:
+            z = layer(z)
+        z = self.quant_embed(z)
+
+        bsq_loss, quantized, z_indices = self.tokenizer(z, half=half, collect_metrics=False)
+        return z_indices
+
+    def decode(self, x, half=False):
+        """
+        Decodes quantized indices back to the input data space.
+
+        Args:
+            x (torch.Tensor): Quantized indices tensor.
+            half (bool, optional): Whether the indices were generated with half quantization. Defaults to False.
+
+        Returns:
+            torch.Tensor: Reconstructed output tensor of shape (batch_size, seq_len, d_in).
+        """
+        quantized = self.indices_to_bits(x, half)
+        z = self.post_quant_embed(quantized)
+        for layer in self.decoder:
+            z = layer(z)
+        z = self.head(z)
+        return z
+
+
+class Kronos(nn.Module, PyTorchModelHubMixin):
+    """
+    Kronos Model.
+
+    Args:
+        s1_bits (int): Number of bits for pre tokens.
+        s2_bits (int): Number of bits for post tokens.
+        n_layers (int): Number of Transformer blocks.
+        d_model (int): Dimension of the model's embeddings and hidden states.
+        n_heads (int): Number of attention heads in the MultiheadAttention layers.
+        ff_dim (int): Dimension of the feedforward network in the Transformer blocks.
+        ffn_dropout_p (float): Dropout probability for the feedforward network.
+        attn_dropout_p (float): Dropout probability for the attention layers.
+        resid_dropout_p (float): Dropout probability for residual connections.
+        token_dropout_p (float): Dropout probability for token embeddings.
+        learn_te (bool): Whether to use learnable temporal embeddings.
+    """
+
+    def __init__(self, s1_bits, s2_bits, n_layers, d_model, n_heads, ff_dim, ffn_dropout_p, attn_dropout_p, resid_dropout_p, token_dropout_p, learn_te):
+        super().__init__()
+        self.s1_bits = s1_bits
+        self.s2_bits = s2_bits
+        self.n_layers = n_layers
+        self.d_model = d_model
+        self.n_heads = n_heads
+        self.learn_te = learn_te
+        self.ff_dim = ff_dim
+        self.ffn_dropout_p = ffn_dropout_p
+        self.attn_dropout_p = attn_dropout_p
+        self.resid_dropout_p = resid_dropout_p
+        self.token_dropout_p = token_dropout_p
+
+        self.s1_vocab_size = 2 ** self.s1_bits
+        self.token_drop = nn.Dropout(self.token_dropout_p)
+        self.embedding = HierarchicalEmbedding(self.s1_bits, self.s2_bits, self.d_model)
+        self.time_emb = TemporalEmbedding(self.d_model, self.learn_te)
+        self.transformer = nn.ModuleList([
+            TransformerBlock(self.d_model, self.n_heads, self.ff_dim, self.ffn_dropout_p, self.attn_dropout_p, self.resid_dropout_p)
+            for _ in range(self.n_layers)
+        ])
+        self.norm = RMSNorm(self.d_model)
+        self.dep_layer = DependencyAwareLayer(self.d_model)
+        self.head = DualHead(self.s1_bits, self.s2_bits, self.d_model)
+        self.apply(self._init_weights)
+
+    def _init_weights(self, module):
+
+        if isinstance(module, nn.Linear):
+            nn.init.xavier_normal_(module.weight)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            nn.init.normal_(module.weight, mean=0, std=self.embedding.d_model ** -0.5)
+        elif isinstance(module, nn.LayerNorm):
+            nn.init.ones_(module.weight)
+            nn.init.zeros_(module.bias)
+        elif isinstance(module, RMSNorm):
+            nn.init.ones_(module.weight)
+
+    def forward(self, s1_ids, s2_ids, stamp=None, padding_mask=None, use_teacher_forcing=False, s1_targets=None):
+        """
+        Args:
+            s1_ids (torch.Tensor): Input tensor of s1 token IDs. Shape: [batch_size, seq_len]
+            s2_ids (torch.Tensor): Input tensor of s2 token IDs. Shape: [batch_size, seq_len]
+            stamp (torch.Tensor, optional): Temporal stamp tensor. Shape: [batch_size, seq_len]. Defaults to None.
+            padding_mask (torch.Tensor, optional): Mask for padding tokens. Shape: [batch_size, seq_len]. Defaults to None.
+            use_teacher_forcing (bool, optional): Whether to use teacher forcing for s1 decoding. Defaults to False.
+            s1_targets (torch.Tensor, optional): Target s1 token IDs for teacher forcing. Shape: [batch_size, seq_len]. Defaults to None.
+
+        Returns:
+            Tuple[torch.Tensor, torch.Tensor]:
+                - s1 logits: Logits for s1 token predictions. Shape: [batch_size, seq_len, s1_vocab_size]
+                - s2_logits: Logits for s2 token predictions, conditioned on s1. Shape: [batch_size, seq_len, s2_vocab_size]
+        """
+        x = self.embedding([s1_ids, s2_ids])
+        if stamp is not None:
+            time_embedding = self.time_emb(stamp)
+            x = x + time_embedding
+        x = self.token_drop(x)
+
+        for layer in self.transformer:
+            x = layer(x, key_padding_mask=padding_mask)
+
+        x = self.norm(x)
+
+        s1_logits = self.head(x)
+
+        if use_teacher_forcing:
+            sibling_embed = self.embedding.emb_s1(s1_targets)
+        else:
+            s1_probs = F.softmax(s1_logits.detach(), dim=-1)
+            sample_s1_ids = torch.multinomial(s1_probs.view(-1, self.s1_vocab_size), 1).view(s1_ids.shape)
+            sibling_embed = self.embedding.emb_s1(sample_s1_ids)
+
+        x2 = self.dep_layer(x, sibling_embed, key_padding_mask=padding_mask) # Dependency Aware Layer: Condition on s1 embeddings
+        s2_logits = self.head.cond_forward(x2)
+        return s1_logits, s2_logits
+
+    def decode_s1(self, s1_ids, s2_ids, stamp=None, padding_mask=None):
+        """
+        Decodes only the s1 tokens.
+
+        This method performs a forward pass to predict only s1 tokens. It returns the s1 logits
+        and the context representation from the Transformer, which can be used for subsequent s2 decoding.
+
+        Args:
+            s1_ids (torch.Tensor): Input tensor of s1 token IDs. Shape: [batch_size, seq_len]
+            s2_ids (torch.Tensor): Input tensor of s2 token IDs. Shape: [batch_size, seq_len]
+            stamp (torch.Tensor, optional): Temporal stamp tensor. Shape: [batch_size, seq_len]. Defaults to None.
+            padding_mask (torch.Tensor, optional): Mask for padding tokens. Shape: [batch_size, seq_len]. Defaults to None.
+
+        Returns:
+            Tuple[torch.Tensor, torch.Tensor]:
+                - s1 logits: Logits for s1 token predictions. Shape: [batch_size, seq_len, s1_vocab_size]
+                - context: Context representation from the Transformer. Shape: [batch_size, seq_len, d_model]
+        """
+        x = self.embedding([s1_ids, s2_ids])
+        if stamp is not None:
+            time_embedding = self.time_emb(stamp)
+            x = x + time_embedding
+        x = self.token_drop(x)
+
+        for layer in self.transformer:
+            x = layer(x, key_padding_mask=padding_mask)
+
+        x = self.norm(x)
+
+        s1_logits = self.head(x)
+        return s1_logits, x
+
+    def decode_s2(self, context, s1_ids, padding_mask=None):
+        """
+        Decodes the s2 tokens, conditioned on the context and s1 tokens.
+
+        This method decodes s2 tokens based on a pre-computed context representation (typically from `decode_s1`)
+        and the s1 token IDs. It uses the dependency-aware layer and the conditional s2 head to predict s2 tokens.
+
+        Args:
+            context (torch.Tensor): Context representation from the transformer (output of decode_s1).
+                                     Shape: [batch_size, seq_len, d_model]
+            s1_ids (torch.Tensor): Input tensor of s1 token IDs. Shape: [batch_size, seq_len]
+            padding_mask (torch.Tensor, optional): Mask for padding tokens. Shape: [batch_size, seq_len]. Defaults to None.
+
+        Returns:
+            torch.Tensor: s2 logits. Shape: [batch_size, seq_len, s2_vocab_size]
+        """
+        sibling_embed = self.embedding.emb_s1(s1_ids)
+        x2 = self.dep_layer(context, sibling_embed, key_padding_mask=padding_mask)
+        return self.head.cond_forward(x2)
+
+
+def top_k_top_p_filtering(
+        logits,
+        top_k: int = 0,
+        top_p: float = 1.0,
+        filter_value: float = -float("Inf"),
+        min_tokens_to_keep: int = 1,
+):
+    """Filter a distribution of logits using top-k and/or nucleus (top-p) filtering
+    Args:
+        logits: logits distribution shape (batch size, vocabulary size)
+        if top_k > 0: keep only top k tokens with highest probability (top-k filtering).
+        if top_p < 1.0: keep the top tokens with cumulative probability >= top_p (nucleus filtering).
+            Nucleus filtering is described in Holtzman et al. (http://arxiv.org/abs/1904.09751)
+        Make sure we keep at least min_tokens_to_keep per batch example in the output
+    From: https://gist.github.com/thomwolf/1a5a29f6962089e871b94cbd09daf317
+    """
+    if top_k > 0:
+        top_k = min(max(top_k, min_tokens_to_keep), logits.size(-1))  # Safety check
+        # Remove all tokens with a probability less than the last token of the top-k
+        indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
+        logits[indices_to_remove] = filter_value
+        return logits
+
+    if top_p < 1.0:
+        sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+        cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
+
+        # Remove tokens with cumulative probability above the threshold (token with 0 are kept)
+        sorted_indices_to_remove = cumulative_probs > top_p
+        if min_tokens_to_keep > 1:
+            # Keep at least min_tokens_to_keep (set to min_tokens_to_keep-1 because we add the first one below)
+            sorted_indices_to_remove[..., :min_tokens_to_keep] = 0
+        # Shift the indices to the right to keep also the first token above the threshold
+        sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
+        sorted_indices_to_remove[..., 0] = 0
+
+        # scatter sorted tensors to original indexing
+        indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
+        logits[indices_to_remove] = filter_value
+        return logits
+
+
+def sample_from_logits(logits, temperature=1.0, top_k=None, top_p=None, sample_logits=True):
+    logits = logits / temperature
+    if top_k is not None or top_p is not None:
+        if top_k > 0 or top_p < 1.0:
+            logits = top_k_top_p_filtering(logits, top_k=top_k, top_p=top_p)
+
+    probs = F.softmax(logits, dim=-1)
+
+    if not sample_logits:
+        _, x = torch.topk(probs, k=1, dim=-1)
+    else:
+        x = torch.multinomial(probs, num_samples=1)
+
+    return x
+
+
+def auto_regressive_inference(tokenizer, model, x, x_stamp, y_stamp, max_context, pred_len, clip=5, T=1.0, top_k=0, top_p=0.99, sample_count=5, verbose=False):
+    with torch.no_grad():
+        x = torch.clip(x, -clip, clip)
+
+        device = x.device
+        x = x.unsqueeze(1).repeat(1, sample_count, 1, 1).reshape(-1, x.size(1), x.size(2)).to(device)
+        x_stamp = x_stamp.unsqueeze(1).repeat(1, sample_count, 1, 1).reshape(-1, x_stamp.size(1), x_stamp.size(2)).to(device)
+        y_stamp = y_stamp.unsqueeze(1).repeat(1, sample_count, 1, 1).reshape(-1, y_stamp.size(1), y_stamp.size(2)).to(device)
+
+        x_token = tokenizer.encode(x, half=True)
+        
+        initial_seq_len = x.size(1)
+        batch_size = x_token[0].size(0)
+        total_seq_len = initial_seq_len + pred_len
+        full_stamp = torch.cat([x_stamp, y_stamp], dim=1)
+
+        generated_pre = x_token[0].new_empty(batch_size, pred_len)
+        generated_post = x_token[1].new_empty(batch_size, pred_len)
+
+        pre_buffer = x_token[0].new_zeros(batch_size, max_context)
+        post_buffer = x_token[1].new_zeros(batch_size, max_context)
+        buffer_len = min(initial_seq_len, max_context)
+        if buffer_len > 0:
+            start_idx = max(0, initial_seq_len - max_context)
+            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]
+
+        if verbose:
+            ran = trange
+        else:
+            ran = range
+        for i in ran(pred_len):
+            current_seq_len = initial_seq_len + i
+            window_len = min(current_seq_len, max_context)
+
+            if current_seq_len <= max_context:
+                input_tokens = [
+                    pre_buffer[:, :window_len],
+                    post_buffer[:, :window_len]
+                ]
+            else:
+                input_tokens = [pre_buffer, post_buffer]
+
+            context_end = current_seq_len
+            context_start = max(0, context_end - max_context)
+            current_stamp = full_stamp[:, context_start:context_end, :].contiguous()
+
+            s1_logits, context = model.decode_s1(input_tokens[0], input_tokens[1], current_stamp)
+            s1_logits = s1_logits[:, -1, :]
+            sample_pre = sample_from_logits(s1_logits, temperature=T, top_k=top_k, top_p=top_p, sample_logits=True)
+
+            s2_logits = model.decode_s2(context, sample_pre)
+            s2_logits = s2_logits[:, -1, :]
+            sample_post = sample_from_logits(s2_logits, temperature=T, top_k=top_k, top_p=top_p, sample_logits=True)
+
+            generated_pre[:, i] = sample_pre.squeeze(-1)
+            generated_post[:, i] = sample_post.squeeze(-1)
+
+            if current_seq_len < max_context:
+                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)
+
+        full_pre = torch.cat([x_token[0], generated_pre], dim=1)
+        full_post = torch.cat([x_token[1], generated_post], dim=1)
+
+        context_start = max(0, total_seq_len - max_context)
+        input_tokens = [
+            full_pre[:, context_start:total_seq_len].contiguous(),
+            full_post[:, context_start:total_seq_len].contiguous()
+        ]
+        z = tokenizer.decode(input_tokens, half=True)
+        z = z.reshape(-1, sample_count, z.size(1), z.size(2))
+        preds = z.cpu().numpy()
+        preds = np.mean(preds, axis=1)
+
+        return preds
+
+
+def calc_time_stamps(x_timestamp):
+    time_df = pd.DataFrame()
+    time_df['minute'] = x_timestamp.dt.minute
+    time_df['hour'] = x_timestamp.dt.hour
+    time_df['weekday'] = x_timestamp.dt.weekday
+    time_df['day'] = x_timestamp.dt.day
+    time_df['month'] = x_timestamp.dt.month
+    return time_df
+
+
+class KronosPredictor:
+
+    def __init__(self, model, tokenizer, device=None, max_context=512, clip=5):
+        self.tokenizer = tokenizer
+        self.model = model
+        self.max_context = max_context
+        self.clip = clip
+        self.price_cols = ['open', 'high', 'low', 'close']
+        self.vol_col = 'volume'
+        self.amt_vol = 'amount'
+        self.time_cols = ['minute', 'hour', 'weekday', 'day', 'month']
+        
+        # Auto-detect device if not specified
+        if device is None:
+            if torch.cuda.is_available():
+                device = "cuda:0"
+            elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
+                device = "mps"
+            else:
+                device = "cpu"
+        
+        self.device = device
+
+        self.tokenizer = self.tokenizer.to(self.device)
+        self.model = self.model.to(self.device)
+
+    def generate(self, x, x_stamp, y_stamp, pred_len, T, top_k, top_p, sample_count, verbose):
+
+        x_tensor = torch.from_numpy(np.array(x).astype(np.float32)).to(self.device)
+        x_stamp_tensor = torch.from_numpy(np.array(x_stamp).astype(np.float32)).to(self.device)
+        y_stamp_tensor = torch.from_numpy(np.array(y_stamp).astype(np.float32)).to(self.device)
+
+        preds = auto_regressive_inference(self.tokenizer, self.model, x_tensor, x_stamp_tensor, y_stamp_tensor, self.max_context, pred_len,
+                                          self.clip, T, top_k, top_p, sample_count, verbose)
+        preds = preds[:, -pred_len:, :]
+        return preds
+
+    def predict(self, df, x_timestamp, y_timestamp, pred_len, T=1.0, top_k=0, top_p=0.9, sample_count=1, verbose=True):
+
+        if not isinstance(df, pd.DataFrame):
+            raise ValueError("Input must be a pandas DataFrame.")
+
+        if not all(col in df.columns for col in self.price_cols):
+            raise ValueError(f"Price columns {self.price_cols} not found in DataFrame.")
+
+        df = df.copy()
+        if self.vol_col not in df.columns:
+            df[self.vol_col] = 0.0  # Fill missing volume with zeros
+            df[self.amt_vol] = 0.0  # Fill missing amount with zeros
+        if self.amt_vol not in df.columns and self.vol_col in df.columns:
+            df[self.amt_vol] = df[self.vol_col] * df[self.price_cols].mean(axis=1)
+
+        if df[self.price_cols + [self.vol_col, self.amt_vol]].isnull().values.any():
+            raise ValueError("Input DataFrame contains NaN values in price or volume columns.")
+
+        x_time_df = calc_time_stamps(x_timestamp)
+        y_time_df = calc_time_stamps(y_timestamp)
+
+        x = df[self.price_cols + [self.vol_col, self.amt_vol]].values.astype(np.float32)
+        x_stamp = x_time_df.values.astype(np.float32)
+        y_stamp = y_time_df.values.astype(np.float32)
+
+        x_mean, x_std = np.mean(x, axis=0), np.std(x, axis=0)
+
+        x = (x - x_mean) / (x_std + 1e-5)
+        x = np.clip(x, -self.clip, self.clip)
+
+        x = x[np.newaxis, :]
+        x_stamp = x_stamp[np.newaxis, :]
+        y_stamp = y_stamp[np.newaxis, :]
+
+        preds = self.generate(x, x_stamp, y_stamp, pred_len, T, top_k, top_p, sample_count, verbose)
+
+        preds = preds.squeeze(0)
+        preds = preds * (x_std + 1e-5) + x_mean
+
+        pred_df = pd.DataFrame(preds, columns=self.price_cols + [self.vol_col, self.amt_vol], index=y_timestamp)
+        return pred_df
+
+
+    def predict_batch(self, df_list, x_timestamp_list, y_timestamp_list, pred_len, T=1.0, top_k=0, top_p=0.9, sample_count=1, verbose=True):
+        """
+        Perform parallel (batch) prediction on multiple time series. All series must have the same historical length and prediction length (pred_len).
+
+        Args:
+            df_list (List[pd.DataFrame]): List of input DataFrames, each containing price columns and optional volume/amount columns.
+            x_timestamp_list (List[pd.DatetimeIndex or Series]): List of timestamps corresponding to historical data, length should match the number of rows in each DataFrame.
+            y_timestamp_list (List[pd.DatetimeIndex or Series]): List of future prediction timestamps, length should equal pred_len.
+            pred_len (int): Number of prediction steps.
+            T (float): Sampling temperature.
+            top_k (int): Top-k filtering threshold.
+            top_p (float): Top-p (nucleus sampling) threshold.
+            sample_count (int): Number of parallel samples per series, automatically averaged internally.
+            verbose (bool): Whether to display autoregressive progress.
+
+        Returns:
+            List[pd.DataFrame]: List of prediction results in the same order as input, each DataFrame contains
+                                `open, high, low, close, volume, amount` columns, indexed by corresponding `y_timestamp`.
+        """
+        # Basic validation
+        if not isinstance(df_list, (list, tuple)) or not isinstance(x_timestamp_list, (list, tuple)) or not isinstance(y_timestamp_list, (list, tuple)):
+            raise ValueError("df_list, x_timestamp_list, y_timestamp_list must be list or tuple types.")
+        if not (len(df_list) == len(x_timestamp_list) == len(y_timestamp_list)):
+            raise ValueError("df_list, x_timestamp_list, y_timestamp_list must have consistent lengths.")
+
+        num_series = len(df_list)
+
+        x_list = []
+        x_stamp_list = []
+        y_stamp_list = []
+        means = []
+        stds = []
+        seq_lens = []
+        y_lens = []
+
+        for i in range(num_series):
+            df = df_list[i]
+            if not isinstance(df, pd.DataFrame):
+                raise ValueError(f"Input at index {i} is not a pandas DataFrame.")
+            if not all(col in df.columns for col in self.price_cols):
+                raise ValueError(f"DataFrame at index {i} is missing price columns {self.price_cols}.")
+
+            df = df.copy()
+            if self.vol_col not in df.columns:
+                df[self.vol_col] = 0.0
+                df[self.amt_vol] = 0.0
+            if self.amt_vol not in df.columns and self.vol_col in df.columns:
+                df[self.amt_vol] = df[self.vol_col] * df[self.price_cols].mean(axis=1)
+
+            if df[self.price_cols + [self.vol_col, self.amt_vol]].isnull().values.any():
+                raise ValueError(f"DataFrame at index {i} contains NaN values in price or volume columns.")
+
+            x_timestamp = x_timestamp_list[i]
+            y_timestamp = y_timestamp_list[i]
+
+            x_time_df = calc_time_stamps(x_timestamp)
+            y_time_df = calc_time_stamps(y_timestamp)
+
+            x = df[self.price_cols + [self.vol_col, self.amt_vol]].values.astype(np.float32)
+            x_stamp = x_time_df.values.astype(np.float32)
+            y_stamp = y_time_df.values.astype(np.float32)
+
+            if x.shape[0] != x_stamp.shape[0]:
+                raise ValueError(f"Inconsistent lengths at index {i}: x has {x.shape[0]} vs x_stamp has {x_stamp.shape[0]}.")
+            if y_stamp.shape[0] != pred_len:
+                raise ValueError(f"y_timestamp length at index {i} should equal pred_len={pred_len}, got {y_stamp.shape[0]}.")
+
+            x_mean, x_std = np.mean(x, axis=0), np.std(x, axis=0)
+            x_norm = (x - x_mean) / (x_std + 1e-5)
+            x_norm = np.clip(x_norm, -self.clip, self.clip)
+
+            x_list.append(x_norm)
+            x_stamp_list.append(x_stamp)
+            y_stamp_list.append(y_stamp)
+            means.append(x_mean)
+            stds.append(x_std)
+
+            seq_lens.append(x_norm.shape[0])
+            y_lens.append(y_stamp.shape[0])
+
+        # Require all series to have consistent historical and prediction lengths for batch processing
+        if len(set(seq_lens)) != 1:
+            raise ValueError(f"Parallel prediction requires all series to have consistent historical lengths, got: {seq_lens}")
+        if len(set(y_lens)) != 1:
+            raise ValueError(f"Parallel prediction requires all series to have consistent prediction lengths, got: {y_lens}")
+
+        x_batch = np.stack(x_list, axis=0).astype(np.float32)           # (B, seq_len, feat)
+        x_stamp_batch = np.stack(x_stamp_list, axis=0).astype(np.float32) # (B, seq_len, time_feat)
+        y_stamp_batch = np.stack(y_stamp_list, axis=0).astype(np.float32) # (B, pred_len, time_feat)
+
+        preds = self.generate(x_batch, x_stamp_batch, y_stamp_batch, pred_len, T, top_k, top_p, sample_count, verbose)
+        # preds: (B, pred_len, feat)
+
+        pred_dfs = []
+        for i in range(num_series):
+            preds_i = preds[i] * (stds[i] + 1e-5) + means[i]
+            pred_df = pd.DataFrame(preds_i, columns=self.price_cols + [self.vol_col, self.amt_vol], index=y_timestamp_list[i])
+            pred_dfs.append(pred_df)
+
+        return pred_dfs
+

kronos/model/module.py

@@ -0,0 +1,570 @@
+import math
+
+from einops import rearrange, reduce
+import torch
+import torch.nn as nn
+from torch.autograd import Function
+import torch.nn.functional as F
+
+
+class DifferentiableEntropyFunction(Function):
+    @staticmethod
+    def forward(ctx, zq, basis, K, eps):
+        zb = (zq + 1) / 2
+        zi = ((zb * basis).sum(-1)).to(torch.int64)
+        cnt = torch.scatter_reduce(torch.zeros(2 ** K, device=zq.device, dtype=zq.dtype),
+                                   0,
+                                   zi.flatten(),
+                                   torch.ones_like(zi.flatten()).to(zq.dtype),
+                                   'sum')
+        prob = (cnt + eps) / (cnt + eps).sum()
+        H = -(prob * torch.log(prob)).sum()
+        ctx.save_for_backward(zq, zi, prob)
+        ctx.K = K
+        return H
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        zq, zi, prob = ctx.saved_tensors
+        grad_array = -grad_output * (torch.log(prob) + 1) / zi.numel() / ctx.K
+        reord_grad = grad_array[zi.flatten()].reshape(zi.shape)
+        grad_input = reord_grad.unsqueeze(-1) * zq
+        return grad_input, None, None, None, None
+
+
+def codebook_entropy(zq, basis, K, eps=1e-4):
+    return DifferentiableEntropyFunction.apply(zq, basis, K, eps)
+
+
+class BinarySphericalQuantizer(nn.Module):
+    def __init__(self, embed_dim, beta, gamma0, gamma, zeta,
+                 input_format='bchw',
+                 soft_entropy=True, group_size=9,
+                 persample_entropy_compute='analytical',
+                 cb_entropy_compute='group',
+                 l2_norm=True,
+                 inv_temperature=1):
+        """
+        Paper link: https://arxiv.org/pdf/2406.07548.pdf
+        Here we use the official implementation of the BinarySphericalQuantizer.
+        """
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.beta = beta  # loss weight for commit loss
+        self.gamma0 = gamma0  # loss weight for entropy penalty
+        self.gamma = gamma  # loss weight for entropy penalty
+        self.zeta = zeta  # loss weight for entire entropy penalty
+        self.input_format = input_format
+        assert self.embed_dim % group_size == 0, "embed_dim must be divisible by group_size"
+        self.num_groups = self.embed_dim // group_size
+        self.group_size = group_size
+        assert persample_entropy_compute in ['group', 'analytical'], "persample_entropy_compute must be either 'group' or 'analytical'"
+        assert cb_entropy_compute in ['group', 'nce'], "cb_entropy_compute must be either 'group' or 'nce'"
+        self.persample_entropy_compute = persample_entropy_compute
+        self.cb_entropy_compute = cb_entropy_compute
+        self.l2_norm = l2_norm
+        self.inv_temperature = inv_temperature
+
+        self.register_buffer('basis', 2 ** torch.arange(embed_dim - 1, -1, -1))
+        self.register_buffer('group_basis', 2 ** torch.arange(group_size - 1, -1, -1))
+
+        self.num_dimensions = 2 ** embed_dim
+        self.bits_per_index = embed_dim
+
+        # we only need to keep the codebook portion up to the group size
+        # because we approximate the H loss with this subcode
+        group_codes = torch.arange(2 ** self.group_size)
+        group_codebook = self.indexes_to_codes(group_codes).float()[:, -group_size:]
+        self.register_buffer('group_codebook', group_codebook, persistent=False)
+
+        self.soft_entropy = soft_entropy  # soft_entropy: Sec 3.2 of https://arxiv.org/pdf/1911.05894.pdf
+
+    def quantize(self, z):
+        assert z.shape[-1] == self.embed_dim, f"Expected {self.embed_dim} dimensions, got {z.shape[-1]}"
+
+        zhat = torch.where(z > 0,
+                           torch.tensor(1, dtype=z.dtype, device=z.device),
+                           torch.tensor(-1, dtype=z.dtype, device=z.device))
+        return z + (zhat - z).detach()
+
+    def forward(self, z, collect_metrics=True):
+        # if self.input_format == 'bchw':
+        #     z = rearrange(z, 'b c h w -> b h w c')
+        zq = self.quantize(z)
+
+        q_scale = 1. / (self.embed_dim ** 0.5) if self.l2_norm else 1.
+
+        zq = zq * q_scale
+
+        if not collect_metrics:
+            return zq, zq.new_zeros(()), {}
+
+        indices = self.codes_to_indexes(zq.detach())
+        group_indices = self.codes_to_group_indexes(zq.detach())
+        if not self.training:
+            used_codes = torch.unique(indices, return_counts=False)
+        else:
+            used_codes = None
+
+        if self.soft_entropy:
+            persample_entropy, cb_entropy, avg_prob = self.soft_entropy_loss(z)
+            entropy_penalty = self.gamma0 * persample_entropy - self.gamma * cb_entropy
+        else:
+            zb_by_sample = ((zq + 1) / 2).reshape(z.shape[0], -1, z.shape[-1]).to(torch.float32)
+            persample_entropy = self.get_hard_per_sample_entropy(zb_by_sample)
+            cb_entropy = codebook_entropy(zq, self.basis, self.embed_dim)
+            entropy_penalty = self.gamma0 * persample_entropy - self.gamma * cb_entropy
+
+        # commit loss
+        commit_loss = self.beta * torch.mean(((zq.detach() - z) ** 2).sum(dim=-1))
+
+        # if self.input_format == 'bchw':
+        #     zq = rearrange(zq, 'b h w c -> b c h w')
+
+        return (
+            zq,
+            commit_loss + self.zeta * entropy_penalty / self.inv_temperature,
+            {"H": cb_entropy, "used_codes": used_codes, "indices": indices, "group_indices": group_indices,
+             "avg_prob": avg_prob}
+        )
+
+    def soft_entropy_loss(self, z):
+        # if we divide the code in subgroups of size group_size, the codebook will be of size 2 ** group_size
+        # the sub-code is the last group_size bits of the full code
+        group_code_book = self.group_codebook / (self.embed_dim ** 0.5 if self.l2_norm else 1)
+        divided_z = rearrange(z, '... (g c) -> ... g c', c=self.group_size)
+
+        # we calculate the distance between the divided_z and the codebook for each subgroup
+        distance = - 2 * torch.einsum('... g c, d c ->... g d', divided_z, group_code_book)
+        prob = (-distance * self.inv_temperature).softmax(dim=-1)
+        if self.persample_entropy_compute == 'analytical':
+            if self.l2_norm:
+                p = torch.sigmoid(-4 * z / (self.embed_dim ** 0.5) * self.inv_temperature)
+            else:
+                p = torch.sigmoid(-4 * z * self.inv_temperature)
+            prob = torch.stack([p, 1 - p], dim=-1)
+            per_sample_entropy = self.get_entropy(prob, dim=-1, normalize=False).sum(dim=-1).mean()
+        else:
+            per_sample_entropy = self.get_entropy(prob, dim=-1, normalize=False).sum(dim=-1).mean()
+
+        # macro average of the probability of each subgroup
+        avg_prob = reduce(prob, '... g d ->g d', 'mean')
+        codebook_entropy = self.get_entropy(avg_prob, dim=-1, normalize=False)
+
+        # the approximation of the entropy is the sum of the entropy of each subgroup
+        return per_sample_entropy, codebook_entropy.sum(), avg_prob
+
+    def get_hard_per_sample_entropy(self, zb_by_sample):
+        probs_per_dim = zb_by_sample.sum(1) / zb_by_sample.shape[1]
+        persample_entropy = - probs_per_dim * torch.log(probs_per_dim + 1e-8) - (1 - probs_per_dim) * torch.log(1 - probs_per_dim + 1e-8)
+        persample_entropy = persample_entropy.sum(-1)
+        return persample_entropy.mean()
+
+    def codes_to_indexes(self, zhat):
+        """Converts a `code` to an index in the codebook.
+        Args:
+            zhat: A tensor of shape (B, ..., C) containing the codes. must be in {-1, 1}
+        """
+        assert zhat.shape[-1] == self.embed_dim, f"Expected {self.embed_dim} dimensions, got {zhat.shape[-1]}"
+        return ((zhat + 1) / 2 * self.basis).sum(axis=-1).to(torch.int64)
+
+    def codes_to_group_indexes(self, zhat):
+        """Converts a `code` to a list of indexes (in groups) in the codebook.
+        Args:
+            zhat: A tensor of shape (B, ..., C) containing the codes. must be in {-1, 1}
+        """
+        zhat_in_group = rearrange(zhat, 'b ... (g c) -> b ... g c', c=self.group_size)
+        return ((zhat_in_group + 1) / 2 * self.group_basis).sum(axis=-1).to(torch.int64)
+
+    def indexes_to_codes(self, indices):
+        """Inverse of `indexes_to_codes`."""
+        indices = indices.unsqueeze(-1)
+        codes_non_centered = torch.remainder(
+            torch.floor_divide(indices, self.basis), 2
+        )
+        return codes_non_centered * 2 - 1
+
+    def group_indexes_to_codes(self, group_indices):
+        """Inverse of `group_indexes_to_codes`."""
+        group_indices = group_indices.unsqueeze(-1)
+        codes_non_centered = torch.remainder(
+            torch.floor_divide(group_indices, self.group_basis), 2
+        )
+        codes_non_centered = rearrange(codes_non_centered, 'b ... g c -> b ... (g c)')
+        return codes_non_centered * 2 - 1
+
+    def get_entropy(self, count, dim=-1, eps=1e-4, normalize=True):
+        if normalize:
+            probs = (count + eps) / (count + eps).sum(dim=dim, keepdim=True)
+        else:
+            probs = count
+        H = -(probs * torch.log(probs + 1e-8)).sum(dim=dim)
+        return H
+
+    def get_group_codebook_entry(self, group_indices):
+        z_q = self.group_indexes_to_codes(group_indices)
+        q_scale = 1. / (self.embed_dim ** 0.5) if self.l2_norm else 1.
+        z_q = z_q * q_scale
+        if self.input_format == 'bchw':
+            h, w = int(z_q.shape[1] ** 0.5)
+            assert h * w == z_q.shape[1], 'Invalid sequence length'
+            z_q = rearrange(z_q, 'b (h w) c -> b c h w', h=h)
+        return z_q
+
+    def get_codebook_entry(self, indices):
+        z_q = self.indexes_to_codes(indices)
+        q_scale = 1. / (self.embed_dim ** 0.5) if self.l2_norm else 1.
+        z_q = z_q * q_scale
+        if self.input_format == 'bchw':
+            h, w = int(z_q.shape[1] ** 0.5)
+            assert h * w == z_q.shape[1], 'Invalid sequence length'
+            z_q = rearrange(z_q, 'b (h w) c -> b c h w', h=h)
+        return z_q
+
+
+class BSQuantizer(nn.Module):
+
+    def __init__(self, s1_bits, s2_bits, beta, gamma0, gamma, zeta, group_size):
+        super().__init__()
+        self.codebook_dim = s1_bits + s2_bits
+        self.s1_bits = s1_bits
+        self.s2_bits = s2_bits
+        self.bsq = BinarySphericalQuantizer(self.codebook_dim, beta, gamma0, gamma, zeta, group_size=group_size)
+
+    def bits_to_indices(self, bits):
+        bits = (bits >= 0).to(torch.long)
+        indices = 2 ** torch.arange(
+            0,
+            bits.shape[-1],
+            1,
+            dtype=torch.long,
+            device=bits.device,
+        )
+        return (bits * indices).sum(-1)
+
+    def forward(self, z, half=False, collect_metrics=True):
+        z = F.normalize(z, dim=-1)
+        quantized, bsq_loss, metrics = self.bsq(z, collect_metrics=collect_metrics)
+        if half:
+            q_pre = quantized[:, :, :self.s1_bits]
+            q_post = quantized[:, :, self.s1_bits:]
+            z_indices = [self.bits_to_indices(q_pre), self.bits_to_indices(q_post)]
+        else:
+            z_indices = self.bits_to_indices(quantized)
+        return bsq_loss, quantized, z_indices
+
+
+class RMSNorm(torch.nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-5):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def _norm(self, x):
+        return x * torch.rsqrt(torch.mean(x * x, dim=-1, keepdim=True) + self.eps)
+
+    def forward(self, x):
+        output = self._norm(x.float()).type_as(x)
+        return output * self.weight
+
+
+class FeedForward(nn.Module):
+    def __init__(self, d_model, ff_dim, ffn_dropout_p=0.0):
+        super().__init__()
+
+        self.w1 = nn.Linear(d_model, ff_dim, bias=False)
+        self.w3 = nn.Linear(d_model, ff_dim, bias=False)
+        self.w2 = nn.Linear(ff_dim, d_model, bias=False)
+        self.ffn_dropout = nn.Dropout(ffn_dropout_p)
+
+    def forward(self, x):
+        return self.ffn_dropout(self.w2(F.silu(self.w1(x)) * self.w3(x)))
+
+
+class RotaryPositionalEmbedding(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2).float() / dim))
+        self.register_buffer("inv_freq", inv_freq)
+        self.seq_len_cached = None
+        self.cos_cached = None
+        self.sin_cached = None
+
+    def _update_cos_sin_cache(self, x, seq_len):
+        if seq_len != self.seq_len_cached:
+            self.seq_len_cached = seq_len
+            t = torch.arange(seq_len, device=x.device).type_as(self.inv_freq)
+            freqs = torch.einsum('i,j->ij', t, self.inv_freq)
+            emb = torch.cat((freqs, freqs), dim=-1).to(x.device)
+            self.cos_cached = emb.cos()[None, None, :, :]
+            self.sin_cached = emb.sin()[None, None, :, :]
+        return self.cos_cached, self.sin_cached
+
+    def forward(self, q, k):
+        cos, sin = self._update_cos_sin_cache(q, q.shape[-2])
+        return (
+            (q * cos) + (self._rotate_half(q) * sin),
+            (k * cos) + (self._rotate_half(k) * sin),
+        )
+
+    def _rotate_half(self, x):
+        x1, x2 = x.chunk(2, dim=-1)
+        return torch.cat((-x2, x1), dim=-1)
+
+
+class MultiHeadAttentionWithRoPE(nn.Module):
+    def __init__(self, d_model, n_heads, attn_dropout_p=0.0, resid_dropout_p=0.0):
+        super().__init__()
+        self.d_model = d_model
+        self.n_heads = n_heads
+        self.head_dim = d_model // n_heads
+
+        self.q_proj = nn.Linear(d_model, d_model)
+        self.k_proj = nn.Linear(d_model, d_model)
+        self.v_proj = nn.Linear(d_model, d_model)
+        self.out_proj = nn.Linear(d_model, d_model)
+        self.rotary = RotaryPositionalEmbedding(self.head_dim)
+        self.attn_dropout_p = attn_dropout_p
+        self.resid_dropout = nn.Dropout(resid_dropout_p)
+
+    def forward(self, x, key_padding_mask=None):
+        batch_size, seq_len, _ = x.shape
+
+        q = self.q_proj(x).view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2)
+        k = self.k_proj(x).view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2)
+        v = self.v_proj(x).view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2)
+
+        q, k = self.rotary(q, k)
+
+        if key_padding_mask is not None:
+            attn_mask = key_padding_mask.unsqueeze(1).unsqueeze(2)  # [batch, 1, 1, seq_len]
+            attn_mask = attn_mask.expand(-1, self.n_heads, seq_len, -1)  # [batch, n_heads, q_len, k_len]
+        else:
+            attn_mask = None
+
+        attn_output = F.scaled_dot_product_attention(
+            q, k, v,
+            attn_mask=attn_mask,
+            dropout_p=self.attn_dropout_p if self.training else 0.0,
+            is_causal=True
+        )
+
+        attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model)
+        return self.resid_dropout(self.out_proj(attn_output))
+
+
+class MultiHeadCrossAttentionWithRoPE(nn.Module):
+    def __init__(self, d_model, n_heads, attn_dropout_p=0.0, resid_dropout=0.0):
+        super().__init__()
+        self.d_model = d_model
+        self.n_heads = n_heads
+        self.head_dim = d_model // n_heads
+
+        self.q_proj = nn.Linear(d_model, d_model)
+        self.k_proj = nn.Linear(d_model, d_model)
+        self.v_proj = nn.Linear(d_model, d_model)
+        self.out_proj = nn.Linear(d_model, d_model)
+        self.rotary = RotaryPositionalEmbedding(self.head_dim)
+        self.attn_dropout_p = attn_dropout_p
+        self.resid_dropout = nn.Dropout(resid_dropout)
+
+    def forward(self, query, key, value, key_padding_mask=None):
+        batch_size, q_len, _ = query.shape
+        _, seq_len, _ = key.shape
+
+        q = self.q_proj(query).view(batch_size, q_len, self.n_heads, self.head_dim).transpose(1, 2)
+        k = self.k_proj(key).view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2)
+        v = self.v_proj(value).view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2)
+
+        q, k = self.rotary(q, k)
+
+        if key_padding_mask is not None:
+            attn_mask = key_padding_mask.unsqueeze(1).unsqueeze(2)
+            attn_mask = attn_mask.expand(-1, self.n_heads, q_len, -1)
+        else:
+            attn_mask = None
+
+        is_causal_flag = self.training
+
+        attn_output = F.scaled_dot_product_attention(
+            q, k, v,
+            attn_mask=attn_mask,
+            dropout_p=self.attn_dropout_p if self.training else 0.0,
+            is_causal=is_causal_flag
+        )
+
+        attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, q_len, self.d_model)
+        return self.resid_dropout(self.out_proj(attn_output))
+
+
+class HierarchicalEmbedding(nn.Module):
+    def __init__(self, s1_bits, s2_bits, d_model=256):
+        super().__init__()
+        self.s1_bits = s1_bits
+        self.s2_bits = s2_bits
+
+        vocab_s1 = 2 ** s1_bits
+        vocab_s2 = 2 ** s2_bits
+
+        self.emb_s1 = nn.Embedding(vocab_s1, d_model)
+        self.emb_s2 = nn.Embedding(vocab_s2, d_model)
+        self.d_model = d_model
+        self.fusion_proj = nn.Linear(d_model * 2, d_model)
+
+        nn.init.normal_(self.emb_s1.weight, mean=0, std=d_model ** -0.5)
+        nn.init.normal_(self.emb_s2.weight, mean=0, std=d_model ** -0.5)
+
+    def split_token(self, token_ids: torch.Tensor, s2_bits: int):
+        """Inputs:
+            token_ids (torch.Tensor): Composite token IDs of shape [batch_size, seq_len] or [N], each in range [0, 2^(s1_bits + s2_bits) - 1].
+            s2_bits (int): Number of low bits used for the fine token (s2).
+        """
+        assert isinstance(s2_bits, int) and s2_bits > 0, "s2_bits must be a positive integer"
+
+        t = token_ids.long()
+        mask = (1 << s2_bits) - 1
+        s2_ids = t & mask           # extract low bits
+        s1_ids = t >> s2_bits       # extract high bits
+        return s1_ids, s2_ids
+
+    def forward(self, token_ids):
+        """Inputs:
+        token_ids:
+            - tuple or list: (s1_ids, s2_ids), each of shape [batch_size, seq_len], or
+            - torch.Tensor: composite token IDs of shape [batch_size, seq_len], which will be split into (s1_ids, s2_ids) internally.
+        Output: [batch_size, seq_len, d_model]
+        """
+        if isinstance(token_ids, tuple) or isinstance(token_ids, list):
+            s1_ids, s2_ids = token_ids
+        else:
+            s1_ids, s2_ids = self.split_token(token_ids, self.s2_bits)
+        s1_emb = self.emb_s1(s1_ids) * math.sqrt(self.d_model)
+        s2_emb = self.emb_s2(s2_ids) * math.sqrt(self.d_model)
+        return self.fusion_proj(torch.cat([s1_emb, s2_emb], dim=-1))
+
+
+class DependencyAwareLayer(nn.Module):
+    def __init__(self, d_model, n_heads=4, attn_dropout_p=0.0, resid_dropout=0.0):
+        super().__init__()
+        self.cross_attn = MultiHeadCrossAttentionWithRoPE(d_model, n_heads, attn_dropout_p, resid_dropout)
+        self.norm = RMSNorm(d_model)
+
+    def forward(self, hidden_states, sibling_embed, key_padding_mask=None):
+        """hidden_states: [batch, seq_len, d_model]
+        sibling_embed: Embedding from another subtoken
+        """
+        attn_out = self.cross_attn(
+            query=sibling_embed,
+            key=hidden_states,
+            value=hidden_states,
+            key_padding_mask=key_padding_mask
+        )
+        return self.norm(hidden_states + attn_out)
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, d_model, n_heads, ff_dim=1024, ffn_dropout_p=0.0, attn_dropout_p=0.0, resid_dropout_p=0.0):
+        super().__init__()
+        self.norm1 = RMSNorm(d_model)
+        self.self_attn = MultiHeadAttentionWithRoPE(d_model, n_heads, attn_dropout_p, resid_dropout_p)
+        self.norm2 = RMSNorm(d_model)
+        self.ffn = FeedForward(d_model, ff_dim, ffn_dropout_p)
+
+    def forward(self, x, key_padding_mask=None):
+        residual = x
+        x = self.norm1(x)
+        attn_out = self.self_attn(x, key_padding_mask=key_padding_mask)
+        x = residual + attn_out
+
+        residual = x
+        x = self.norm2(x)
+        ffn_out = self.ffn(x)
+        x = residual + ffn_out
+        return x
+
+
+class DualHead(nn.Module):
+    def __init__(self, s1_bits, s2_bits, d_model):
+        super().__init__()
+        self.vocab_s1 = 2 ** s1_bits
+        self.vocab_s2 = 2 ** s2_bits
+        self.proj_s1 = nn.Linear(d_model, self.vocab_s1)
+        self.proj_s2 = nn.Linear(d_model, self.vocab_s2)
+
+    def compute_loss(self, s1_logits, s2_logits, s1_targets, s2_targets, padding_mask=None):
+        if padding_mask is not None:
+            valid_mask = (padding_mask == 0)
+            s1_logits = s1_logits[valid_mask]
+            s2_logits = s2_logits[valid_mask]
+            s1_targets = s1_targets[valid_mask]
+            s2_targets = s2_targets[valid_mask]
+            ce_s1 = F.cross_entropy(s1_logits, s1_targets)
+            ce_s2 = F.cross_entropy(s2_logits, s2_targets)
+        else:
+            ce_s1 = F.cross_entropy(s1_logits.reshape(-1, self.vocab_s1), s1_targets.reshape(-1))
+            ce_s2 = F.cross_entropy(s2_logits.reshape(-1, self.vocab_s2), s2_targets.reshape(-1))
+        ce_loss = (ce_s1 + ce_s2) / 2
+        return ce_loss, ce_s1, ce_s2
+
+    def forward(self, x):
+        return self.proj_s1(x)
+
+    def cond_forward(self, x2):
+        return self.proj_s2(x2)
+
+
+class FixedEmbedding(nn.Module):
+    def __init__(self, c_in, d_model):
+        super(FixedEmbedding, self).__init__()
+
+        w = torch.zeros(c_in, d_model).float()
+        w.require_grad = False
+
+        position = torch.arange(0, c_in).float().unsqueeze(1)
+        div_term = (torch.arange(0, d_model, 2).float() * -(math.log(10000.0) / d_model)).exp()
+
+        w[:, 0::2] = torch.sin(position * div_term)
+        w[:, 1::2] = torch.cos(position * div_term)
+
+        self.emb = nn.Embedding(c_in, d_model)
+        self.emb.weight = nn.Parameter(w, requires_grad=False)
+
+    def forward(self, x):
+        return self.emb(x).detach()
+
+
+class TemporalEmbedding(nn.Module):
+    def __init__(self, d_model, learn_pe):
+        super(TemporalEmbedding, self).__init__()
+
+        minute_size = 60
+        hour_size = 24
+        weekday_size = 7
+        day_size = 32
+        month_size = 13
+
+        Embed = FixedEmbedding if not learn_pe else nn.Embedding
+        self.minute_embed = Embed(minute_size, d_model)
+        self.hour_embed = Embed(hour_size, d_model)
+        self.weekday_embed = Embed(weekday_size, d_model)
+        self.day_embed = Embed(day_size, d_model)
+        self.month_embed = Embed(month_size, d_model)
+
+    def forward(self, x):
+        x = x.long()
+
+        minute_x = self.minute_embed(x[:, :, 0])
+        hour_x = self.hour_embed(x[:, :, 1])
+        weekday_x = self.weekday_embed(x[:, :, 2])
+        day_x = self.day_embed(x[:, :, 3])
+        month_x = self.month_embed(x[:, :, 4])
+
+        return hour_x + weekday_x + day_x + month_x + minute_x
+
+
+
+
+
+
+
+