Initial deployment: Kronos stock predictor REST API

Monte-Carlo probabilistic forecasting using NeoQuasar/Kronos-base.
Async task queue (POST submit / GET poll), Tushare qfq data source.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

.gitignore

@@ -0,0 +1,5 @@
+__pycache__/
+*.pyc
+*.pyo
+.env
+docs/

Dockerfile

@@ -0,0 +1,27 @@
+# ── Build stage ───────────────────────────────────────────────────────────────
+FROM python:3.11-slim
+
+# System deps
+RUN apt-get update && apt-get install -y --no-install-recommends \
+        git \
+    && rm -rf /var/lib/apt/lists/*
+
+WORKDIR /app
+
+# Clone Kronos source at build time (avoids runtime clone delay)
+RUN git clone --depth 1 https://github.com/shiyu-coder/Kronos /app/Kronos
+
+# Install Python deps first (layer cache friendly)
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Copy application source
+COPY app.py predictor.py data_fetcher.py ./
+
+# HuggingFace Spaces default port
+EXPOSE 7860
+
+# KRONOS_DIR tells predictor.py where the source lives (already cloned above)
+ENV KRONOS_DIR=/app/Kronos
+
+CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]

README.md

@@ -0,0 +1,107 @@
+---
+title: Kronos Stock Predictor API
+emoji: 📈
+colorFrom: blue
+colorTo: indigo
+sdk: docker
+pinned: false
+---
+
+# Kronos Stock Predictor API
+
+Monte-Carlo probabilistic stock forecasting powered by
+[Kronos](https://arxiv.org/abs/2508.02739) — Tsinghua University's open-source
+financial K-line foundation model.
+
+Data source: **Tushare Pro** (front-adjusted / qfq).
+
+---
+
+## Endpoints
+
+### `POST /api/v1/predict`
+
+Submit a prediction job. Returns a `task_id` immediately.
+
+**Request body**
+
+| Field | Type | Default | Description |
+|-------|------|---------|-------------|
+| `ts_code` | string | — | Tushare stock code, e.g. `"600900.SH"` |
+| `lookback` | int | 512 | Historical bars to feed the model (1–512) |
+| `pred_len` | int | 5 | Future trading days to predict (1–60) |
+| `sample_count` | int | 30 | MC sampling iterations (1–100) |
+| `mode` | string | `"simple"` | `"simple"` or `"advanced"` |
+| `include_volume` | bool | false | Include volume CI in `advanced` mode |
+
+```json
+{
+  "ts_code": "600900.SH",
+  "lookback": 512,
+  "pred_len": 5,
+  "sample_count": 30,
+  "mode": "simple"
+}
+```
+
+**Response**
+
+```json
+{ "task_id": "550e8400-e29b-41d4-a716-446655440000" }
+```
+
+---
+
+### `GET /api/v1/predict/{task_id}`
+
+Poll for results.
+
+```json
+{
+  "status": "done",
+  "error": null,
+  "result": {
+    "ts_code": "600900.SH",
+    "base_date": "2026-03-13",
+    "pred_len": 5,
+    "confidence": 95,
+    "confidence_warning": false,
+    "direction": { "signal": "bullish", "probability": 0.73 },
+    "summary": {
+      "mean_close": 27.05,
+      "range_low": 25.80,
+      "range_high": 28.30,
+      "range_width": 2.50
+    },
+    "bands": [
+      {
+        "date": "2026-03-14",
+        "step": 1,
+        "mean_close": 26.88,
+        "trading_low": 26.20,
+        "trading_high": 27.55,
+        "uncertainty": 0.0504
+      }
+    ]
+  }
+}
+```
+
+`status` is one of `"pending"` / `"done"` / `"failed"`.
+
+---
+
+### `GET /health`
+
+```json
+{ "status": "ok" }
+```
+
+---
+
+## Notes
+
+- **Direction signal**: based on the last predicted close vs. the last historical close across all MC samples.
+- **95 % trading band**: `trading_low` = q2.5 of daily predicted lows; `trading_high` = q97.5 of daily predicted highs.
+- **`confidence_warning: true`** when `pred_len > 30` (model uncertainty grows significantly beyond ~30 days).
+- CPU inference: ~3–5 s/sample → 30 samples ≈ 2–5 min. Consider selecting GPU hardware for production use.

app.py

@@ -0,0 +1,211 @@
+"""
+Kronos Stock Predictor — RESTful API
+=====================================
+POST /api/v1/predict          → { "task_id": "uuid" }
+GET  /api/v1/predict/{id}     → { "status": "pending|done|failed", "result": {...} }
+GET  /health                  → { "status": "ok" }
+"""
+
+import asyncio
+import logging
+import uuid
+from concurrent.futures import ThreadPoolExecutor
+from contextlib import asynccontextmanager
+from typing import Literal
+
+import pandas as pd
+from fastapi import FastAPI, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel, Field
+
+import data_fetcher
+import predictor as pred_module
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# ── Task store (in-process; sufficient for single-worker deployments) ─────────
+_tasks: dict[str, dict] = {}
+_executor = ThreadPoolExecutor(max_workers=2)
+
+
+# ── Startup: eagerly load the model so the first request isn't slow ───────────
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    loop = asyncio.get_event_loop()
+    logger.info("Pre-loading Kronos predictor …")
+    await loop.run_in_executor(_executor, pred_module.get_predictor)
+    logger.info("Kronos predictor ready.")
+    yield
+
+
+app = FastAPI(
+    title="Kronos Stock Predictor API",
+    version="1.0.0",
+    description=(
+        "Monte-Carlo probabilistic stock forecasting powered by the "
+        "Kronos foundation model (Tsinghua University)."
+    ),
+    lifespan=lifespan,
+)
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+
+# ── Request / Response schemas ────────────────────────────────────────────────
+class PredictRequest(BaseModel):
+    ts_code: str = Field(..., examples=["600900.SH"], description="Tushare 股票代码")
+    lookback: int = Field(
+        default=512,
+        ge=20,
+        le=512,
+        description="回看历史 K 线根数（最多 512，不足时自动截断）",
+    )
+    pred_len: int = Field(
+        default=5,
+        ge=1,
+        le=60,
+        description="预测未来交易日数（建议 ≤ 30，超过时返回 confidence_warning）",
+    )
+    sample_count: int = Field(
+        default=30,
+        ge=1,
+        le=100,
+        description="MC 蒙特卡洛采样次数",
+    )
+    mode: Literal["simple", "advanced"] = Field(
+        default="simple",
+        description="simple: 仅返回均值 + 交易区间；advanced: 追加 OHLC 均值及收盘 CI",
+    )
+    include_volume: bool = Field(
+        default=False,
+        description="mode=advanced 时是否额外返回成交量预测（默认关闭）",
+    )
+
+
+# ── Response builder ──────────────────────────────────────────────────────────
+def _build_response(req: PredictRequest, base_date: str, pred_mean, ci,
+                    trading_low, trading_high, direction_prob, last_close,
+                    y_timestamp) -> dict:
+    bands = []
+    for i in range(req.pred_len):
+        band: dict = {
+            "date":         str(y_timestamp.iloc[i].date()),
+            "step":         i + 1,
+            "mean_close":   round(float(pred_mean["close"].iloc[i]), 4),
+            "trading_low":  round(float(trading_low[i]), 4),
+            "trading_high": round(float(trading_high[i]), 4),
+            "uncertainty":  round(
+                float((trading_high[i] - trading_low[i]) / last_close), 4
+            ),
+        }
+        if req.mode == "advanced":
+            band.update({
+                "mean_open":     round(float(pred_mean["open"].iloc[i]), 4),
+                "mean_high":     round(float(pred_mean["high"].iloc[i]), 4),
+                "mean_low":      round(float(pred_mean["low"].iloc[i]), 4),
+                "close_ci_low":  round(float(ci["close"]["low"][i]), 4),
+                "close_ci_high": round(float(ci["close"]["high"][i]), 4),
+            })
+        bands.append(band)
+
+    result: dict = {
+        "ts_code":             req.ts_code,
+        "base_date":           base_date,
+        "pred_len":            req.pred_len,
+        "confidence":          95,
+        "confidence_warning":  req.pred_len > 30,
+        "direction": {
+            "signal":      "bullish" if direction_prob >= 0.5 else "bearish",
+            "probability": round(direction_prob, 4),
+        },
+        "summary": {
+            "mean_close":  round(float(pred_mean["close"].iloc[-1]), 4),
+            "range_low":   round(float(trading_low.min()), 4),
+            "range_high":  round(float(trading_high.max()), 4),
+            "range_width": round(float(trading_high.max() - trading_low.min()), 4),
+        },
+        "bands": bands,
+    }
+
+    if req.mode == "advanced" and req.include_volume:
+        result["volume"] = [
+            {
+                "date":            str(y_timestamp.iloc[i].date()),
+                "mean_volume":     round(float(pred_mean["volume"].iloc[i])),
+                "volume_ci_low":   round(float(ci["volume"]["low"][i])),
+                "volume_ci_high":  round(float(ci["volume"]["high"][i])),
+            }
+            for i in range(req.pred_len)
+        ]
+
+    return result
+
+
+# ── Background task ───────────────────────────────────────────────────────────
+def _run_prediction(task_id: str, req: PredictRequest) -> None:
+    try:
+        x_df, x_timestamp, last_trade_date = data_fetcher.fetch_stock_data(
+            req.ts_code, req.lookback
+        )
+        y_timestamp = data_fetcher.get_future_trading_dates(last_trade_date, req.pred_len)
+
+        pred_mean, ci, trading_low, trading_high, direction_prob, last_close = (
+            pred_module.run_mc_prediction(
+                x_df, x_timestamp, y_timestamp, req.pred_len, req.sample_count
+            )
+        )
+
+        base_date = str(pd.to_datetime(last_trade_date, format="%Y%m%d").date())
+        result = _build_response(
+            req, base_date, pred_mean, ci,
+            trading_low, trading_high, direction_prob, last_close, y_timestamp,
+        )
+        _tasks[task_id] = {"status": "done", "result": result, "error": None}
+    except Exception as exc:
+        logger.exception("Task %s failed", task_id)
+        _tasks[task_id] = {"status": "failed", "result": None, "error": str(exc)}
+
+
+# ── Routes ────────────────────────────────────────────────────────────────────
+@app.post(
+    "/api/v1/predict",
+    summary="提交预测任务",
+    response_description="任务 ID，用于轮询结果",
+)
+async def submit_predict(req: PredictRequest):
+    """
+    提交一个蒙特卡洛预测任务，立即返回 `task_id`。
+    通过 `GET /api/v1/predict/{task_id}` 轮询结果。
+    """
+    task_id = str(uuid.uuid4())
+    _tasks[task_id] = {"status": "pending", "result": None, "error": None}
+    _executor.submit(_run_prediction, task_id, req)
+    return {"task_id": task_id}
+
+
+@app.get(
+    "/api/v1/predict/{task_id}",
+    summary="查询预测结果",
+)
+async def get_predict_result(task_id: str):
+    """
+    轮询预测任务状态。
+
+    - `status: "pending"` — 正在计算
+    - `status: "done"` — 完成，`result` 字段包含预测数据
+    - `status: "failed"` — 失败，`error` 字段包含错误信息
+    """
+    task = _tasks.get(task_id)
+    if task is None:
+        raise HTTPException(status_code=404, detail=f"Task {task_id!r} not found")
+    return task
+
+
+@app.get("/health", summary="健康检查")
+async def health():
+    return {"status": "ok"}

data_fetcher.py

@@ -0,0 +1,77 @@
+import os
+from datetime import datetime, timedelta
+
+import pandas as pd
+import tushare as ts
+
+TUSHARE_TOKEN = os.environ.get(
+    "TUSHARE_TOKEN",
+)
+
+ts.set_token(TUSHARE_TOKEN)
+_pro = ts.pro_api()
+
+
+def fetch_stock_data(
+    ts_code: str, lookback: int
+) -> tuple[pd.DataFrame, pd.Series, str]:
+    """
+    Returns:
+        x_df         : DataFrame with columns [open, high, low, close, volume, amount]
+        x_timestamp  : pd.Series[datetime], aligned to x_df
+        last_trade_date: str "YYYYMMDD", the most recent bar date
+    """
+    end_date = datetime.today().strftime("%Y%m%d")
+    # 2× buffer to account for weekends/holidays
+    start_date = (datetime.today() - timedelta(days=lookback * 2)).strftime("%Y%m%d")
+
+    df = ts.pro_bar(
+        ts_code=ts_code,
+        adj="qfq",
+        start_date=start_date,
+        end_date=end_date,
+        asset="E",
+    )
+
+    if df is None or df.empty:
+        raise ValueError(f"No data returned for ts_code={ts_code!r}")
+
+    df = df.sort_values("trade_date").reset_index(drop=True)
+    df = df.rename(columns={"vol": "volume"})
+    df["timestamps"] = pd.to_datetime(df["trade_date"], format="%Y%m%d")
+
+    # Keep the most recent `lookback` bars
+    df = df.tail(lookback).reset_index(drop=True)
+
+    x_df = df[["open", "high", "low", "close", "volume", "amount"]].copy()
+    x_timestamp = df["timestamps"].copy()
+    last_trade_date = df["trade_date"].iloc[-1]
+
+    return x_df, x_timestamp, last_trade_date
+
+
+def get_future_trading_dates(last_trade_date: str, pred_len: int) -> pd.Series:
+    """
+    Return a pd.Series of `pred_len` future SSE trading dates (datetime) that
+    follow `last_trade_date` (format: YYYYMMDD).
+    """
+    last_dt = datetime.strptime(last_trade_date, "%Y%m%d")
+    # 3× buffer so we always have enough dates even over a long holiday
+    end_dt = last_dt + timedelta(days=pred_len * 3)
+
+    cal = _pro.trade_cal(
+        exchange="SSE",
+        start_date=(last_dt + timedelta(days=1)).strftime("%Y%m%d"),
+        end_date=end_dt.strftime("%Y%m%d"),
+        is_open="1",
+    )
+    cal = cal.sort_values("cal_date")
+    dates = pd.to_datetime(cal["cal_date"].values[:pred_len], format="%Y%m%d")
+
+    if len(dates) < pred_len:
+        raise ValueError(
+            f"Could only obtain {len(dates)} future trading dates; "
+            f"increase buffer or check Tushare calendar coverage."
+        )
+
+    return pd.Series(dates)

predictor.py

@@ -0,0 +1,121 @@
+"""
+Kronos model singleton + Monte-Carlo prediction logic.
+
+On import this module:
+  1. Clones shiyu-coder/Kronos from GitHub if not already present at KRONOS_DIR.
+  2. Adds KRONOS_DIR to sys.path so `from model import ...` works.
+  3. Does NOT load the model weights yet (lazy, first-request).
+"""
+
+import logging
+import os
+import subprocess
+import sys
+from typing import Tuple
+
+import numpy as np
+import pandas as pd
+import torch
+
+logger = logging.getLogger(__name__)
+
+# ── Paths / IDs ─────────────────────────────────────────────────────────────
+KRONOS_DIR = os.environ.get("KRONOS_DIR", "/app/Kronos")
+MODEL_ID = "NeoQuasar/Kronos-base"
+TOKENIZER_ID = "NeoQuasar/Kronos-Tokenizer-base"
+
+
+# ── Bootstrap Kronos source ──────────────────────────────────────────────────
+def _ensure_kronos_source() -> None:
+    if not os.path.isdir(KRONOS_DIR):
+        logger.info("Cloning Kronos source to %s …", KRONOS_DIR)
+        subprocess.run(
+            [
+                "git", "clone", "--depth", "1",
+                "https://github.com/shiyu-coder/Kronos",
+                KRONOS_DIR,
+            ],
+            check=True,
+        )
+    if KRONOS_DIR not in sys.path:
+        sys.path.insert(0, KRONOS_DIR)
+
+
+_ensure_kronos_source()
+
+from model import Kronos, KronosPredictor, KronosTokenizer  # noqa: E402 (after sys.path setup)
+
+# ── Global singleton ─────────────────────────────────────────────────────────
+_predictor: KronosPredictor | None = None
+
+
+def get_predictor() -> KronosPredictor:
+    global _predictor
+    if _predictor is None:
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        logger.info("Loading Kronos model on %s …", device)
+        tokenizer = KronosTokenizer.from_pretrained(TOKENIZER_ID)
+        model = Kronos.from_pretrained(MODEL_ID)
+        _predictor = KronosPredictor(model, tokenizer, device=device, max_context=512)
+        logger.info("Kronos predictor ready.")
+    return _predictor
+
+
+# ── Monte-Carlo prediction ────────────────────────────────────────────────────
+def run_mc_prediction(
+    x_df: pd.DataFrame,
+    x_timestamp: pd.Series,
+    y_timestamp: pd.Series,
+    pred_len: int,
+    sample_count: int,
+) -> Tuple[pd.DataFrame, dict, np.ndarray, np.ndarray, float, float]:
+    """
+    Run `sample_count` independent samples (each with sample_count=1) to build
+    MC statistics.
+
+    Returns:
+        pred_mean      : DataFrame (index=y_timestamp, cols=OHLCVA), 均值轨迹
+        ci             : dict[field]["low"/"high"] → ndarray(pred_len,), 95% CI
+        trading_low    : ndarray(pred_len,), q2.5 of predicted_low
+        trading_high   : ndarray(pred_len,), q97.5 of predicted_high
+        direction_prob : float ∈ [0,1], fraction of samples where final close > last close
+        last_close     : float, closing price of the last historical bar
+    """
+    predictor = get_predictor()
+    samples = []
+
+    for _ in range(sample_count):
+        s = predictor.predict(
+            df=x_df,
+            x_timestamp=x_timestamp,
+            y_timestamp=y_timestamp,
+            pred_len=pred_len,
+            T=0.8,
+            top_p=0.9,
+            sample_count=1,
+            verbose=False,
+        )
+        samples.append(s)
+
+    pred_mean = pd.concat(samples).groupby(level=0).mean()
+
+    def stack(field: str) -> np.ndarray:
+        return np.stack([s[field].values for s in samples])  # (sample_count, pred_len)
+
+    alpha = 2.5  # → 95 % CI
+    ci = {
+        field: {
+            "low":  np.percentile(stack(field), alpha,       axis=0),
+            "high": np.percentile(stack(field), 100 - alpha, axis=0),
+        }
+        for field in ["open", "high", "low", "close", "volume"]
+    }
+
+    trading_low  = ci["low"]["low"]    # q2.5  of the predicted daily low
+    trading_high = ci["high"]["high"]  # q97.5 of the predicted daily high
+
+    last_close = float(x_df["close"].iloc[-1])
+    bull_count = sum(float(s["close"].iloc[-1]) > last_close for s in samples)
+    direction_prob = bull_count / sample_count
+
+    return pred_mean, ci, trading_low, trading_high, direction_prob, last_close

requirements.txt

@@ -0,0 +1,12 @@
+fastapi>=0.111.0
+uvicorn[standard]>=0.29.0
+pydantic>=2.0.0
+numpy
+pandas==2.2.2
+torch>=2.0.0
+einops==0.8.1
+huggingface_hub==0.33.1
+matplotlib==3.9.3
+tqdm==4.67.1
+safetensors==0.6.2
+tushare