279 lines
12 KiB
Python
279 lines
12 KiB
Python
"""
|
|
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(),
|
|
)
|