Spaces:
Running
Running
| import os | |
| import sys | |
| import numpy as np | |
| import pandas as pd | |
| import torch | |
| import gradio as gr | |
| import plotly.graph_objects as go | |
| from plotly.subplots import make_subplots | |
| # Use all available CPU cores (HF free Spaces are CPU-only). | |
| try: | |
| torch.set_num_threads(max(1, os.cpu_count() or 1)) | |
| except Exception: | |
| pass | |
| # Make the vendored `model/` package importable. | |
| sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) | |
| from model import Kronos, KronosTokenizer, KronosPredictor # noqa: E402 | |
| # label -> (tokenizer_repo, model_repo, max_context) | |
| MODELS = { | |
| "Kronos-mini (4.1M, ctx 2048)": ("NeoQuasar/Kronos-Tokenizer-2k", "NeoQuasar/Kronos-mini", 2048), | |
| "Kronos-small (24.7M, ctx 512)": ("NeoQuasar/Kronos-Tokenizer-base", "NeoQuasar/Kronos-small", 512), | |
| "Kronos-base (102.3M, ctx 512)": ("NeoQuasar/Kronos-Tokenizer-base", "NeoQuasar/Kronos-base", 512), | |
| } | |
| DEFAULT_MODEL = "Kronos-small (24.7M, ctx 512)" | |
| PRICE_COLS = ["open", "high", "low", "close"] | |
| TIME_CANDIDATES = ["timestamps", "timestamp", "date", "datetime", "time"] | |
| _predictors = {} # cache loaded models so we only download/instantiate once | |
| def get_predictor(choice): | |
| if choice not in _predictors: | |
| tok_repo, mdl_repo, max_ctx = MODELS[choice] | |
| tokenizer = KronosTokenizer.from_pretrained(tok_repo) | |
| model = Kronos.from_pretrained(mdl_repo) | |
| _predictors[choice] = KronosPredictor(model, tokenizer, device="cpu", max_context=max_ctx) | |
| return _predictors[choice] | |
| def _read_df(csv_file): | |
| if csv_file is None: | |
| path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "sample_data.csv") | |
| else: | |
| path = csv_file if isinstance(csv_file, str) else getattr(csv_file, "name", csv_file) | |
| df = pd.read_csv(path) | |
| df.columns = [str(c).strip().lower() for c in df.columns] | |
| return df | |
| def _get_timestamps(df): | |
| for c in TIME_CANDIDATES: | |
| if c in df.columns: | |
| ts = pd.to_datetime(df[c], errors="coerce") | |
| if ts.notna().any(): | |
| return ts.reset_index(drop=True) | |
| # No timestamp column -> synthesize a regular 5-minute index. | |
| return pd.Series(pd.date_range("2024-01-01", periods=len(df), freq="5min")) | |
| def forecast(model_choice, lookback, pred_len, T, top_p, sample_count, csv_file): | |
| lookback, pred_len, sample_count = int(lookback), int(pred_len), int(sample_count) | |
| df = _read_df(csv_file) | |
| missing = [c for c in PRICE_COLS if c not in df.columns] | |
| if missing: | |
| raise gr.Error(f"CSV must contain columns {PRICE_COLS}. Missing: {missing}") | |
| if "volume" not in df.columns: | |
| df["volume"] = 0.0 | |
| if "amount" not in df.columns: | |
| df["amount"] = 0.0 | |
| _, _, max_ctx = MODELS[model_choice] | |
| lookback = min(lookback, max_ctx) | |
| if len(df) < lookback + 1: | |
| raise gr.Error(f"Need at least {lookback + 1} rows of history; the file has only {len(df)}.") | |
| df = df.reset_index(drop=True) | |
| ts_all = _get_timestamps(df) | |
| x_df = df.loc[: lookback - 1, ["open", "high", "low", "close", "volume", "amount"]] | |
| x_ts = ts_all.loc[: lookback - 1] | |
| have_future = len(df) >= lookback + pred_len | |
| if have_future: | |
| y_ts = ts_all.loc[lookback : lookback + pred_len - 1].reset_index(drop=True) | |
| else: | |
| deltas = ts_all.diff().dropna() | |
| step = deltas.median() if len(deltas) else pd.Timedelta(minutes=5) | |
| last = ts_all.iloc[lookback - 1] | |
| y_ts = pd.Series([last + step * (i + 1) for i in range(pred_len)]) | |
| predictor = get_predictor(model_choice) | |
| pred_df = predictor.predict( | |
| df=x_df, x_timestamp=x_ts, y_timestamp=y_ts, pred_len=pred_len, | |
| T=float(T), top_p=float(top_p), sample_count=sample_count, verbose=False, | |
| ).reset_index(drop=True) | |
| # Build interactive chart: history + forecast (+ actual future when available). | |
| hist_x = list(range(lookback)) | |
| fut_x = list(range(lookback, lookback + pred_len)) | |
| fig = make_subplots( | |
| rows=2, cols=1, shared_xaxes=True, vertical_spacing=0.07, | |
| row_heights=[0.7, 0.3], subplot_titles=("Close price", "Volume"), | |
| ) | |
| fig.add_trace(go.Scatter(x=hist_x, y=df.loc[: lookback - 1, "close"], | |
| name="History", line=dict(color="#1f77b4")), row=1, col=1) | |
| fig.add_trace(go.Scatter(x=fut_x, y=pred_df["close"], | |
| name="Kronos forecast", line=dict(color="#d62728", width=2)), row=1, col=1) | |
| if have_future: | |
| fig.add_trace(go.Scatter(x=fut_x, y=df.loc[lookback : lookback + pred_len - 1, "close"], | |
| name="Actual (future)", line=dict(color="#2ca02c", dash="dot")), row=1, col=1) | |
| fig.add_trace(go.Scatter(x=hist_x, y=df.loc[: lookback - 1, "volume"], | |
| line=dict(color="#1f77b4"), showlegend=False), row=2, col=1) | |
| fig.add_trace(go.Scatter(x=fut_x, y=pred_df["volume"], | |
| line=dict(color="#d62728"), showlegend=False), row=2, col=1) | |
| fig.add_vline(x=lookback - 1, line=dict(color="gray", dash="dash")) | |
| fig.update_layout(height=600, hovermode="x unified", | |
| legend=dict(orientation="h", yanchor="bottom", y=1.04, xanchor="left", x=0), | |
| margin=dict(l=50, r=20, t=60, b=40)) | |
| out = pred_df.copy() | |
| out.insert(0, "timestamp", list(y_ts)) | |
| return fig, out.round(4) | |
| DESCRIPTION = """ | |
| # 📈 Kronos — Financial K-line Forecasting | |
| Interactive demo of [**Kronos**](https://github.com/shiyu-coder/Kronos), the first open-source | |
| foundation model for financial candlesticks (K-lines), trained on data from 45+ global exchanges. | |
| Upload an OHLCV CSV (or use the bundled sample), choose a model size, and generate a probabilistic | |
| forecast of future candles. Required columns: **open, high, low, close**; *volume*, *amount* and a | |
| *timestamp* column are optional. | |
| > ⚠️ **Not financial advice.** This is a research/demo tool. Forecasts are probabilistic samples from | |
| > a model and must not be used as the sole basis for any trading or investment decision. | |
| """ | |
| with gr.Blocks(title="Kronos Forecast") as demo: | |
| gr.Markdown(DESCRIPTION) | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| model_choice = gr.Dropdown(list(MODELS.keys()), value=DEFAULT_MODEL, label="Model") | |
| csv_file = gr.File(label="OHLCV CSV (optional — uses sample if empty)", file_types=[".csv"]) | |
| lookback = gr.Slider(64, 512, value=256, step=8, label="Lookback (history length)") | |
| pred_len = gr.Slider(10, 120, value=30, step=5, label="Forecast length (candles)") | |
| with gr.Accordion("Sampling parameters", open=False): | |
| T = gr.Slider(0.1, 2.0, value=1.0, step=0.1, label="Temperature (T)") | |
| top_p = gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="Top-p (nucleus)") | |
| sample_count = gr.Slider(1, 5, value=1, step=1, label="Sample count (averaged)") | |
| run_btn = gr.Button("Generate forecast", variant="primary") | |
| gr.Markdown("*CPU inference: a 30-step forecast takes roughly 10-30s.*") | |
| with gr.Column(scale=2): | |
| plot = gr.Plot(label="Forecast") | |
| table = gr.Dataframe(label="Forecast values", wrap=True) | |
| run_btn.click( | |
| forecast, | |
| inputs=[model_choice, lookback, pred_len, T, top_p, sample_count, csv_file], | |
| outputs=[plot, table], | |
| ) | |
| gr.Markdown("Model weights: [NeoQuasar on Hugging Face](https://huggingface.co/NeoQuasar) - " | |
| "Code: [shiyu-coder/Kronos](https://github.com/shiyu-coder/Kronos) (MIT)") | |
| if __name__ == "__main__": | |
| demo.queue(max_size=8).launch( | |
| theme=gr.themes.Soft(), | |
| server_name="0.0.0.0", | |
| server_port=int(os.environ.get("PORT", 7860)), | |
| ) | |