kronos-mc / app.py
Chia-Ming's picture
Upload folder using huggingface_hub
4c93cbb verified
Raw History Blame Contribute Delete
11.3 kB
"""
Kronos Monte-Carlo 預測 —— Gradio + ZeroGPU 版。
原本是 Docker + Flask 跑在 HF 免費 CPU 上,一次 6-15 分鐘。
ZeroGPU 只支援 Gradio SDK,所以外殼改寫成 Gradio;預測的數學一行都沒動。
三個 ZeroGPU 的硬性要求,對應到下面的寫法:
1. Gradio SDK → README.md 的 sdk: gradio,本檔用 gr.Interface
2. 模型 module level 載入 → 下面的 tokenizer / model / predictor 都在模組層
3. 推論包 @spaces.GPU → _monte_carlo()
上游的 Kronos 套件(github.com/shiyu-coder/Kronos)沒有 setup.py 也沒有
pyproject.toml,所以 pip 裝不了。原本的 Dockerfile 是 build 時 git clone 進去,
Gradio Space 沒有 Dockerfile,改成啟動時把 model/ 那三個檔抓下來。
"""
import json
import os
import sys
import time
import urllib.request
from datetime import timedelta
# ─── 先把上游的 Kronos 套件放進 sys.path ─────────────────────────────────────
# 上游 repo 沒有 setup.py / pyproject.toml,pip 裝不了,所以直接抓 model/ 那三個檔
# (__init__.py / kronos.py / module.py,合計約 54KB)。
#
# ⚠️ shiyu-coder/Kronos 的預設分支是 **master**,不是 main。填錯會在開機時 404。
# KRONOS_REF 沒設的話會依序試 master、main;正式用建議直接釘成 commit SHA,
# 免得上游一改動你的 Space 就跟著壞。
KRONOS_FILES = ("__init__.py", "kronos.py", "module.py")
KRONOS_DIR = os.getenv("KRONOS_DIR", "/tmp/kronos_src")
_ref_env = os.getenv("KRONOS_REF", "").strip()
KRONOS_REFS = [_ref_env] if _ref_env else ["master", "main"]
def _raw_url(ref: str, fname: str) -> str:
return f"https://raw.githubusercontent.com/shiyu-coder/Kronos/{ref}/model/{fname}"
def _bootstrap_kronos() -> None:
pkg = os.path.join(KRONOS_DIR, "model")
os.makedirs(pkg, exist_ok=True)
missing = [f for f in KRONOS_FILES
if not (os.path.exists(os.path.join(pkg, f))
and os.path.getsize(os.path.join(pkg, f)) > 100)]
if not missing:
if KRONOS_DIR not in sys.path:
sys.path.insert(0, KRONOS_DIR)
return
errors = []
for ref in KRONOS_REFS:
try:
for fname in missing:
url = _raw_url(ref, fname)
print(f"[bootstrap] fetching {fname} @ {ref}")
with urllib.request.urlopen(url, timeout=30) as resp:
body = resp.read()
# raw.githubusercontent 對不存在的檔案是回 404,但保險起見也擋一下
# 「拿到 HTML 錯誤頁卻當成 Python」這種更難查的情況。
if len(body) < 100 or body.lstrip()[:1] == b"<":
raise RuntimeError(f"{url} 回傳的內容不像 Python({len(body)} bytes)")
with open(os.path.join(pkg, fname), "wb") as fh:
fh.write(body)
print(f"[bootstrap] ok @ {ref}")
if KRONOS_DIR not in sys.path:
sys.path.insert(0, KRONOS_DIR)
return
except Exception as exc: # noqa: BLE001 —— 換下一個 ref 再試
print(f"[bootstrap] ref={ref} 失敗:{exc}")
errors.append(f"{ref}: {exc}")
raise RuntimeError(
"抓不到 Kronos 原始碼。試過的 ref:" + "、".join(KRONOS_REFS) + "\n"
+ "\n".join(errors)
+ "\n請確認 https://github.com/shiyu-coder/Kronos 的分支名稱,"
"或用環境變數 KRONOS_REF 指定分支/commit SHA。"
)
_bootstrap_kronos()
# spaces 要在 torch 之前 import(ZeroGPU 會去接管 CUDA 的初始化)
import spaces # noqa: E402
import numpy as np # noqa: E402
import pandas as pd # noqa: E402
import torch # noqa: E402
import gradio as gr # noqa: E402
from model import Kronos, KronosTokenizer, KronosPredictor # noqa: E402
# ─── 設定 ────────────────────────────────────────────────────────────────────
# 模型尺寸做成環境變數,不寫死。三種選擇(參數量 / context / 搭配的 tokenizer):
# NeoQuasar/Kronos-mini 4.1M 2048 NeoQuasar/Kronos-Tokenizer-2k
# NeoQuasar/Kronos-small 24.7M 512 NeoQuasar/Kronos-Tokenizer-base
# NeoQuasar/Kronos-base 102.3M 512 NeoQuasar/Kronos-Tokenizer-base ← 預設
# 換 mini 的話 tokenizer 一定要跟著換,只改 MODEL_REPO 會對不起來。
MODEL_REPO = os.getenv("KRONOS_MODEL", "NeoQuasar/Kronos-base")
TOKENIZER_REPO = os.getenv("KRONOS_TOKENIZER", "NeoQuasar/Kronos-Tokenizer-base")
MAX_CONTEXT = int(os.getenv("KRONOS_MAX_CONTEXT", "512"))
# ZeroGPU 上,module level 的時候 GPU 還沒配給你,torch.cuda.is_available() 會回
# False —— 用它判斷會誤選 cpu,GPU 就白開了。所以看 HF 注入的旗標,不看 torch。
ZERO_GPU = os.getenv("SPACES_ZERO_GPU", "").lower() in ("1", "true")
if ZERO_GPU:
DEVICE = "cuda"
else:
DEVICE = os.getenv("KRONOS_DEVICE") or ("cuda" if torch.cuda.is_available() else "cpu")
# @spaces.GPU 的額度是按**宣告值**整包扣的(先扣不退),不是按實際跑多久。
# 免費帳號一天 300 秒:宣告 60 秒 → 一天 5 次;宣告 30 秒 → 10 次。
# 所以這個值要盡量貼近實測時間,不要留一大堆保險。
GPU_DURATION = int(os.getenv("KRONOS_GPU_DURATION", "60"))
print(f"[init] model={MODEL_REPO} tokenizer={TOKENIZER_REPO} "
f"device={DEVICE} zero_gpu={ZERO_GPU} max_context={MAX_CONTEXT}")
tokenizer = KronosTokenizer.from_pretrained(TOKENIZER_REPO)
model = Kronos.from_pretrained(MODEL_REPO)
predictor = KronosPredictor(
model=model, tokenizer=tokenizer, device=DEVICE, max_context=MAX_CONTEXT
)
print("[init] predictor ready")
@spaces.GPU(duration=GPU_DURATION)
def _monte_carlo(df, ts, y_timestamp, pred_len, sample_count, T, top_p):
"""實際吃 GPU 的那一段。跑 sample_count 條各自獨立的軌跡。
這裡刻意維持「跑 N 次、每次 sample_count=1」的寫法 —— 扇形圖需要 N 條
各自獨立的路徑,predict() 的 sample_count 參數是內部合成後只回一條。
"""
t0 = time.time()
paths_close = []
paths_ohlc = []
for _ in range(sample_count):
pred_df = predictor.predict(
df=df[["open", "high", "low", "close", "volume", "amount"]],
x_timestamp=pd.Series(ts),
y_timestamp=y_timestamp,
pred_len=pred_len,
T=T,
top_p=top_p,
sample_count=1,
verbose=False,
)
paths_close.append(pred_df["close"].astype(float).tolist())
paths_ohlc.append([
{
"ts": int(y_timestamp.iloc[i].timestamp() * 1000),
"open": float(pred_df["open"].iloc[i]),
"high": float(pred_df["high"].iloc[i]),
"low": float(pred_df["low"].iloc[i]),
"close": float(pred_df["close"].iloc[i]),
}
for i in range(pred_len)
])
return paths_close, paths_ohlc, time.time() - t0
def predict(k_lines_json, pred_len=24, sample_count=10, interval_ms=3600000,
T=1.0, top_p=0.9):
"""Gradio 的對外函式。回傳的 JSON 欄位跟舊版 Flask 完全一致。"""
try:
k_lines = json.loads(k_lines_json) if isinstance(k_lines_json, str) else k_lines_json
except (TypeError, ValueError) as exc:
return {"error": f"k_lines 不是合法的 JSON:{exc}"}
if not k_lines or len(k_lines) < 50:
return {"error": "need at least 50 k_lines"}
pred_len = int(pred_len)
sample_count = int(sample_count)
interval_ms = int(interval_ms)
T = float(T)
top_p = float(top_p)
# k_lines: [[ts_ms, open, high, low, close, volume], ...]
arr = np.array(k_lines, dtype=float)
ts = pd.to_datetime(arr[:, 0], unit="ms")
df = pd.DataFrame({
"open": arr[:, 1],
"high": arr[:, 2],
"low": arr[:, 3],
"close": arr[:, 4],
"volume": arr[:, 5] if arr.shape[1] > 5 else np.zeros(len(arr)),
"amount": np.zeros(len(arr)), # some versions expect this col
})
last_close = float(df["close"].iloc[-1])
# Baseline volatility: realized std over a window of the SAME length as
# the forecast horizon (pred_len). Comparing a 360-bar std to a 24-bar
# forecast std is unfair and systematically makes vol_amp_prob ~= 0.
tail_closes = df["close"].values[-(pred_len + 1):]
hist_returns = np.diff(np.log(tail_closes))
hist_std = float(np.std(hist_returns))
# Guard — if history was perfectly flat, fall back to full-series std
if hist_std == 0.0:
full_rets = np.diff(np.log(df["close"].values))
hist_std = float(np.std(full_rets)) if len(full_rets) else 1e-9
last_ts = ts[-1]
y_timestamp = pd.Series(
[last_ts + timedelta(milliseconds=interval_ms * (i + 1)) for i in range(pred_len)]
)
paths_close, paths_ohlc, elapsed = _monte_carlo(
df, ts, y_timestamp, pred_len, sample_count, T, top_p
)
finals = np.array([p[-1] for p in paths_close])
upside_prob = float(np.mean(finals > last_close))
path_stds = np.array([float(np.std(np.diff(np.log(np.array(p))))) for p in paths_close])
# Ratio form: mean predicted vol / baseline vol. >1 means amplified, <1 dampened.
vol_amp_ratio = float(np.mean(path_stds) / hist_std) if hist_std > 0 else 0.0
# Probability form: fraction of paths more volatile than baseline (kept for reference).
vol_amp_prob = float(np.mean(path_stds > hist_std))
return {
"ok": True,
"model": MODEL_REPO,
"device": DEVICE,
"sample_count": sample_count,
"pred_len": pred_len,
"elapsed_s": round(elapsed, 2),
"last_close": last_close,
"hist_std": hist_std,
"upside_prob": upside_prob,
"vol_amp": vol_amp_ratio, # legacy alias
"vol_amp_ratio": vol_amp_ratio,
"vol_amp_prob": vol_amp_prob,
"mean_path_close": np.mean(np.array(paths_close), axis=0).tolist(),
"paths": paths_ohlc,
"future_ts_ms": [int(t.timestamp() * 1000) for t in y_timestamp],
}
demo = gr.Interface(
fn=predict,
inputs=[
gr.Textbox(label="k_lines (JSON)", lines=3,
info="[[ts_ms, open, high, low, close, volume], ...],至少 50 根"),
gr.Number(label="pred_len", value=24, precision=0),
gr.Number(label="sample_count", value=10, precision=0),
gr.Number(label="interval_ms", value=3600000, precision=0),
gr.Number(label="T", value=1.0),
gr.Number(label="top_p", value=0.9),
],
outputs=gr.JSON(label="prediction"),
api_name="predict",
title="Kronos Monte-Carlo",
description=(f"model={MODEL_REPO} · device={DEVICE} · "
f"gpu_duration={GPU_DURATION}s\n\n"
"API: POST /gradio_api/call/predict → 取 event_id → "
"GET /gradio_api/call/predict/{event_id}(SSE)"),
)
if __name__ == "__main__":
demo.queue().launch(server_name="0.0.0.0", server_port=7860)