Files

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(),
)