Download app.py from Chia-Ming/kronos-mc: direct link, hf CLI and curl.
- Browser
- Download file 11.3 kB
-
https://huggingface.co/spaces/Chia-Ming/kronos-mc/resolve/main/app.py
- Command line
-
hf download hf://spaces/Chia-Ming/kronos-mc/app.py
-
curl -L -o app.py https://huggingface.co/spaces/Chia-Ming/kronos-mc/resolve/main/app.py
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") | |
| 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) | |