Spaces:
Running
Running
| """ | |
| Ensemble inference engine for the hosted dashboard. | |
| Same logic as ../engine.py but with flat imports (Space repo is self-contained). | |
| """ | |
| import json | |
| import time | |
| from pathlib import Path | |
| import numpy as np | |
| from model import Kronos, KronosTokenizer, KronosPredictor | |
| from predict_us import fetch_ohlcv, future_timestamps, report_metrics | |
| _COLS = ["open", "high", "low", "close", "volume"] | |
| _PROFILE_DIR = Path(__file__).resolve().parent / "out" / "profiles" | |
| def load_profile(ticker: str, interval: str = "1d"): | |
| """Return the cached trust profile for a ticker (from autoconfig.py), or None.""" | |
| p = _PROFILE_DIR / f"{ticker.upper()}_{interval}.json" | |
| if p.exists(): | |
| try: | |
| return json.loads(p.read_text(encoding="utf-8")) | |
| except Exception: | |
| return None | |
| return None | |
| # When the lookback window's max/min close exceeds this, Kronos' per-window | |
| # normalisation gets skewed by the trend and forecasts can collapse. Warn above it. | |
| _RANGE_THRESH = 2.0 | |
| _MIN_LOOKBACK = 40 | |
| _PREDICTOR = None | |
| _PREDICTOR_KEY = None | |
| def get_predictor(device: str = "cpu", max_context: int = 512, | |
| model_id: str = "NeoQuasar/Kronos-small", | |
| tokenizer_id: str = "NeoQuasar/Kronos-Tokenizer-base"): | |
| global _PREDICTOR, _PREDICTOR_KEY | |
| key = (device, max_context, model_id, tokenizer_id) | |
| if _PREDICTOR is None or _PREDICTOR_KEY != key: | |
| tokenizer = KronosTokenizer.from_pretrained(tokenizer_id) | |
| model = Kronos.from_pretrained(model_id) | |
| tokenizer.eval() | |
| model.eval() | |
| _PREDICTOR = KronosPredictor(model, tokenizer, device=device, max_context=max_context) | |
| _PREDICTOR_KEY = key | |
| return _PREDICTOR | |
| def is_loaded() -> bool: | |
| return _PREDICTOR is not None | |
| def run_ensemble(ticker: str = "AAPL", interval: str = "1d", period: str = "3y", | |
| lookback: int = 120, pred_len: int = 20, mode: str = "backtest", | |
| n_paths: int = 20, T: float = 0.7, top_p: float = 0.9, top_k: int = 0, | |
| seed: int = 123, device: str = "cpu", ctx_tail: int = 120, | |
| band_scale: float = 1.0) -> dict: | |
| import random | |
| import torch | |
| random.seed(seed); np.random.seed(seed); torch.manual_seed(seed) | |
| df = fetch_ohlcv(ticker, interval, period) | |
| need = lookback + (pred_len if mode == "backtest" else 0) | |
| if len(df) < need: | |
| raise ValueError( | |
| f"Not enough bars for {ticker}: have {len(df)}, need {need}. " | |
| f"Lower lookback/pred_len or raise period." | |
| ) | |
| if mode == "backtest": | |
| hist = df.iloc[-(lookback + pred_len):-pred_len] | |
| fut = df.iloc[-pred_len:] | |
| y_ts = fut["timestamps"].reset_index(drop=True) | |
| else: | |
| hist = df.iloc[-lookback:] | |
| fut = None | |
| y_ts = future_timestamps(hist["timestamps"].iloc[-1], interval, pred_len) | |
| x_df = hist[_COLS].reset_index(drop=True) | |
| x_ts = hist["timestamps"].reset_index(drop=True) | |
| predictor = get_predictor(device=device) | |
| t0 = time.time() | |
| paths = [] | |
| for i in range(n_paths): | |
| torch.manual_seed(seed + i); np.random.seed(seed + i) | |
| pred_df = predictor.predict( | |
| df=x_df, x_timestamp=x_ts, y_timestamp=y_ts, pred_len=pred_len, | |
| T=T, top_k=top_k, top_p=top_p, sample_count=1, verbose=False, | |
| ) | |
| paths.append(pred_df["close"].values.astype(float)) | |
| paths = np.stack(paths, axis=0) | |
| elapsed = time.time() - t0 | |
| median = np.median(paths, axis=0) | |
| lo = np.quantile(paths, 0.1, axis=0) | |
| hi = np.quantile(paths, 0.9, axis=0) | |
| if band_scale and band_scale != 1.0: | |
| lo = median - band_scale * (median - lo) | |
| hi = median + band_scale * (hi - median) | |
| def _iso(ts_series): | |
| return [pd_ts.isoformat() for pd_ts in ts_series] | |
| # Kronos normalises by the lookback window's mean/std; if that window spans a huge | |
| # price range (a strong trend), the latest price is an extreme outlier and forecasts | |
| # collapse toward the window mean. Flag it and suggest a shorter, tighter window. | |
| win = hist["close"].values.astype(float) | |
| range_ratio = float(win.max() / win.min()) if win.min() > 0 else float("inf") | |
| warning = None | |
| if range_ratio >= _RANGE_THRESH: | |
| suggest = None | |
| for L in range(len(win), _MIN_LOOKBACK, -1): | |
| w = win[-L:] | |
| if w.min() > 0 and w.max() / w.min() < _RANGE_THRESH: | |
| suggest = L # monotonic: first hit scanning down = largest safe window | |
| break | |
| warning = {"type": "wide_range", "range_ratio": round(range_ratio, 2), | |
| "suggest_lookback": suggest} | |
| ctx = hist.iloc[-ctx_tail:] | |
| out = { | |
| "meta": { | |
| "ticker": ticker, "interval": interval, "period": period, "mode": mode, | |
| "lookback": lookback, "pred_len": pred_len, "n_paths": n_paths, | |
| "T": T, "top_p": top_p, "top_k": top_k, "elapsed": round(elapsed, 1), | |
| "range_ratio": round(range_ratio, 2), "band_scale": round(band_scale, 3), | |
| }, | |
| "warning": warning, | |
| "history": { | |
| "t": _iso(ctx["timestamps"]), | |
| "close": [float(v) for v in ctx["close"].values], | |
| }, | |
| "forecast": { | |
| "t": _iso(y_ts), | |
| "median": [float(v) for v in median], | |
| "lo": [float(v) for v in lo], | |
| "hi": [float(v) for v in hi], | |
| "paths": [[float(v) for v in p] for p in paths], | |
| }, | |
| "truth": None, | |
| "metrics": None, | |
| } | |
| if mode == "backtest": | |
| truth = fut["close"].values.astype(float) | |
| out["truth"] = {"t": _iso(fut["timestamps"]), "close": [float(v) for v in truth]} | |
| m = report_metrics(truth, median) | |
| cover = float(np.mean((truth >= lo) & (truth <= hi)) * 100) | |
| m["coverage%"] = cover | |
| out["metrics"] = {k: round(v, 3) for k, v in m.items()} | |
| return out | |