Commit ·
b8b028f
0
Parent(s):
Kronos market-forecast app + surface-roughness lab
Browse filesWeb app (crypto_ui): multi-asset forecasting UI, portfolio analyzer with
forecast matrix + saved reports, and a leverage trade optimizer. Research
lab (roughness_lab): ISO-16610 surface-roughness toolkit, Kronos sampling
calibration, and a GPU-ready roughness-aware fine-tune pipeline. Vendors the
MIT-licensed Kronos model/ package for self-contained deployment.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
- .gitignore +29 -0
- Dockerfile +24 -0
- README.md +42 -0
- crypto_ui/analyzer.html +445 -0
- crypto_ui/app.py +857 -0
- crypto_ui/index.html +502 -0
- crypto_ui/requirements.txt +9 -0
- model/LICENSE +21 -0
- model/__init__.py +17 -0
- model/kronos.py +662 -0
- model/module.py +570 -0
- roughness_lab/calibrate.py +253 -0
- roughness_lab/fetch_data.py +102 -0
- roughness_lab/gpu_finetune/README_GPU.md +86 -0
- roughness_lab/gpu_finetune/config_aapl_1h.yaml +65 -0
- roughness_lab/gpu_finetune/config_btc_1h.yaml +64 -0
- roughness_lab/gpu_finetune/config_smoke_cpu.yaml +60 -0
- roughness_lab/gpu_finetune/evaluate_texture.py +58 -0
- roughness_lab/gpu_finetune/train_rough.py +274 -0
- roughness_lab/roughness.py +187 -0
- run_prediction.py +116 -0
.gitignore
ADDED
|
@@ -0,0 +1,29 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Python
|
| 2 |
+
.venv/
|
| 3 |
+
__pycache__/
|
| 4 |
+
*.pyc
|
| 5 |
+
*.pyo
|
| 6 |
+
|
| 7 |
+
# Local tooling / secrets
|
| 8 |
+
.claude/settings.local.json
|
| 9 |
+
|
| 10 |
+
# Generated artifacts (regenerated at runtime)
|
| 11 |
+
output/
|
| 12 |
+
crypto_ui/analyses/
|
| 13 |
+
roughness_lab/results/
|
| 14 |
+
roughness_lab/gpu_finetune/finetuned/
|
| 15 |
+
*.log
|
| 16 |
+
|
| 17 |
+
# Large fetched datasets (re-fetch with roughness_lab/fetch_data.py)
|
| 18 |
+
roughness_lab/data/*.csv
|
| 19 |
+
|
| 20 |
+
# Upstream clone keeps its own git history; the app vendors only model/ at deploy time
|
| 21 |
+
Kronos/
|
| 22 |
+
|
| 23 |
+
# Hugging Face / model caches
|
| 24 |
+
.cache/
|
| 25 |
+
*.safetensors
|
| 26 |
+
|
| 27 |
+
# Tunnel binary + logs
|
| 28 |
+
cloudflared.exe
|
| 29 |
+
tunnel.*.log
|
Dockerfile
ADDED
|
@@ -0,0 +1,24 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Hugging Face Space (Docker SDK) for the Kronos forecast app.
|
| 2 |
+
# Mirrors the local layout so crypto_ui/app.py's imports resolve unchanged:
|
| 3 |
+
# /app/crypto_ui/app.py -> `from model import ...` via /app/Kronos/model
|
| 4 |
+
# optional texture metric via /app/roughness_lab
|
| 5 |
+
FROM python:3.11-slim
|
| 6 |
+
|
| 7 |
+
WORKDIR /app
|
| 8 |
+
ENV HOST=0.0.0.0 \
|
| 9 |
+
PORT=7860 \
|
| 10 |
+
HF_HOME=/app/.cache/huggingface \
|
| 11 |
+
PYTHONUNBUFFERED=1
|
| 12 |
+
|
| 13 |
+
COPY crypto_ui/requirements.txt ./requirements.txt
|
| 14 |
+
RUN pip install --no-cache-dir --extra-index-url https://download.pytorch.org/whl/cpu -r requirements.txt
|
| 15 |
+
|
| 16 |
+
COPY crypto_ui/ ./crypto_ui/
|
| 17 |
+
COPY model/ ./Kronos/model/
|
| 18 |
+
COPY roughness_lab/roughness.py ./roughness_lab/roughness.py
|
| 19 |
+
|
| 20 |
+
# Models download from the HF Hub on first request; this dir must be writable.
|
| 21 |
+
RUN mkdir -p /app/.cache/huggingface /app/crypto_ui/analyses && chmod -R 777 /app/.cache
|
| 22 |
+
|
| 23 |
+
EXPOSE 7860
|
| 24 |
+
CMD ["python", "crypto_ui/app.py"]
|
README.md
ADDED
|
@@ -0,0 +1,42 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Kronos Market Forecast
|
| 2 |
+
|
| 3 |
+
A small web app that turns the open-source [Kronos](https://github.com/shiyu-coder/Kronos)
|
| 4 |
+
financial foundation model into an interactive, multi-asset forecasting tool — plus a
|
| 5 |
+
surface-roughness research lab for calibrating and fine-tuning it.
|
| 6 |
+
|
| 7 |
+
> ⚠️ **Research / educational demo — not financial advice.** Forecasts are probabilistic
|
| 8 |
+
> samples from a model pre-trained largely on crypto K-lines, and are frequently wrong.
|
| 9 |
+
> The leverage optimizer is a mechanical translation of a forecast into trade structure,
|
| 10 |
+
> not a recommendation. Leverage can lose your entire margin (and more). Do not trade on this.
|
| 11 |
+
|
| 12 |
+
## What's inside
|
| 13 |
+
|
| 14 |
+
| Component | Path | What it does |
|
| 15 |
+
|---|---|---|
|
| 16 |
+
| **Chart UI** | `crypto_ui/` (`/`) | Search any asset (crypto via Binance, stocks/ETFs/FX/indices/commodities via Yahoo), forecast it with Kronos-small/base, see ghost-candle predictions + a p10–p90 uncertainty band. |
|
| 17 |
+
| **Portfolio Analyzer** | `crypto_ui/` (`/analyzer`) | Forecast one ticker across a matrix of intervals × horizons × both models; color-graded heatmaps, model-agreement consensus, save to JSON + standalone HTML report. |
|
| 18 |
+
| **Leverage Trade Optimizer** | `crypto_ui/app.py` (`/api/optimize`) | Turns a forecast cell into Conservative/Balanced/Aggressive trade setups (entry, stop, take-profits, leverage sized so the stop sits inside liquidation) with full reasoning. Suppresses setups when the two models disagree on direction. |
|
| 19 |
+
| **Surface-roughness lab** | `roughness_lab/` | Treats price as a measured surface profile (ISO-16610 Gaussian waviness/roughness split; Ra/Rq/Rz/RSm/Rsk/Rku/Wa). Calibrates Kronos sampling settings so forecast *texture* matches real market texture. |
|
| 20 |
+
| **GPU fine-tune pipeline** | `roughness_lab/gpu_finetune/` | Roughness-aware fine-tuning: per-epoch checkpoints selected by texture realism, ready to run on a GPU box. |
|
| 21 |
+
|
| 22 |
+
## Run locally
|
| 23 |
+
|
| 24 |
+
```bash
|
| 25 |
+
python -m venv .venv
|
| 26 |
+
.venv/Scripts/python -m pip install -r crypto_ui/requirements.txt # Windows
|
| 27 |
+
# source .venv/bin/activate && pip install -r crypto_ui/requirements.txt # macOS/Linux
|
| 28 |
+
python crypto_ui/app.py
|
| 29 |
+
# open http://127.0.0.1:8765 (analyzer at /analyzer)
|
| 30 |
+
```
|
| 31 |
+
|
| 32 |
+
Models (`NeoQuasar/Kronos-small`, `NeoQuasar/Kronos-base`, `NeoQuasar/Kronos-Tokenizer-base`)
|
| 33 |
+
download automatically from the Hugging Face Hub on first use. CPU works; a forecast takes
|
| 34 |
+
~15 s (small) to a few minutes (base). The app imports the Kronos `model/` package — when
|
| 35 |
+
deployed, that package is vendored alongside the app (see deploy notes).
|
| 36 |
+
|
| 37 |
+
## Credits & license
|
| 38 |
+
|
| 39 |
+
Built on [shiyu-coder/Kronos](https://github.com/shiyu-coder/Kronos) (MIT). Kronos:
|
| 40 |
+
*A Foundation Model for the Language of Financial Markets*, Shi et al., AAAI 2026
|
| 41 |
+
([arXiv:2508.02739](https://arxiv.org/abs/2508.02739)). This project is MIT-licensed; the
|
| 42 |
+
vendored `model/` package retains its upstream MIT license and copyright.
|
crypto_ui/analyzer.html
ADDED
|
@@ -0,0 +1,445 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
<!doctype html>
|
| 2 |
+
<html lang="en">
|
| 3 |
+
<head>
|
| 4 |
+
<meta charset="utf-8">
|
| 5 |
+
<meta name="viewport" content="width=device-width, initial-scale=1">
|
| 6 |
+
<title>Kronos — Portfolio Analyzer</title>
|
| 7 |
+
<link rel="preconnect" href="https://fonts.googleapis.com">
|
| 8 |
+
<link href="https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600&display=swap" rel="stylesheet">
|
| 9 |
+
<style>
|
| 10 |
+
:root {
|
| 11 |
+
--bg:#0a0b10; --panel:#10121a; --border:#1b1e2a; --text:#e8eaf2; --dim:#767d96;
|
| 12 |
+
--up:#16c784; --down:#ea3943; --accent:#8b5cf6;
|
| 13 |
+
}
|
| 14 |
+
* { box-sizing:border-box; }
|
| 15 |
+
[hidden] { display:none !important; }
|
| 16 |
+
body {
|
| 17 |
+
margin:0; background:var(--bg); color:var(--text); min-height:100vh;
|
| 18 |
+
font:14px/1.45 'Inter',system-ui,-apple-system,sans-serif; -webkit-font-smoothing:antialiased;
|
| 19 |
+
}
|
| 20 |
+
header {
|
| 21 |
+
display:flex; align-items:baseline; justify-content:space-between;
|
| 22 |
+
padding:18px 28px 14px; border-bottom:1px solid var(--border);
|
| 23 |
+
}
|
| 24 |
+
.brand { font-size:13px; font-weight:600; letter-spacing:.38em; }
|
| 25 |
+
.brand em { font-style:normal; color:var(--accent); }
|
| 26 |
+
.brand span { letter-spacing:.04em; font-weight:400; color:var(--dim); margin-left:14px; }
|
| 27 |
+
.nav a { color:var(--dim); text-decoration:none; font-size:12px; margin-left:18px; }
|
| 28 |
+
.nav a:hover { color:var(--text); }
|
| 29 |
+
.nav a.on { color:var(--accent); }
|
| 30 |
+
|
| 31 |
+
main { padding:18px 28px 40px; max-width:1100px; }
|
| 32 |
+
.row { display:flex; align-items:center; gap:14px; flex-wrap:wrap; margin-bottom:14px; }
|
| 33 |
+
.lbl { font-size:10px; text-transform:uppercase; letter-spacing:.12em; color:var(--dim); }
|
| 34 |
+
|
| 35 |
+
.search-wrap { position:relative; flex:1 1 320px; max-width:460px; }
|
| 36 |
+
#q { width:100%; background:var(--panel); border:1px solid var(--border); color:var(--text);
|
| 37 |
+
border-radius:999px; padding:9px 18px; font:400 13px 'Inter',sans-serif; outline:none; }
|
| 38 |
+
#q:focus { border-color:rgba(139,92,246,.55); }
|
| 39 |
+
#q::placeholder { color:#4d5469; }
|
| 40 |
+
.dropdown { position:absolute; top:calc(100% + 6px); left:0; right:0; z-index:20;
|
| 41 |
+
background:var(--panel); border:1px solid var(--border); border-radius:12px;
|
| 42 |
+
max-height:320px; overflow-y:auto; box-shadow:0 14px 40px rgba(0,0,0,.5); }
|
| 43 |
+
.opt { display:flex; align-items:center; gap:10px; padding:9px 14px; cursor:pointer; font-size:12px; }
|
| 44 |
+
.opt:hover,.opt.sel { background:rgba(139,92,246,.10); }
|
| 45 |
+
.opt .sym { font-weight:600; min-width:84px; }
|
| 46 |
+
.opt .nm { color:var(--dim); flex:1; overflow:hidden; text-overflow:ellipsis; white-space:nowrap; }
|
| 47 |
+
.opt .ex { color:#4d5469; font-size:10px; }
|
| 48 |
+
.none { padding:12px 14px; color:var(--dim); font-size:12px; }
|
| 49 |
+
|
| 50 |
+
.badge { font-size:9px; font-weight:600; text-transform:uppercase; letter-spacing:.1em;
|
| 51 |
+
padding:2px 8px; border-radius:999px; border:1px solid var(--dim); color:var(--dim); white-space:nowrap; }
|
| 52 |
+
.badge[data-k="Crypto"]{border-color:#8b5cf6;color:#a78bfa}.badge[data-k="Equity"]{border-color:#16c784;color:#34d399}
|
| 53 |
+
.badge[data-k="ETF"]{border-color:#2dd4bf;color:#5eead4}.badge[data-k="Forex"]{border-color:#60a5fa;color:#93c5fd}
|
| 54 |
+
.badge[data-k="Index"]{border-color:#f59e0b;color:#fbbf24}.badge[data-k="Commodity"]{border-color:#fb923c;color:#fdba74}
|
| 55 |
+
|
| 56 |
+
.staged { display:flex; align-items:center; gap:9px; padding:6px 8px 6px 13px;
|
| 57 |
+
border:1px dashed rgba(139,92,246,.5); border-radius:999px; font-size:12px; }
|
| 58 |
+
.staged .x { border:none; background:transparent; color:var(--dim); cursor:pointer; font-size:14px; padding:0 5px; }
|
| 59 |
+
.staged .x:hover { color:var(--text); }
|
| 60 |
+
|
| 61 |
+
.pill { border:1px solid var(--border); background:transparent; color:var(--dim); border-radius:999px;
|
| 62 |
+
padding:5px 13px; font:500 12px 'Inter',sans-serif; cursor:pointer; transition:all .15s; }
|
| 63 |
+
.pill:hover { border-color:#2b3044; color:var(--text); }
|
| 64 |
+
.pill.on { background:rgba(139,92,246,.13); border-color:rgba(139,92,246,.55); color:var(--text); }
|
| 65 |
+
.pill .px { opacity:.6; font-size:10px; margin-left:5px; }
|
| 66 |
+
|
| 67 |
+
#run { border:none; border-radius:999px; padding:9px 26px; background:var(--accent); color:#fff;
|
| 68 |
+
font:600 13px 'Inter',sans-serif; cursor:pointer; min-width:190px; transition:filter .15s; }
|
| 69 |
+
#run:hover:not(:disabled){ filter:brightness(1.12); } #run:disabled{ opacity:.5; cursor:default; }
|
| 70 |
+
.warn { font-size:11px; color:#fbbf24; }
|
| 71 |
+
|
| 72 |
+
.progress { height:5px; background:var(--border); border-radius:999px; overflow:hidden; margin:6px 0 4px; }
|
| 73 |
+
.progress > div { height:100%; width:0; background:var(--accent); transition:width .3s; }
|
| 74 |
+
|
| 75 |
+
.banner { display:flex; flex-wrap:wrap; gap:10px; margin:8px 0 4px; }
|
| 76 |
+
.chip { border:1px solid var(--border); border-radius:8px; padding:6px 12px; font-size:12px;
|
| 77 |
+
color:var(--dim); font-variant-numeric:tabular-nums; }
|
| 78 |
+
.chip b { font-weight:600; color:var(--text); }
|
| 79 |
+
.chip.good b { color:var(--up); } .chip.bad b { color:var(--down); }
|
| 80 |
+
|
| 81 |
+
h2.sec { font-size:11px; text-transform:uppercase; letter-spacing:.12em; color:var(--dim);
|
| 82 |
+
margin:26px 0 8px; font-weight:500; }
|
| 83 |
+
table.matrix { border-collapse:collapse; }
|
| 84 |
+
table.matrix th { color:var(--dim); font-weight:500; font-size:11px; padding:7px 12px; text-align:center; }
|
| 85 |
+
table.matrix th.rowh { text-align:right; color:var(--text); }
|
| 86 |
+
table.matrix td { border:1px solid var(--border); padding:8px 14px; text-align:center; min-width:96px;
|
| 87 |
+
font-variant-numeric:tabular-nums; }
|
| 88 |
+
td.cell { font-weight:600; cursor:default; }
|
| 89 |
+
td .sub { font-size:10px; color:#aab; font-weight:400; margin-top:2px; }
|
| 90 |
+
td.pending { color:#3a4056; }
|
| 91 |
+
td.pending .dot { animation:pulse 1.2s ease-in-out infinite; }
|
| 92 |
+
@keyframes pulse { 0%,100%{opacity:.3} 50%{opacity:.9} }
|
| 93 |
+
|
| 94 |
+
.save-row { margin-top:24px; display:flex; align-items:center; gap:14px; }
|
| 95 |
+
#save { border:1px solid var(--accent); background:transparent; color:#a78bfa; border-radius:999px;
|
| 96 |
+
padding:8px 22px; font:600 12px 'Inter',sans-serif; cursor:pointer; }
|
| 97 |
+
#save:hover:not(:disabled){ background:rgba(139,92,246,.15); } #save:disabled{ opacity:.4; cursor:default; }
|
| 98 |
+
#savemsg { font-size:12px; color:var(--dim); }
|
| 99 |
+
#status { font-size:12px; color:var(--dim); margin-top:6px; }
|
| 100 |
+
#status.err { color:var(--down); }
|
| 101 |
+
.foot { color:var(--dim); font-size:11px; margin-top:30px; }
|
| 102 |
+
|
| 103 |
+
/* leverage optimizer */
|
| 104 |
+
#optcell { background:var(--panel); border:1px solid var(--border); color:var(--text);
|
| 105 |
+
border-radius:8px; padding:7px 12px; font:13px 'Inter',sans-serif; min-width:260px; }
|
| 106 |
+
#optbtn { border:none; border-radius:999px; padding:8px 20px; background:var(--accent); color:#fff;
|
| 107 |
+
font:600 12px 'Inter',sans-serif; cursor:pointer; }
|
| 108 |
+
#optbtn:hover:not(:disabled){ filter:brightness(1.12); } #optbtn:disabled{ opacity:.5; }
|
| 109 |
+
.sigbar { display:flex; flex-wrap:wrap; gap:10px; align-items:center; margin:10px 0 4px; }
|
| 110 |
+
.sigbar .hl { font-size:13px; }
|
| 111 |
+
.warnd { color:var(--down); font-weight:600; }
|
| 112 |
+
.tiers { display:flex; gap:14px; flex-wrap:wrap; margin-top:14px; }
|
| 113 |
+
.tcard { flex:1 1 280px; min-width:262px; border:1px solid var(--border); border-radius:14px;
|
| 114 |
+
padding:14px 16px; background:var(--panel); }
|
| 115 |
+
.tcard h3 { margin:0 0 10px; font-size:14px; display:flex; align-items:center; gap:10px; font-weight:600; }
|
| 116 |
+
.dir { font-size:10px; font-weight:600; padding:2px 8px; border-radius:999px; letter-spacing:.04em; }
|
| 117 |
+
.dir.long { background:rgba(22,199,132,.15); color:var(--up); }
|
| 118 |
+
.dir.short { background:rgba(234,57,67,.15); color:var(--down); }
|
| 119 |
+
.lev { margin-left:auto; font-size:13px; color:#a78bfa; font-weight:600; }
|
| 120 |
+
.tk { display:flex; justify-content:space-between; gap:12px; padding:6px 0; border-top:1px solid var(--border);
|
| 121 |
+
font-variant-numeric:tabular-nums; font-size:12px; }
|
| 122 |
+
.tk .k { color:var(--dim); white-space:nowrap; } .tk .v { text-align:right; font-weight:600; }
|
| 123 |
+
.tk .v .sub { color:#9aa0b4; font-size:10px; font-weight:400; margin-top:1px; }
|
| 124 |
+
.tk.entry .v { color:var(--accent); }
|
| 125 |
+
ul.reason { margin:10px 0 0; padding-left:16px; font-size:11px; color:#9aa0b4; line-height:1.5; }
|
| 126 |
+
ul.reason li { margin-bottom:5px; }
|
| 127 |
+
.disclaimer { margin-top:16px; font-size:11px; color:#fbbf24; border:1px solid rgba(251,191,36,.3);
|
| 128 |
+
border-radius:10px; padding:10px 14px; line-height:1.55; }
|
| 129 |
+
</style>
|
| 130 |
+
</head>
|
| 131 |
+
<body>
|
| 132 |
+
<header>
|
| 133 |
+
<div class="brand">KRONOS<em>.</em><span>portfolio analyzer</span></div>
|
| 134 |
+
<div class="nav"><a href="/">Chart</a><a href="/analyzer" class="on">Analyzer</a></div>
|
| 135 |
+
</header>
|
| 136 |
+
|
| 137 |
+
<main>
|
| 138 |
+
<div class="row">
|
| 139 |
+
<div class="search-wrap">
|
| 140 |
+
<input id="q" placeholder="Search a ticker to analyze — BTC, AAPL, gold, EUR/USD…"
|
| 141 |
+
autocomplete="off" spellcheck="false">
|
| 142 |
+
<div class="dropdown" id="dd" hidden></div>
|
| 143 |
+
</div>
|
| 144 |
+
<div class="staged" id="staged" hidden>
|
| 145 |
+
<span class="badge" id="stagedClass"></span>
|
| 146 |
+
<span id="stagedLabel"></span>
|
| 147 |
+
<button class="x" id="unstage" title="clear">×</button>
|
| 148 |
+
</div>
|
| 149 |
+
</div>
|
| 150 |
+
|
| 151 |
+
<div class="row">
|
| 152 |
+
<span class="lbl">Intervals</span><span id="intervals"></span>
|
| 153 |
+
</div>
|
| 154 |
+
<div class="row">
|
| 155 |
+
<span class="lbl">Horizons</span><span id="horizons"></span>
|
| 156 |
+
</div>
|
| 157 |
+
<div class="row">
|
| 158 |
+
<span class="lbl">Models</span><span id="models"></span>
|
| 159 |
+
<button id="run" disabled>Run Complete Analysis</button>
|
| 160 |
+
</div>
|
| 161 |
+
<div class="row" id="estimate" style="margin-top:-6px"></div>
|
| 162 |
+
|
| 163 |
+
<div class="progress" id="progwrap" hidden><div id="prog"></div></div>
|
| 164 |
+
<div id="status">pick a ticker, choose intervals · horizons · models, then run</div>
|
| 165 |
+
|
| 166 |
+
<div class="banner" id="banner"></div>
|
| 167 |
+
<div id="matrices"></div>
|
| 168 |
+
<div id="divergence"></div>
|
| 169 |
+
|
| 170 |
+
<div id="optpanel" hidden>
|
| 171 |
+
<h2 class="sec">⚡ Leverage Trade Optimizer</h2>
|
| 172 |
+
<div class="row">
|
| 173 |
+
<span class="lbl">Forecast</span>
|
| 174 |
+
<select id="optcell"></select>
|
| 175 |
+
<button id="optbtn">Optimize Trades</button>
|
| 176 |
+
</div>
|
| 177 |
+
<div id="optsignal"></div>
|
| 178 |
+
<div class="tiers" id="opttiers"></div>
|
| 179 |
+
<div class="disclaimer" id="optdisc" hidden>⚠ Educational / research demo — <b>not financial advice</b>.
|
| 180 |
+
Leverage can wipe out your entire margin (and trigger liabilities beyond it) on a small adverse move.
|
| 181 |
+
These setups are a mechanical translation of a probabilistic forecast into trade structure; the model
|
| 182 |
+
is crypto-pretrained and frequently wrong. Treat every level as a hypothesis to stress-test, never a
|
| 183 |
+
recommendation. Size so a full stop-out is a loss you can absorb.</div>
|
| 184 |
+
</div>
|
| 185 |
+
|
| 186 |
+
<div class="save-row" hidden id="saverow">
|
| 187 |
+
<button id="save">Save Analysis</button>
|
| 188 |
+
<span id="savemsg"></span>
|
| 189 |
+
</div>
|
| 190 |
+
|
| 191 |
+
<div class="foot">Each cell: mean expected move over the horizon with its p10–p90 band, from
|
| 192 |
+
independently sampled forecast paths. Research demo · not financial advice.</div>
|
| 193 |
+
</main>
|
| 194 |
+
|
| 195 |
+
<script>
|
| 196 |
+
const ALL_INTERVALS = ['15m','1h','4h','1d'];
|
| 197 |
+
const ALL_HORIZONS = [12,24,48];
|
| 198 |
+
const ALL_MODELS = ['small','base'];
|
| 199 |
+
const MODEL_PARAMS = { small:'24.7M', base:'102.3M' };
|
| 200 |
+
const state = {
|
| 201 |
+
asset:null, staged:null,
|
| 202 |
+
intervals:new Set(['1h']), horizons:new Set([24]), models:new Set(['small']),
|
| 203 |
+
jobId:null, polling:null, results:[], summary:null,
|
| 204 |
+
};
|
| 205 |
+
const $ = (s)=>document.querySelector(s);
|
| 206 |
+
const esc = (s)=>String(s).replace(/[&<>"']/g,c=>({'&':'&','<':'<','>':'>','"':'"',"'":'''}[c]));
|
| 207 |
+
const fp = (v)=>(v>=0?'+':'')+v.toFixed(2)+'%';
|
| 208 |
+
|
| 209 |
+
/* ---------- search & stage (nothing computes until Run) ---------- */
|
| 210 |
+
const qEl=$('#q'), dd=$('#dd');
|
| 211 |
+
let ddItems=[], ddSel=-1, debounce=null, inflight=null;
|
| 212 |
+
qEl.addEventListener('input',()=>{clearTimeout(debounce);debounce=setTimeout(doSearch,280);});
|
| 213 |
+
qEl.addEventListener('keydown',(e)=>{
|
| 214 |
+
if(dd.hidden)return;
|
| 215 |
+
if(e.key==='ArrowDown'){e.preventDefault();move(1);}
|
| 216 |
+
else if(e.key==='ArrowUp'){e.preventDefault();move(-1);}
|
| 217 |
+
else if(e.key==='Enter'){e.preventDefault();if(ddItems.length)stage(ddItems[Math.max(ddSel,0)]);}
|
| 218 |
+
else if(e.key==='Escape')dd.hidden=true;
|
| 219 |
+
});
|
| 220 |
+
qEl.addEventListener('blur',()=>setTimeout(()=>dd.hidden=true,150));
|
| 221 |
+
async function doSearch(){
|
| 222 |
+
const term=qEl.value.trim(); if(!term){dd.hidden=true;return;}
|
| 223 |
+
inflight?.abort(); inflight=new AbortController();
|
| 224 |
+
try{
|
| 225 |
+
const r=await fetch(`/api/search?q=${encodeURIComponent(term)}`,{signal:inflight.signal});
|
| 226 |
+
const j=await r.json(); if(term!==qEl.value.trim())return;
|
| 227 |
+
ddItems=j.results||[]; ddSel=-1; renderDD();
|
| 228 |
+
}catch(e){ if(e.name!=='AbortError')dd.hidden=true; }
|
| 229 |
+
}
|
| 230 |
+
function renderDD(){
|
| 231 |
+
if(!ddItems.length){dd.innerHTML="<div class='none'>no matches</div>";dd.hidden=false;return;}
|
| 232 |
+
dd.innerHTML=ddItems.map((it,i)=>`<div class="opt${i===ddSel?' sel':''}" data-i="${i}">
|
| 233 |
+
<span class="badge" data-k="${esc(it.klass)}">${esc(it.klass)}</span>
|
| 234 |
+
<span class="sym">${esc(it.symbol)}</span><span class="nm">${esc(it.name)}</span>
|
| 235 |
+
<span class="ex">${esc(it.exchange)}</span></div>`).join('');
|
| 236 |
+
dd.hidden=false;
|
| 237 |
+
dd.querySelectorAll('.opt').forEach(el=>el.addEventListener('mousedown',()=>stage(ddItems[+el.dataset.i])));
|
| 238 |
+
}
|
| 239 |
+
function move(d){ddSel=(ddSel+d+ddItems.length)%ddItems.length;renderDD();
|
| 240 |
+
dd.querySelector('.opt.sel')?.scrollIntoView({block:'nearest'});}
|
| 241 |
+
function stage(it){
|
| 242 |
+
state.staged=it; dd.hidden=true; qEl.value='';
|
| 243 |
+
$('#stagedClass').textContent=it.klass; $('#stagedClass').dataset.k=it.klass;
|
| 244 |
+
$('#stagedLabel').textContent=`${it.symbol} · ${it.name}`;
|
| 245 |
+
$('#staged').hidden=false; updateRun();
|
| 246 |
+
setStatus(`${it.symbol} staged — choose options and run`);
|
| 247 |
+
}
|
| 248 |
+
$('#unstage').onclick=()=>{ state.staged=null; $('#staged').hidden=true; updateRun();
|
| 249 |
+
setStatus('pick a ticker, choose intervals · horizons · models, then run'); };
|
| 250 |
+
|
| 251 |
+
/* ---------- multi-select pills ---------- */
|
| 252 |
+
function multipills(elId, items, set, fmt){
|
| 253 |
+
const wrap=$(elId);
|
| 254 |
+
items.forEach(v=>{
|
| 255 |
+
const b=document.createElement('button');
|
| 256 |
+
b.className='pill'+(set.has(v)?' on':''); b.innerHTML=fmt(v);
|
| 257 |
+
b.onclick=()=>{ set.has(v)?set.delete(v):set.add(v); b.classList.toggle('on'); updateRun(); };
|
| 258 |
+
wrap.appendChild(b);
|
| 259 |
+
});
|
| 260 |
+
}
|
| 261 |
+
multipills('#intervals', ALL_INTERVALS, state.intervals, v=>v.toUpperCase());
|
| 262 |
+
multipills('#horizons', ALL_HORIZONS, state.horizons, v=>`${v} bars`);
|
| 263 |
+
multipills('#models', ALL_MODELS, state.models, v=>`${v[0].toUpperCase()+v.slice(1)}<span class="px">${MODEL_PARAMS[v]}</span>`);
|
| 264 |
+
|
| 265 |
+
function updateRun(){
|
| 266 |
+
const ok = state.staged && state.intervals.size && state.horizons.size && state.models.size && !state.polling;
|
| 267 |
+
$('#run').disabled=!ok;
|
| 268 |
+
const nForecasts = state.intervals.size*state.models.size;
|
| 269 |
+
const nCells = nForecasts*state.horizons.size;
|
| 270 |
+
const slow = state.models.has('base');
|
| 271 |
+
$('#estimate').innerHTML = state.staged
|
| 272 |
+
? `<span class="warn">${nForecasts} forecast${nForecasts!==1?'s':''} → ${nCells} cells.`
|
| 273 |
+
+ (slow?' Kronos-base is ~5–8× slower than small on CPU — this can take several minutes.':'')+'</span>'
|
| 274 |
+
: '';
|
| 275 |
+
}
|
| 276 |
+
updateRun();
|
| 277 |
+
|
| 278 |
+
/* ---------- run + poll ---------- */
|
| 279 |
+
function setStatus(m,err=false){ const e=$('#status'); e.textContent=m; e.className=err?'err':''; }
|
| 280 |
+
|
| 281 |
+
$('#run').onclick=async()=>{
|
| 282 |
+
if($('#run').disabled)return;
|
| 283 |
+
const body={ provider:state.staged.provider, symbol:state.staged.symbol,
|
| 284 |
+
intervals:[...state.intervals], horizons:[...state.horizons].sort((a,b)=>a-b),
|
| 285 |
+
models:[...state.models] };
|
| 286 |
+
state.asset=state.staged;
|
| 287 |
+
$('#banner').innerHTML=''; $('#divergence').innerHTML=''; $('#saverow').hidden=true; $('#savemsg').textContent='';
|
| 288 |
+
$('#optpanel').hidden=true; $('#opttiers').innerHTML=''; $('#optsignal').innerHTML='';
|
| 289 |
+
try{
|
| 290 |
+
const r=await fetch('/api/analyze',{method:'POST',headers:{'Content-Type':'application/json'},body:JSON.stringify(body)});
|
| 291 |
+
const j=await r.json(); if(!r.ok)throw new Error(j.error||r.statusText);
|
| 292 |
+
state.jobId=j.job_id; state.results=[]; state.summary=null;
|
| 293 |
+
renderSkeleton(body);
|
| 294 |
+
$('#progwrap').hidden=false; $('#prog').style.width='0%';
|
| 295 |
+
updateRun(); poll();
|
| 296 |
+
}catch(e){ setStatus(e.message,true); }
|
| 297 |
+
};
|
| 298 |
+
|
| 299 |
+
function poll(){
|
| 300 |
+
state.polling=setInterval(async()=>{
|
| 301 |
+
try{
|
| 302 |
+
const r=await fetch(`/api/analyze/${state.jobId}`);
|
| 303 |
+
const j=await r.json(); if(!r.ok)throw new Error(j.error||r.statusText);
|
| 304 |
+
state.results=j.results; state.summary=j.summary;
|
| 305 |
+
fillCells(j);
|
| 306 |
+
const pct=j.total_forecasts?Math.round(j.done_forecasts/j.total_forecasts*100):0;
|
| 307 |
+
$('#prog').style.width=pct+'%';
|
| 308 |
+
if(j.status==='running')
|
| 309 |
+
setStatus(`computing ${esc(j.current)} · ${j.done_forecasts}/${j.total_forecasts} forecasts · ${pct}%`);
|
| 310 |
+
else { clearInterval(state.polling); state.polling=null; updateRun();
|
| 311 |
+
if(j.status==='error'){ setStatus(j.error,true); }
|
| 312 |
+
else { renderSummary(j); setupOptimizer(j); $('#saverow').hidden=false;
|
| 313 |
+
setStatus(`analysis complete · ${j.done_forecasts} forecasts in ${Math.round(j.elapsed_total)}s compute`); } }
|
| 314 |
+
}catch(e){ clearInterval(state.polling); state.polling=null; updateRun(); setStatus(e.message,true); }
|
| 315 |
+
},1500);
|
| 316 |
+
}
|
| 317 |
+
|
| 318 |
+
/* ---------- matrices ---------- */
|
| 319 |
+
function cellColor(d){ const t=Math.max(-1,Math.min(1,d/3)); const a=(0.10+Math.abs(t)*0.45).toFixed(3);
|
| 320 |
+
return t>=0?`rgba(22,199,132,${a})`:`rgba(234,57,67,${a})`; }
|
| 321 |
+
function cid(iv,h,m){ return `c_${m}_${iv}_${h}`.replace(/[^a-z0-9_]/gi,''); }
|
| 322 |
+
|
| 323 |
+
function renderSkeleton(body){
|
| 324 |
+
let html='';
|
| 325 |
+
body.models.forEach(m=>{
|
| 326 |
+
html+=`<h2 class="sec">Kronos-${m} — expected move · p10…p90</h2>`;
|
| 327 |
+
html+='<table class="matrix"><tr><th class="rowh">interval \\ horizon</th>'
|
| 328 |
+
+ body.horizons.map(h=>`<th>${h} bars</th>`).join('')+'</tr>';
|
| 329 |
+
body.intervals.forEach(iv=>{
|
| 330 |
+
html+=`<tr><th class="rowh">${iv.toUpperCase()}</th>`;
|
| 331 |
+
body.horizons.forEach(h=>{ html+=`<td class="pending" id="${cid(iv,h,m)}"><span class="dot">•••</span></td>`; });
|
| 332 |
+
html+='</tr>';
|
| 333 |
+
});
|
| 334 |
+
html+='</table>';
|
| 335 |
+
});
|
| 336 |
+
$('#matrices').innerHTML=html;
|
| 337 |
+
}
|
| 338 |
+
function fillCells(j){
|
| 339 |
+
j.results.forEach(r=>{
|
| 340 |
+
if(r.horizon===undefined)return;
|
| 341 |
+
const td=document.getElementById(cid(r.interval,r.horizon,r.model)); if(!td)return;
|
| 342 |
+
if(r.error){ td.className=''; td.innerHTML=`<span class="sub">err</span>`; td.title=r.error; return; }
|
| 343 |
+
td.className='cell'; td.style.background=cellColor(r.delta_pct);
|
| 344 |
+
td.innerHTML=`<div>${fp(r.delta_pct)}</div><div class="sub">${fp(r.band_lo_pct)}…${fp(r.band_hi_pct)}</div>`;
|
| 345 |
+
td.title=`${r.interval} · ${r.horizon} bars · Kronos-${r.model}\nend close ${r.end_close.toFixed(4)}\n`
|
| 346 |
+
+`expected ${fp(r.delta_pct)} (band ${fp(r.band_lo_pct)}…${fp(r.band_hi_pct)})\ntexture Ra ${r.texture_pct.toFixed(3)}%`;
|
| 347 |
+
});
|
| 348 |
+
}
|
| 349 |
+
|
| 350 |
+
function renderSummary(j){
|
| 351 |
+
const S=j.summary; if(!S||!S.n)return;
|
| 352 |
+
const chips=[];
|
| 353 |
+
chips.push(`<div class="chip good">bullish <b>${S.bullish}</b>/${S.n}</div>`);
|
| 354 |
+
chips.push(`<div class="chip bad">bearish <b>${S.bearish}</b>/${S.n}</div>`);
|
| 355 |
+
chips.push(`<div class="chip">avg move <b class="${S.avg_delta_pct>=0?'':'bad'}">${fp(S.avg_delta_pct)}</b></div>`);
|
| 356 |
+
chips.push(`<div class="chip">avg band <b>${S.avg_band_width_pct.toFixed(2)}%</b></div>`);
|
| 357 |
+
if(S.model_agreement!=null)
|
| 358 |
+
chips.push(`<div class="chip">small/base agree <b>${Math.round(S.model_agreement*100)}%</b></div>`);
|
| 359 |
+
if(S.most_bullish) chips.push(`<div class="chip">most bullish <b class="">${S.most_bullish.interval}·${S.most_bullish.horizon}b·${S.most_bullish.model} ${fp(S.most_bullish.delta_pct)}</b></div>`);
|
| 360 |
+
$('#banner').innerHTML=chips.join('');
|
| 361 |
+
if(S.max_divergence){
|
| 362 |
+
const d=S.max_divergence;
|
| 363 |
+
$('#divergence').innerHTML=`<h2 class="sec">Largest model disagreement</h2>`
|
| 364 |
+
+`<div class="chip">${d.interval.toUpperCase()} · ${d.horizon} bars — `
|
| 365 |
+
+`small <b class="${d.small>=0?'':'bad'}">${fp(d.small)}</b> vs base <b class="${d.base>=0?'':'bad'}">${fp(d.base)}</b> `
|
| 366 |
+
+`(gap ${d.gap.toFixed(2)}%)</div>`;
|
| 367 |
+
}
|
| 368 |
+
}
|
| 369 |
+
|
| 370 |
+
/* ---------- save ---------- */
|
| 371 |
+
$('#save').onclick=async()=>{
|
| 372 |
+
if(!state.jobId)return;
|
| 373 |
+
$('#save').disabled=true; $('#savemsg').textContent='saving…';
|
| 374 |
+
try{
|
| 375 |
+
const r=await fetch(`/api/analyze/${state.jobId}/save`,{method:'POST'});
|
| 376 |
+
const j=await r.json(); if(!r.ok)throw new Error(j.error||r.statusText);
|
| 377 |
+
$('#savemsg').innerHTML=`saved · <a href="/analyses/${encodeURIComponent(j.html)}" target="_blank" style="color:#a78bfa">view report</a>`
|
| 378 |
+
+` · <a id="dljson" href="#" style="color:#a78bfa">download JSON</a>`
|
| 379 |
+
+` · <span style="color:#6b7180">${esc(j.dir)}</span>`;
|
| 380 |
+
$('#dljson').onclick=(e)=>{ e.preventDefault();
|
| 381 |
+
const blob=new Blob([JSON.stringify({symbol:state.asset.symbol,results:state.results,summary:state.summary},null,2)],{type:'application/json'});
|
| 382 |
+
const a=document.createElement('a'); a.href=URL.createObjectURL(blob);
|
| 383 |
+
a.download=`${state.asset.symbol}_analysis.json`; a.click(); };
|
| 384 |
+
}catch(e){ $('#savemsg').textContent=e.message; }
|
| 385 |
+
finally{ $('#save').disabled=false; }
|
| 386 |
+
};
|
| 387 |
+
|
| 388 |
+
/* ---------- leverage trade optimizer ---------- */
|
| 389 |
+
function setupOptimizer(j){
|
| 390 |
+
const groups={};
|
| 391 |
+
j.results.filter(r=>'delta_pct' in r).forEach(r=>{ (groups[r.interval+'|'+r.horizon]=groups[r.interval+'|'+r.horizon]||[]).push(r); });
|
| 392 |
+
const opts=Object.entries(groups).map(([k,cells])=>{
|
| 393 |
+
const [iv,h]=k.split('|');
|
| 394 |
+
const avg=(f)=>cells.reduce((s,c)=>s+c[f],0)/cells.length;
|
| 395 |
+
const mv=avg('delta_pct'), band=Math.max((avg('band_hi_pct')-avg('band_lo_pct'))/2,0.1);
|
| 396 |
+
return { v:`${iv}|${h}`, iv, h:+h, q:Math.abs(mv)/band,
|
| 397 |
+
label:`${iv.toUpperCase()} · ${h} bars — ${fp(mv)} (signal ${(Math.abs(mv)/band).toFixed(2)})` };
|
| 398 |
+
}).sort((a,b)=>b.q-a.q);
|
| 399 |
+
$('#optcell').innerHTML=opts.map(o=>`<option value="${o.v}">${o.label}</option>`).join('');
|
| 400 |
+
$('#optpanel').hidden=false; $('#optdisc').hidden=false;
|
| 401 |
+
}
|
| 402 |
+
function priceFmt(p){ return p>=1000?p.toLocaleString('en-US',{maximumFractionDigits:2})
|
| 403 |
+
:p>=1?p.toFixed(4):Number(p.toPrecision(5)).toString(); }
|
| 404 |
+
$('#optbtn').onclick=async()=>{
|
| 405 |
+
if(!state.jobId)return;
|
| 406 |
+
const [iv,h]=$('#optcell').value.split('|');
|
| 407 |
+
$('#optbtn').disabled=true; $('#optsignal').innerHTML='optimizing…';
|
| 408 |
+
try{
|
| 409 |
+
const r=await fetch('/api/optimize',{method:'POST',headers:{'Content-Type':'application/json'},
|
| 410 |
+
body:JSON.stringify({job_id:state.jobId,interval:iv,horizon:+h})});
|
| 411 |
+
const j=await r.json(); if(!r.ok)throw new Error(j.error||r.statusText);
|
| 412 |
+
renderOptimizer(j);
|
| 413 |
+
}catch(e){ $('#optsignal').innerHTML=`<span class="warnd">${esc(e.message)}</span>`; $('#opttiers').innerHTML=''; }
|
| 414 |
+
finally{ $('#optbtn').disabled=false; }
|
| 415 |
+
};
|
| 416 |
+
function renderOptimizer(j){
|
| 417 |
+
const o=j.optimizer, s=j.signal;
|
| 418 |
+
const pm=Object.entries(s.per_model).map(([m,d])=>`${m} ${fp(d)}`).join(' · ');
|
| 419 |
+
let sig=`<div class="sigbar"><span class="hl">${esc(o.headline)}</span>`;
|
| 420 |
+
if(o.confidence) sig+=`<span class="chip">confidence <b>${o.confidence}</b></span>`;
|
| 421 |
+
sig+=`<span class="chip">models: ${esc(pm)}</span>`;
|
| 422 |
+
if(o.disagree) sig+=`<span class="chip"><span class="warnd">⚠ small & base disagree — no consensus</span></span>`;
|
| 423 |
+
sig+='</div>';
|
| 424 |
+
$('#optsignal').innerHTML=sig;
|
| 425 |
+
$('#opttiers').innerHTML=(o.tiers||[]).map(tierCard).join('');
|
| 426 |
+
}
|
| 427 |
+
function tierCard(t){
|
| 428 |
+
const line=(label,price,sub,cls='')=>`<div class="tk ${cls}"><span class="k">${label}</span>`
|
| 429 |
+
+`<span class="v">${priceFmt(price)}<div class="sub">${sub}</div></span></div>`;
|
| 430 |
+
const tps=t.targets.map((x,i)=>line(`Take profit ${i+1}`,x.price,
|
| 431 |
+
`${fp(x.pct)} · R:R ${x.rr.toFixed(2)} · +${x.gain_margin_pct.toFixed(0)}% margin`)).join('');
|
| 432 |
+
return `<div class="tcard">
|
| 433 |
+
<h3>${t.name}<span class="dir ${t.direction}">${t.direction.toUpperCase()}</span><span class="lev">${t.leverage}×</span></h3>
|
| 434 |
+
${line('Entry',t.entry,esc(t.entry_kind),'entry')}
|
| 435 |
+
${line('Stop loss',t.stop,`-${t.stop_pct.toFixed(2)}% · -${t.stop_margin_loss_pct.toFixed(0)}% margin`)}
|
| 436 |
+
${tps}
|
| 437 |
+
${line('Liquidation ≈',t.liq,`-${t.liq_pct.toFixed(1)}% away`)}
|
| 438 |
+
<ul class="reason">${t.reasoning.map(x=>`<li>${esc(x)}</li>`).join('')}</ul>
|
| 439 |
+
</div>`;
|
| 440 |
+
}
|
| 441 |
+
|
| 442 |
+
qEl.focus();
|
| 443 |
+
</script>
|
| 444 |
+
</body>
|
| 445 |
+
</html>
|
crypto_ui/app.py
ADDED
|
@@ -0,0 +1,857 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Minimal multi-asset forecasting UI powered by Kronos.
|
| 2 |
+
|
| 3 |
+
Flask server that:
|
| 4 |
+
* searches assets across classes — crypto via Binance's full spot universe,
|
| 5 |
+
equities/ETFs/forex/indices/commodities/funds via Yahoo Finance;
|
| 6 |
+
* fetches live OHLCV from the matching provider (Yahoo timestamps shifted to
|
| 7 |
+
exchange-local time; 4h bars resampled from 1h since Yahoo lacks them);
|
| 8 |
+
* runs Kronos (small or base, CPU) for probabilistic candlestick forecasts:
|
| 9 |
+
N sampled paths -> mean candles + p10/p90 close band, with future
|
| 10 |
+
timestamps generated session-aware so stock forecasts skip closed hours;
|
| 11 |
+
* serves a single-page UI. Data is only fetched on explicit user action.
|
| 12 |
+
"""
|
| 13 |
+
import datetime
|
| 14 |
+
import json
|
| 15 |
+
import os
|
| 16 |
+
import re
|
| 17 |
+
import sys
|
| 18 |
+
import threading
|
| 19 |
+
import time
|
| 20 |
+
import urllib.parse
|
| 21 |
+
import uuid
|
| 22 |
+
from pathlib import Path
|
| 23 |
+
|
| 24 |
+
import numpy as np
|
| 25 |
+
import pandas as pd
|
| 26 |
+
import requests
|
| 27 |
+
from flask import Flask, jsonify, request, send_file
|
| 28 |
+
|
| 29 |
+
REPO_ROOT = Path(__file__).resolve().parent.parent / "Kronos"
|
| 30 |
+
sys.path.insert(0, str(REPO_ROOT))
|
| 31 |
+
sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "roughness_lab"))
|
| 32 |
+
|
| 33 |
+
from model import Kronos, KronosTokenizer, KronosPredictor
|
| 34 |
+
|
| 35 |
+
try: # surface-roughness texture metric for forecasts (ties into roughness_lab)
|
| 36 |
+
from roughness import roughness_params
|
| 37 |
+
except Exception:
|
| 38 |
+
roughness_params = None
|
| 39 |
+
|
| 40 |
+
BINANCE_HOSTS = [
|
| 41 |
+
"https://data-api.binance.vision", # public market-data domain
|
| 42 |
+
"https://api.binance.com",
|
| 43 |
+
]
|
| 44 |
+
YAHOO_UA = {"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64)"}
|
| 45 |
+
YAHOO_CLASS = {
|
| 46 |
+
"EQUITY": "Equity", "ETF": "ETF", "INDEX": "Index", "CURRENCY": "Forex",
|
| 47 |
+
"FUTURE": "Commodity", "MUTUALFUND": "Fund", "CRYPTOCURRENCY": "Crypto",
|
| 48 |
+
}
|
| 49 |
+
INTERVALS = {"15m": 15 * 60, "1h": 3600, "4h": 4 * 3600, "1d": 24 * 3600}
|
| 50 |
+
YAHOO_FETCH = { # ui interval -> (yahoo interval, range); 4h is resampled from 1h
|
| 51 |
+
"15m": ("15m", "60d"),
|
| 52 |
+
"1h": ("60m", "730d"),
|
| 53 |
+
"4h": ("60m", "730d"),
|
| 54 |
+
"1d": ("1d", "5y"),
|
| 55 |
+
}
|
| 56 |
+
LOOKBACK = 400 # context bars fed to the model (max_context is 512)
|
| 57 |
+
MAX_HORIZON = 96
|
| 58 |
+
N_PATHS = 5 # sampled forecast paths per request
|
| 59 |
+
ANALYZER_PATHS = 4 # sampled paths per cell in the portfolio analyzer matrix
|
| 60 |
+
ANALYZER_T = 1.0
|
| 61 |
+
ANALYZER_TOP_P = 0.9
|
| 62 |
+
ANALYSES_DIR = Path(__file__).resolve().parent / "analyses"
|
| 63 |
+
|
| 64 |
+
MODELS = {
|
| 65 |
+
"small": {"model_id": "NeoQuasar/Kronos-small", "tokenizer_id": "NeoQuasar/Kronos-Tokenizer-base"},
|
| 66 |
+
"base": {"model_id": "NeoQuasar/Kronos-base", "tokenizer_id": "NeoQuasar/Kronos-Tokenizer-base"},
|
| 67 |
+
}
|
| 68 |
+
|
| 69 |
+
app = Flask(__name__)
|
| 70 |
+
_predict_lock = threading.Lock()
|
| 71 |
+
_load_lock = threading.Lock()
|
| 72 |
+
_predictors: dict = {}
|
| 73 |
+
|
| 74 |
+
|
| 75 |
+
def get_predictor(name: str) -> KronosPredictor:
|
| 76 |
+
"""Return the predictor for a model size, loading it on first use."""
|
| 77 |
+
with _load_lock:
|
| 78 |
+
if name not in _predictors:
|
| 79 |
+
cfg = MODELS[name]
|
| 80 |
+
print(f"Loading Kronos-{name} (CPU)...")
|
| 81 |
+
tok = KronosTokenizer.from_pretrained(cfg["tokenizer_id"])
|
| 82 |
+
mdl = Kronos.from_pretrained(cfg["model_id"])
|
| 83 |
+
tok.eval()
|
| 84 |
+
mdl.eval()
|
| 85 |
+
_predictors[name] = KronosPredictor(mdl, tok, device="cpu", max_context=512)
|
| 86 |
+
print(f"Kronos-{name} ready.")
|
| 87 |
+
return _predictors[name]
|
| 88 |
+
|
| 89 |
+
|
| 90 |
+
get_predictor("small") # warm the default model at startup; "base" loads lazily
|
| 91 |
+
|
| 92 |
+
|
| 93 |
+
# --------------------------------------------------------------------------
|
| 94 |
+
# Asset search
|
| 95 |
+
# --------------------------------------------------------------------------
|
| 96 |
+
_binance_cache = {"ts": 0.0, "symbols": []}
|
| 97 |
+
|
| 98 |
+
|
| 99 |
+
def binance_universe() -> list:
|
| 100 |
+
"""All TRADING spot symbols on Binance, cached for an hour."""
|
| 101 |
+
if time.time() - _binance_cache["ts"] > 3600 or not _binance_cache["symbols"]:
|
| 102 |
+
for host in BINANCE_HOSTS:
|
| 103 |
+
try:
|
| 104 |
+
r = requests.get(f"{host}/api/v3/exchangeInfo", timeout=20)
|
| 105 |
+
r.raise_for_status()
|
| 106 |
+
_binance_cache["symbols"] = [
|
| 107 |
+
{"symbol": s["symbol"], "base": s["baseAsset"], "quote": s["quoteAsset"]}
|
| 108 |
+
for s in r.json()["symbols"]
|
| 109 |
+
if s.get("status") == "TRADING" and s.get("isSpotTradingAllowed")
|
| 110 |
+
]
|
| 111 |
+
_binance_cache["ts"] = time.time()
|
| 112 |
+
break
|
| 113 |
+
except Exception:
|
| 114 |
+
continue
|
| 115 |
+
return _binance_cache["symbols"]
|
| 116 |
+
|
| 117 |
+
|
| 118 |
+
def search_binance(q: str, limit: int = 8) -> list:
|
| 119 |
+
q = q.upper().replace("/", "")
|
| 120 |
+
quote_rank = {"USDT": 0, "USDC": 1, "FDUSD": 2, "BTC": 3, "ETH": 4}
|
| 121 |
+
scored = []
|
| 122 |
+
for s in binance_universe():
|
| 123 |
+
sym, base = s["symbol"], s["base"]
|
| 124 |
+
if q == base:
|
| 125 |
+
score = 0
|
| 126 |
+
elif base.startswith(q):
|
| 127 |
+
score = 1
|
| 128 |
+
elif q == sym:
|
| 129 |
+
score = 2
|
| 130 |
+
elif sym.startswith(q):
|
| 131 |
+
score = 3
|
| 132 |
+
elif q in sym:
|
| 133 |
+
score = 4
|
| 134 |
+
else:
|
| 135 |
+
continue
|
| 136 |
+
scored.append((score, quote_rank.get(s["quote"], 5), len(sym), s))
|
| 137 |
+
scored.sort(key=lambda t: t[:3])
|
| 138 |
+
return [
|
| 139 |
+
{"provider": "binance", "symbol": s["symbol"], "name": f"{s['base']}/{s['quote']}",
|
| 140 |
+
"klass": "Crypto", "exchange": "Binance"}
|
| 141 |
+
for *_, s in scored[:limit]
|
| 142 |
+
]
|
| 143 |
+
|
| 144 |
+
|
| 145 |
+
def search_yahoo(q: str, limit: int = 8) -> list:
|
| 146 |
+
try:
|
| 147 |
+
r = requests.get(
|
| 148 |
+
"https://query1.finance.yahoo.com/v1/finance/search",
|
| 149 |
+
params={"q": q, "quotesCount": limit, "newsCount": 0},
|
| 150 |
+
headers=YAHOO_UA, timeout=10,
|
| 151 |
+
)
|
| 152 |
+
r.raise_for_status()
|
| 153 |
+
quotes = r.json().get("quotes", [])
|
| 154 |
+
except Exception:
|
| 155 |
+
return []
|
| 156 |
+
out = []
|
| 157 |
+
for it in quotes:
|
| 158 |
+
sym, qt = it.get("symbol"), it.get("quoteType", "")
|
| 159 |
+
if not sym or qt == "CRYPTOCURRENCY": # crypto is served by Binance
|
| 160 |
+
continue
|
| 161 |
+
out.append({
|
| 162 |
+
"provider": "yahoo", "symbol": sym,
|
| 163 |
+
"name": it.get("shortname") or it.get("longname") or sym,
|
| 164 |
+
"klass": YAHOO_CLASS.get(qt, qt.title() or "Other"),
|
| 165 |
+
"exchange": it.get("exchDisp") or it.get("exchange") or "Yahoo",
|
| 166 |
+
})
|
| 167 |
+
return out
|
| 168 |
+
|
| 169 |
+
|
| 170 |
+
# --------------------------------------------------------------------------
|
| 171 |
+
# Market data
|
| 172 |
+
# --------------------------------------------------------------------------
|
| 173 |
+
def validate(provider: str, symbol: str, interval: str) -> None:
|
| 174 |
+
if interval not in INTERVALS:
|
| 175 |
+
raise ValueError("Invalid interval")
|
| 176 |
+
if provider == "binance":
|
| 177 |
+
if not re.fullmatch(r"[A-Z0-9]{5,14}", symbol):
|
| 178 |
+
raise ValueError("Invalid symbol")
|
| 179 |
+
elif provider == "yahoo":
|
| 180 |
+
if not re.fullmatch(r"[A-Za-z0-9.^=\-]{1,20}", symbol):
|
| 181 |
+
raise ValueError("Invalid symbol")
|
| 182 |
+
else:
|
| 183 |
+
raise ValueError("Invalid provider")
|
| 184 |
+
|
| 185 |
+
|
| 186 |
+
def fetch_binance(symbol: str, interval: str, limit: int) -> pd.DataFrame:
|
| 187 |
+
params = {"symbol": symbol, "interval": interval, "limit": min(limit + 1, 1000)}
|
| 188 |
+
last_err = None
|
| 189 |
+
for host in BINANCE_HOSTS:
|
| 190 |
+
try:
|
| 191 |
+
r = requests.get(f"{host}/api/v3/klines", params=params, timeout=15)
|
| 192 |
+
if r.status_code == 400:
|
| 193 |
+
raise ValueError(f"Unknown symbol {symbol}")
|
| 194 |
+
r.raise_for_status()
|
| 195 |
+
rows = r.json()
|
| 196 |
+
break
|
| 197 |
+
except ValueError:
|
| 198 |
+
raise
|
| 199 |
+
except Exception as e:
|
| 200 |
+
last_err = e
|
| 201 |
+
else:
|
| 202 |
+
raise RuntimeError(f"Market data unavailable: {last_err}")
|
| 203 |
+
|
| 204 |
+
df = pd.DataFrame(
|
| 205 |
+
[
|
| 206 |
+
{
|
| 207 |
+
"time": int(row[0] // 1000),
|
| 208 |
+
"open": float(row[1]), "high": float(row[2]),
|
| 209 |
+
"low": float(row[3]), "close": float(row[4]),
|
| 210 |
+
"volume": float(row[5]),
|
| 211 |
+
"amount": float(row[7]), # quote-asset volume
|
| 212 |
+
}
|
| 213 |
+
for row in rows
|
| 214 |
+
]
|
| 215 |
+
)
|
| 216 |
+
# Drop the still-forming newest candle so the context is closed bars only.
|
| 217 |
+
if len(df) and df["time"].iloc[-1] + INTERVALS[interval] > time.time():
|
| 218 |
+
df = df.iloc[:-1]
|
| 219 |
+
return df.tail(limit).reset_index(drop=True)
|
| 220 |
+
|
| 221 |
+
|
| 222 |
+
def fetch_yahoo(symbol: str, interval: str, limit: int) -> pd.DataFrame:
|
| 223 |
+
y_itv, y_range = YAHOO_FETCH[interval]
|
| 224 |
+
url = f"https://query1.finance.yahoo.com/v8/finance/chart/{urllib.parse.quote(symbol)}"
|
| 225 |
+
r = requests.get(
|
| 226 |
+
url,
|
| 227 |
+
params={"interval": y_itv, "range": y_range, "includePrePost": "false"},
|
| 228 |
+
headers=YAHOO_UA, timeout=20,
|
| 229 |
+
)
|
| 230 |
+
if r.status_code in (400, 404):
|
| 231 |
+
raise ValueError(f"Unknown symbol {symbol}")
|
| 232 |
+
r.raise_for_status()
|
| 233 |
+
chart = r.json()["chart"]
|
| 234 |
+
if chart.get("error"):
|
| 235 |
+
raise ValueError(chart["error"].get("description", "Yahoo error"))
|
| 236 |
+
res = chart["result"][0]
|
| 237 |
+
ts = res.get("timestamp") or []
|
| 238 |
+
if not ts:
|
| 239 |
+
raise ValueError(f"No data for {symbol} at this interval")
|
| 240 |
+
quote = res["indicators"]["quote"][0]
|
| 241 |
+
gmtoff = int(res["meta"].get("gmtoffset", 0))
|
| 242 |
+
|
| 243 |
+
# Shift to exchange-local time: sessions render sanely and the model's
|
| 244 |
+
# temporal embedding sees natural trading hours.
|
| 245 |
+
df = pd.DataFrame({
|
| 246 |
+
"time": [int(t) + gmtoff for t in ts],
|
| 247 |
+
"open": quote["open"], "high": quote["high"],
|
| 248 |
+
"low": quote["low"], "close": quote["close"],
|
| 249 |
+
"volume": quote["volume"],
|
| 250 |
+
})
|
| 251 |
+
df = df.dropna(subset=["open", "high", "low", "close"]).reset_index(drop=True)
|
| 252 |
+
df["volume"] = df["volume"].astype(float).fillna(0.0)
|
| 253 |
+
df["amount"] = df["volume"] * df[["open", "high", "low", "close"]].mean(axis=1)
|
| 254 |
+
|
| 255 |
+
if interval == "4h": # Yahoo has no native 4h: aggregate hourly bars
|
| 256 |
+
d = df.set_index(pd.to_datetime(df["time"], unit="s"))
|
| 257 |
+
d = d.resample("4h").agg(
|
| 258 |
+
{"open": "first", "high": "max", "low": "min",
|
| 259 |
+
"close": "last", "volume": "sum", "amount": "sum"}
|
| 260 |
+
).dropna(subset=["open"])
|
| 261 |
+
d["time"] = d.index.astype("int64") // 10 ** 9
|
| 262 |
+
df = d.reset_index(drop=True)[["time", "open", "high", "low", "close", "volume", "amount"]]
|
| 263 |
+
|
| 264 |
+
# Drop the still-forming bar (compare in exchange-local time).
|
| 265 |
+
if len(df) and df["time"].iloc[-1] + INTERVALS[interval] > time.time() + gmtoff:
|
| 266 |
+
df = df.iloc[:-1]
|
| 267 |
+
|
| 268 |
+
# Yahoo appends a session-close snapshot (e.g. a 16:00 bar after the
|
| 269 |
+
# 15:30 hourly bar). Its phase is off the bar grid and would corrupt
|
| 270 |
+
# both the chart and the forecast-timestamp pattern, so drop trailing
|
| 271 |
+
# bars that are neither step-spaced nor phase-aligned with the grid.
|
| 272 |
+
step = INTERVALS[interval]
|
| 273 |
+
while len(df) > 1:
|
| 274 |
+
t_last, t_prev = int(df["time"].iloc[-1]), int(df["time"].iloc[-2])
|
| 275 |
+
if t_last - t_prev == step or t_last % step == t_prev % step:
|
| 276 |
+
break
|
| 277 |
+
df = df.iloc[:-1]
|
| 278 |
+
return df.tail(limit).reset_index(drop=True)
|
| 279 |
+
|
| 280 |
+
|
| 281 |
+
def fetch(provider: str, symbol: str, interval: str, limit: int = LOOKBACK) -> pd.DataFrame:
|
| 282 |
+
if provider == "binance":
|
| 283 |
+
return fetch_binance(symbol, interval, limit)
|
| 284 |
+
return fetch_yahoo(symbol, interval, limit)
|
| 285 |
+
|
| 286 |
+
|
| 287 |
+
def future_timestamps(times: pd.Series, horizon: int, step: int) -> list:
|
| 288 |
+
"""Continue the series' own session pattern: only emit future times whose
|
| 289 |
+
(weekday, time-of-day) slot occurs in history. 24/7 markets pass through
|
| 290 |
+
unchanged; stock forecasts skip nights, weekends and other closed hours."""
|
| 291 |
+
dt = pd.to_datetime(times, unit="s")
|
| 292 |
+
counts = pd.Series(list(zip(dt.dt.weekday, dt.dt.hour, dt.dt.minute))).value_counts()
|
| 293 |
+
# Ignore one-off slots (data anomalies, half-days) so they can't skew the pattern.
|
| 294 |
+
slots = set(counts[counts >= 2].index) or set(counts.index)
|
| 295 |
+
out, t = [], int(times.iloc[-1])
|
| 296 |
+
for _ in range(horizon * 80):
|
| 297 |
+
t += step
|
| 298 |
+
d = pd.Timestamp(t, unit="s")
|
| 299 |
+
if (d.weekday(), d.hour, d.minute) in slots:
|
| 300 |
+
out.append(t)
|
| 301 |
+
if len(out) == horizon:
|
| 302 |
+
break
|
| 303 |
+
while len(out) < horizon: # safety net for sparse/irregular histories
|
| 304 |
+
out.append((out[-1] if out else int(times.iloc[-1])) + step)
|
| 305 |
+
return out
|
| 306 |
+
|
| 307 |
+
|
| 308 |
+
# --------------------------------------------------------------------------
|
| 309 |
+
# Routes
|
| 310 |
+
# --------------------------------------------------------------------------
|
| 311 |
+
@app.route("/")
|
| 312 |
+
def index():
|
| 313 |
+
return send_file(Path(__file__).resolve().parent / "index.html")
|
| 314 |
+
|
| 315 |
+
|
| 316 |
+
@app.route("/api/search")
|
| 317 |
+
def api_search():
|
| 318 |
+
q = request.args.get("q", "").strip()
|
| 319 |
+
if not q:
|
| 320 |
+
return jsonify({"results": []})
|
| 321 |
+
return jsonify({"results": (search_binance(q) + search_yahoo(q))[:14]})
|
| 322 |
+
|
| 323 |
+
|
| 324 |
+
@app.route("/api/klines")
|
| 325 |
+
def api_klines():
|
| 326 |
+
provider = request.args.get("provider", "binance")
|
| 327 |
+
symbol = request.args.get("symbol", "").strip()
|
| 328 |
+
interval = request.args.get("interval", "1h")
|
| 329 |
+
try:
|
| 330 |
+
validate(provider, symbol, interval)
|
| 331 |
+
df = fetch(provider, symbol, interval)
|
| 332 |
+
except ValueError as e:
|
| 333 |
+
return jsonify({"error": str(e)}), 400
|
| 334 |
+
except Exception as e:
|
| 335 |
+
return jsonify({"error": str(e)}), 502
|
| 336 |
+
return jsonify({
|
| 337 |
+
"symbol": symbol,
|
| 338 |
+
"interval": interval,
|
| 339 |
+
"candles": df.drop(columns="amount").to_dict("records"),
|
| 340 |
+
})
|
| 341 |
+
|
| 342 |
+
|
| 343 |
+
@app.route("/api/predict", methods=["POST"])
|
| 344 |
+
def api_predict():
|
| 345 |
+
body = request.get_json(force=True)
|
| 346 |
+
provider = str(body.get("provider", "binance"))
|
| 347 |
+
symbol = str(body.get("symbol", "")).strip()
|
| 348 |
+
interval = str(body.get("interval", "1h"))
|
| 349 |
+
horizon = int(body.get("horizon", 24))
|
| 350 |
+
model_name = str(body.get("model", "small"))
|
| 351 |
+
|
| 352 |
+
try:
|
| 353 |
+
validate(provider, symbol, interval)
|
| 354 |
+
if not 1 <= horizon <= MAX_HORIZON:
|
| 355 |
+
raise ValueError(f"Horizon must be 1..{MAX_HORIZON}")
|
| 356 |
+
if model_name not in MODELS:
|
| 357 |
+
raise ValueError(f"Unknown model '{model_name}'")
|
| 358 |
+
df = fetch(provider, symbol, interval)
|
| 359 |
+
except ValueError as e:
|
| 360 |
+
return jsonify({"error": str(e)}), 400
|
| 361 |
+
except Exception as e:
|
| 362 |
+
return jsonify({"error": str(e)}), 502
|
| 363 |
+
|
| 364 |
+
if len(df) < 64:
|
| 365 |
+
return jsonify({"error": "Not enough history for this asset/interval"}), 400
|
| 366 |
+
|
| 367 |
+
fut = future_timestamps(df["time"], horizon, INTERVALS[interval])
|
| 368 |
+
x_df = df[["open", "high", "low", "close", "volume", "amount"]]
|
| 369 |
+
x_ts = pd.Series(pd.to_datetime(df["time"], unit="s"))
|
| 370 |
+
y_ts = pd.Series(pd.to_datetime(fut, unit="s"))
|
| 371 |
+
|
| 372 |
+
predictor = get_predictor(model_name) # may download/load on first use
|
| 373 |
+
t0 = time.time()
|
| 374 |
+
with _predict_lock:
|
| 375 |
+
pred_dfs = predictor.predict_batch(
|
| 376 |
+
df_list=[x_df] * N_PATHS,
|
| 377 |
+
x_timestamp_list=[x_ts] * N_PATHS,
|
| 378 |
+
y_timestamp_list=[y_ts] * N_PATHS,
|
| 379 |
+
pred_len=horizon,
|
| 380 |
+
T=1.0,
|
| 381 |
+
top_p=0.9,
|
| 382 |
+
sample_count=1,
|
| 383 |
+
verbose=False,
|
| 384 |
+
)
|
| 385 |
+
elapsed = time.time() - t0
|
| 386 |
+
|
| 387 |
+
cols = ["open", "high", "low", "close", "volume"]
|
| 388 |
+
paths = np.stack([p[cols].to_numpy(dtype=np.float64) for p in pred_dfs]) # (N, H, 5)
|
| 389 |
+
paths[:, :, 4] = np.clip(paths[:, :, 4], 0, None)
|
| 390 |
+
mean = paths.mean(axis=0)
|
| 391 |
+
|
| 392 |
+
fc_candles = []
|
| 393 |
+
for i in range(horizon):
|
| 394 |
+
o, h, l, c, v = mean[i]
|
| 395 |
+
h, l = max(h, o, c), min(l, o, c) # keep wicks enclosing the body after averaging
|
| 396 |
+
fc_candles.append({
|
| 397 |
+
"time": fut[i],
|
| 398 |
+
"open": float(o), "high": float(h), "low": float(l),
|
| 399 |
+
"close": float(c), "volume": float(v),
|
| 400 |
+
})
|
| 401 |
+
|
| 402 |
+
close_paths = paths[:, :, 3]
|
| 403 |
+
p10 = np.percentile(close_paths, 10, axis=0)
|
| 404 |
+
p90 = np.percentile(close_paths, 90, axis=0)
|
| 405 |
+
|
| 406 |
+
last_close = float(df["close"].iloc[-1])
|
| 407 |
+
pct = lambda v: (v / last_close - 1.0) * 100.0
|
| 408 |
+
|
| 409 |
+
return jsonify({
|
| 410 |
+
"context": df.drop(columns="amount").to_dict("records"),
|
| 411 |
+
"forecast": {
|
| 412 |
+
"candles": fc_candles,
|
| 413 |
+
"p10": [{"time": t, "value": float(v)} for t, v in zip(fut, p10)],
|
| 414 |
+
"p90": [{"time": t, "value": float(v)} for t, v in zip(fut, p90)],
|
| 415 |
+
},
|
| 416 |
+
"stats": {
|
| 417 |
+
"last_close": last_close,
|
| 418 |
+
"end_close": float(mean[-1, 3]),
|
| 419 |
+
"delta_pct": pct(float(mean[-1, 3])),
|
| 420 |
+
"band_lo_pct": pct(float(p10[-1])),
|
| 421 |
+
"band_hi_pct": pct(float(p90[-1])),
|
| 422 |
+
"paths": N_PATHS,
|
| 423 |
+
"model": model_name,
|
| 424 |
+
"elapsed_s": round(elapsed, 1),
|
| 425 |
+
},
|
| 426 |
+
})
|
| 427 |
+
|
| 428 |
+
|
| 429 |
+
# --------------------------------------------------------------------------
|
| 430 |
+
# Portfolio analyzer: forecast a ticker across intervals x horizons x models
|
| 431 |
+
# --------------------------------------------------------------------------
|
| 432 |
+
_jobs: dict = {}
|
| 433 |
+
_jobs_lock = threading.Lock()
|
| 434 |
+
|
| 435 |
+
# Standalone matrix renderer embedded in saved HTML reports.
|
| 436 |
+
_REPORT_JS = r"""
|
| 437 |
+
function cellColor(d){const t=Math.max(-1,Math.min(1,d/3));
|
| 438 |
+
const a=(0.10+Math.abs(t)*0.45).toFixed(3);
|
| 439 |
+
return t>=0?`rgba(22,199,132,${a})`:`rgba(234,57,67,${a})`;}
|
| 440 |
+
function fp(v){return (v>=0?'+':'')+v.toFixed(2)+'%';}
|
| 441 |
+
function render(){const app=document.getElementById('app');const S=DATA.summary||{};let h='';
|
| 442 |
+
if(S.n){h+="<h2>Consensus</h2><div class='meta'>"+S.bullish+" bullish · "+S.bearish+
|
| 443 |
+
" bearish of "+S.n+" · avg "+fp(S.avg_delta_pct)+" · avg band "+S.avg_band_width_pct.toFixed(2)+"%";
|
| 444 |
+
if(S.model_agreement!=null)h+=" · models agree "+Math.round(S.model_agreement*100)+"%";h+="</div>";}
|
| 445 |
+
const idx={};DATA.results.forEach(r=>{idx[r.interval+'|'+r.horizon+'|'+r.model]=r;});
|
| 446 |
+
DATA.models.forEach(m=>{h+="<h2>Kronos-"+m+" — expected move (p10…p90)</h2>";
|
| 447 |
+
h+="<table><tr><th>interval \\ horizon</th>"+DATA.horizons.map(x=>"<th>"+x+" bars</th>").join('')+"</tr>";
|
| 448 |
+
DATA.intervals.forEach(iv=>{h+="<tr><th>"+iv+"</th>";DATA.horizons.forEach(x=>{
|
| 449 |
+
const r=idx[iv+'|'+x+'|'+m];
|
| 450 |
+
if(!r||r.error){h+="<td class='sub'>"+(r&&r.error?'err':'—')+"</td>";return;}
|
| 451 |
+
h+="<td style='background:"+cellColor(r.delta_pct)+"'><div class='cell'>"+fp(r.delta_pct)+
|
| 452 |
+
"</div><div class='sub'>"+fp(r.band_lo_pct)+"…"+fp(r.band_hi_pct)+"</div></td>";});h+="</tr>";});
|
| 453 |
+
h+="</table>";});
|
| 454 |
+
app.innerHTML=h;}
|
| 455 |
+
render();
|
| 456 |
+
"""
|
| 457 |
+
|
| 458 |
+
|
| 459 |
+
def _texture_pct(close_paths: np.ndarray) -> float:
|
| 460 |
+
"""Median surface-roughness Ra (%) across forecast paths — 'predicted
|
| 461 |
+
choppiness'. Falls back to std of per-bar % moves if roughness_lab is
|
| 462 |
+
unavailable or the path is too short."""
|
| 463 |
+
vals = []
|
| 464 |
+
for path in close_paths:
|
| 465 |
+
if roughness_params is not None and len(path) >= 6:
|
| 466 |
+
try:
|
| 467 |
+
vals.append(roughness_params(path, max(4, len(path) // 6)).ra)
|
| 468 |
+
continue
|
| 469 |
+
except Exception:
|
| 470 |
+
pass
|
| 471 |
+
rets = np.diff(path) / path[:-1] * 100.0
|
| 472 |
+
vals.append(float(np.std(rets)) if len(rets) else 0.0)
|
| 473 |
+
return float(np.median(vals)) if vals else 0.0
|
| 474 |
+
|
| 475 |
+
|
| 476 |
+
def _sign(x: float, flat: float = 0.05) -> int:
|
| 477 |
+
return 1 if x > flat else (-1 if x < -flat else 0)
|
| 478 |
+
|
| 479 |
+
|
| 480 |
+
def _summarize(results: list, intervals: list, horizons: list, models: list) -> dict:
|
| 481 |
+
cells = [r for r in results if "delta_pct" in r]
|
| 482 |
+
if not cells:
|
| 483 |
+
return {"n": 0}
|
| 484 |
+
deltas = [r["delta_pct"] for r in cells]
|
| 485 |
+
bull = sum(1 for d in deltas if d > 0.05)
|
| 486 |
+
bear = sum(1 for d in deltas if d < -0.05)
|
| 487 |
+
agreement, max_div = None, None
|
| 488 |
+
if "small" in models and "base" in models:
|
| 489 |
+
idx = {(r["interval"], r["horizon"], r["model"]): r for r in cells}
|
| 490 |
+
agree = comp = 0
|
| 491 |
+
divs = []
|
| 492 |
+
for iv in intervals:
|
| 493 |
+
for h in horizons:
|
| 494 |
+
a, b = idx.get((iv, h, "small")), idx.get((iv, h, "base"))
|
| 495 |
+
if a and b:
|
| 496 |
+
comp += 1
|
| 497 |
+
if _sign(a["delta_pct"]) == _sign(b["delta_pct"]):
|
| 498 |
+
agree += 1
|
| 499 |
+
divs.append({"interval": iv, "horizon": h,
|
| 500 |
+
"small": a["delta_pct"], "base": b["delta_pct"],
|
| 501 |
+
"gap": abs(a["delta_pct"] - b["delta_pct"])})
|
| 502 |
+
agreement = (agree / comp) if comp else None
|
| 503 |
+
max_div = max(divs, key=lambda d: d["gap"]) if divs else None
|
| 504 |
+
most_bull = max(cells, key=lambda r: r["delta_pct"])
|
| 505 |
+
most_bear = min(cells, key=lambda r: r["delta_pct"])
|
| 506 |
+
return {
|
| 507 |
+
"n": len(cells), "bullish": bull, "bearish": bear,
|
| 508 |
+
"avg_delta_pct": float(np.mean(deltas)),
|
| 509 |
+
"median_delta_pct": float(np.median(deltas)),
|
| 510 |
+
"avg_band_width_pct": float(np.mean([r["band_hi_pct"] - r["band_lo_pct"] for r in cells])),
|
| 511 |
+
"model_agreement": agreement,
|
| 512 |
+
"max_divergence": max_div,
|
| 513 |
+
"most_bullish": {k: most_bull[k] for k in ("interval", "horizon", "model", "delta_pct")},
|
| 514 |
+
"most_bearish": {k: most_bear[k] for k in ("interval", "horizon", "model", "delta_pct")},
|
| 515 |
+
}
|
| 516 |
+
|
| 517 |
+
|
| 518 |
+
def _run_analysis(job_id, provider, symbol, intervals, horizons, models):
|
| 519 |
+
job = _jobs[job_id]
|
| 520 |
+
try:
|
| 521 |
+
max_h = max(horizons)
|
| 522 |
+
# Fetch each interval once; record per-interval failures without aborting.
|
| 523 |
+
data, fetch_err = {}, {}
|
| 524 |
+
for iv in intervals:
|
| 525 |
+
try:
|
| 526 |
+
df = fetch(provider, symbol, iv)
|
| 527 |
+
if len(df) < 64:
|
| 528 |
+
raise ValueError("not enough history")
|
| 529 |
+
data[iv] = df
|
| 530 |
+
except Exception as e:
|
| 531 |
+
fetch_err[iv] = str(e)
|
| 532 |
+
|
| 533 |
+
for iv in intervals:
|
| 534 |
+
if iv not in data:
|
| 535 |
+
for model_name in models:
|
| 536 |
+
job["results"].append({"interval": iv, "model": model_name,
|
| 537 |
+
"error": fetch_err.get(iv, "unavailable")})
|
| 538 |
+
job["done_forecasts"] += 1
|
| 539 |
+
continue
|
| 540 |
+
df = data[iv]
|
| 541 |
+
last_close = float(df["close"].iloc[-1])
|
| 542 |
+
fut = future_timestamps(df["time"], max_h, INTERVALS[iv])
|
| 543 |
+
x_df = df[["open", "high", "low", "close", "volume", "amount"]]
|
| 544 |
+
x_ts = pd.Series(pd.to_datetime(df["time"], unit="s"))
|
| 545 |
+
y_ts = pd.Series(pd.to_datetime(fut, unit="s"))
|
| 546 |
+
|
| 547 |
+
for model_name in models:
|
| 548 |
+
job["current"] = f"{symbol} · {iv} · Kronos-{model_name}"
|
| 549 |
+
try:
|
| 550 |
+
predictor = get_predictor(model_name)
|
| 551 |
+
t0 = time.time()
|
| 552 |
+
with _predict_lock:
|
| 553 |
+
preds = predictor.predict_batch(
|
| 554 |
+
df_list=[x_df] * ANALYZER_PATHS,
|
| 555 |
+
x_timestamp_list=[x_ts] * ANALYZER_PATHS,
|
| 556 |
+
y_timestamp_list=[y_ts] * ANALYZER_PATHS,
|
| 557 |
+
pred_len=max_h, T=ANALYZER_T, top_p=ANALYZER_TOP_P,
|
| 558 |
+
sample_count=1, verbose=False,
|
| 559 |
+
)
|
| 560 |
+
elapsed = time.time() - t0
|
| 561 |
+
# (paths, max_h) close matrix; slice the prefix for each horizon
|
| 562 |
+
closes = np.stack([p["close"].to_numpy(dtype=np.float64) for p in preds])
|
| 563 |
+
for h in horizons:
|
| 564 |
+
sub = closes[:, :h]
|
| 565 |
+
mean_close = sub.mean(axis=0)
|
| 566 |
+
end = float(mean_close[-1])
|
| 567 |
+
lo, hi = np.percentile(sub[:, -1], [10, 90])
|
| 568 |
+
delta = (end / last_close - 1.0) * 100.0
|
| 569 |
+
job["results"].append({
|
| 570 |
+
"interval": iv, "horizon": h, "model": model_name,
|
| 571 |
+
"last_close": last_close, "end_close": end,
|
| 572 |
+
"delta_pct": delta,
|
| 573 |
+
"band_lo_pct": (lo / last_close - 1.0) * 100.0,
|
| 574 |
+
"band_hi_pct": (hi / last_close - 1.0) * 100.0,
|
| 575 |
+
"texture_pct": _texture_pct(sub),
|
| 576 |
+
"trend": "up" if delta > 0.05 else ("down" if delta < -0.05 else "flat"),
|
| 577 |
+
})
|
| 578 |
+
job["done_forecasts"] += 1
|
| 579 |
+
job["elapsed_total"] += elapsed
|
| 580 |
+
except Exception as e:
|
| 581 |
+
for h in horizons:
|
| 582 |
+
job["results"].append({"interval": iv, "horizon": h,
|
| 583 |
+
"model": model_name, "error": str(e)})
|
| 584 |
+
job["done_forecasts"] += 1
|
| 585 |
+
|
| 586 |
+
job["summary"] = _summarize(job["results"], intervals, horizons, models)
|
| 587 |
+
job["status"] = "done"
|
| 588 |
+
job["current"] = "complete"
|
| 589 |
+
except Exception as e:
|
| 590 |
+
job["status"] = "error"
|
| 591 |
+
job["error"] = str(e)
|
| 592 |
+
|
| 593 |
+
|
| 594 |
+
@app.route("/analyzer")
|
| 595 |
+
def analyzer_page():
|
| 596 |
+
return send_file(Path(__file__).resolve().parent / "analyzer.html")
|
| 597 |
+
|
| 598 |
+
|
| 599 |
+
@app.route("/api/analyze", methods=["POST"])
|
| 600 |
+
def api_analyze():
|
| 601 |
+
body = request.get_json(force=True)
|
| 602 |
+
provider = str(body.get("provider", "binance"))
|
| 603 |
+
symbol = str(body.get("symbol", "")).strip()
|
| 604 |
+
intervals = [i for i in body.get("intervals", []) if i in INTERVALS]
|
| 605 |
+
horizons = sorted({int(h) for h in body.get("horizons", []) if 1 <= int(h) <= MAX_HORIZON})
|
| 606 |
+
models = [m for m in body.get("models", []) if m in MODELS]
|
| 607 |
+
|
| 608 |
+
if not symbol or not intervals or not horizons or not models:
|
| 609 |
+
return jsonify({"error": "Need a symbol and at least one interval, horizon, and model"}), 400
|
| 610 |
+
try:
|
| 611 |
+
for iv in intervals:
|
| 612 |
+
validate(provider, symbol, iv)
|
| 613 |
+
except ValueError as e:
|
| 614 |
+
return jsonify({"error": str(e)}), 400
|
| 615 |
+
|
| 616 |
+
job_id = uuid.uuid4().hex[:12]
|
| 617 |
+
_jobs[job_id] = {
|
| 618 |
+
"id": job_id, "status": "running", "provider": provider, "symbol": symbol,
|
| 619 |
+
"intervals": intervals, "horizons": horizons, "models": models,
|
| 620 |
+
"total_forecasts": len(intervals) * len(models),
|
| 621 |
+
"done_forecasts": 0, "results": [], "summary": None, "error": None,
|
| 622 |
+
"current": "starting…", "elapsed_total": 0.0,
|
| 623 |
+
"started": datetime.datetime.now().isoformat(timespec="seconds"),
|
| 624 |
+
}
|
| 625 |
+
threading.Thread(target=_run_analysis,
|
| 626 |
+
args=(job_id, provider, symbol, intervals, horizons, models),
|
| 627 |
+
daemon=True).start()
|
| 628 |
+
return jsonify({"job_id": job_id, "total_forecasts": _jobs[job_id]["total_forecasts"]})
|
| 629 |
+
|
| 630 |
+
|
| 631 |
+
@app.route("/api/analyze/<job_id>")
|
| 632 |
+
def api_analyze_status(job_id):
|
| 633 |
+
job = _jobs.get(job_id)
|
| 634 |
+
if not job:
|
| 635 |
+
return jsonify({"error": "unknown job"}), 404
|
| 636 |
+
return jsonify({k: job[k] for k in (
|
| 637 |
+
"id", "status", "symbol", "provider", "intervals", "horizons", "models",
|
| 638 |
+
"total_forecasts", "done_forecasts", "results", "summary", "error",
|
| 639 |
+
"current", "elapsed_total", "started")})
|
| 640 |
+
|
| 641 |
+
|
| 642 |
+
@app.route("/analyses/<path:name>")
|
| 643 |
+
def serve_analysis(name):
|
| 644 |
+
"""Serve a previously saved analysis (HTML report or JSON)."""
|
| 645 |
+
target = (ANALYSES_DIR / name).resolve()
|
| 646 |
+
if ANALYSES_DIR.resolve() not in target.parents or not target.exists():
|
| 647 |
+
return jsonify({"error": "not found"}), 404
|
| 648 |
+
return send_file(target)
|
| 649 |
+
|
| 650 |
+
|
| 651 |
+
@app.route("/api/analyze/<job_id>/save", methods=["POST"])
|
| 652 |
+
def api_analyze_save(job_id):
|
| 653 |
+
job = _jobs.get(job_id)
|
| 654 |
+
if not job:
|
| 655 |
+
return jsonify({"error": "unknown job"}), 404
|
| 656 |
+
if job["status"] != "done":
|
| 657 |
+
return jsonify({"error": "analysis not finished"}), 400
|
| 658 |
+
ANALYSES_DIR.mkdir(exist_ok=True)
|
| 659 |
+
stamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
|
| 660 |
+
base = f"{job['symbol']}_{stamp}"
|
| 661 |
+
payload = {k: job[k] for k in (
|
| 662 |
+
"id", "symbol", "provider", "intervals", "horizons", "models",
|
| 663 |
+
"results", "summary", "started")}
|
| 664 |
+
payload["saved"] = datetime.datetime.now().isoformat(timespec="seconds")
|
| 665 |
+
(ANALYSES_DIR / f"{base}.json").write_text(json.dumps(payload, indent=2), encoding="utf-8")
|
| 666 |
+
(ANALYSES_DIR / f"{base}.html").write_text(_render_report(payload), encoding="utf-8")
|
| 667 |
+
return jsonify({"saved": f"{base}.json", "html": f"{base}.html",
|
| 668 |
+
"dir": str(ANALYSES_DIR)})
|
| 669 |
+
|
| 670 |
+
|
| 671 |
+
# --------------------------------------------------------------------------
|
| 672 |
+
# Leverage trade optimizer: turn a forecast cell into risk-tiered setups
|
| 673 |
+
# --------------------------------------------------------------------------
|
| 674 |
+
# Each tier sets how much margin you accept losing at the stop, the stop
|
| 675 |
+
# width (in multiples of the forecast's adverse edge), and a leverage cap.
|
| 676 |
+
# Leverage falls out as accept_loss / stop_distance, so wider stops or higher
|
| 677 |
+
# uncertainty automatically reduce leverage, and the stop is always inside
|
| 678 |
+
# liquidation by construction (accept_loss < 100%).
|
| 679 |
+
TIERS = [
|
| 680 |
+
{"name": "Conservative", "stop_mult": 1.6, "accept_loss": 8.0, "max_lev": 3.0, "entry": "limit"},
|
| 681 |
+
{"name": "Balanced", "stop_mult": 1.1, "accept_loss": 18.0, "max_lev": 10.0, "entry": "limit"},
|
| 682 |
+
{"name": "Aggressive", "stop_mult": 0.8, "accept_loss": 30.0, "max_lev": 20.0, "entry": "market"},
|
| 683 |
+
]
|
| 684 |
+
MIN_STOP_PCT = 0.30
|
| 685 |
+
MAINT_MARGIN = 0.005
|
| 686 |
+
FLAT_PCT = 0.10 # |expected move| below this = no directional edge
|
| 687 |
+
|
| 688 |
+
|
| 689 |
+
def _signal_for(job: dict, interval: str, horizon: int) -> dict:
|
| 690 |
+
"""Model-averaged forecast for one (interval, horizon) cell of a job."""
|
| 691 |
+
cells = [r for r in job["results"]
|
| 692 |
+
if r.get("interval") == interval and r.get("horizon") == horizon and "delta_pct" in r]
|
| 693 |
+
if not cells:
|
| 694 |
+
return None
|
| 695 |
+
avg = lambda k: float(np.mean([c[k] for c in cells]))
|
| 696 |
+
models = sorted({c["model"] for c in cells})
|
| 697 |
+
agree = None
|
| 698 |
+
if len(models) > 1:
|
| 699 |
+
signs = {_sign(c["delta_pct"]) for c in cells}
|
| 700 |
+
# conflict only when signs are strictly opposite (one up, one down);
|
| 701 |
+
# flat-vs-directional is weak agreement, not a conflict.
|
| 702 |
+
agree = not (1 in signs and -1 in signs)
|
| 703 |
+
return {
|
| 704 |
+
"interval": interval, "horizon": horizon,
|
| 705 |
+
"last_close": avg("last_close"), "mv": avg("delta_pct"),
|
| 706 |
+
"lo": avg("band_lo_pct"), "hi": avg("band_hi_pct"), "texture": avg("texture_pct"),
|
| 707 |
+
"models": models, "agree": agree,
|
| 708 |
+
"per_model": {c["model"]: c["delta_pct"] for c in cells},
|
| 709 |
+
}
|
| 710 |
+
|
| 711 |
+
|
| 712 |
+
def _optimize_trade(sig: dict) -> dict:
|
| 713 |
+
mv, lo, hi, last = sig["mv"], sig["lo"], sig["hi"], sig["last_close"]
|
| 714 |
+
|
| 715 |
+
# Two models pointing strictly opposite ways: the average masks a real
|
| 716 |
+
# conflict, so refuse to emit leveraged directional setups — that would be
|
| 717 |
+
# false confidence at exactly the wrong moment.
|
| 718 |
+
if sig["agree"] is False:
|
| 719 |
+
parts = ", ".join(f"{m} {d:+.2f}%" for m, d in sig["per_model"].items())
|
| 720 |
+
return {"direction": "conflict", "disagree": True, "tiers": [],
|
| 721 |
+
"headline": f"Models disagree on direction ({parts}) — no consensus. Leveraged setups "
|
| 722 |
+
"are suppressed; stand aside, pick a forecast where the models agree, or "
|
| 723 |
+
"run a single model."}
|
| 724 |
+
|
| 725 |
+
if abs(mv) < FLAT_PCT:
|
| 726 |
+
return {"direction": "neutral", "tiers": [],
|
| 727 |
+
"headline": f"Expected move {mv:+.2f}% is within noise (±{FLAT_PCT:.2f}%). "
|
| 728 |
+
"No directional edge — stand aside."}
|
| 729 |
+
|
| 730 |
+
long = mv > 0
|
| 731 |
+
edge = abs(mv) # expected move magnitude (%)
|
| 732 |
+
vol = max((hi - lo) / 2.0, 0.10) # half p10–p90 band ≈ 1.28σ, volatility proxy
|
| 733 |
+
adverse = max((-lo) if long else hi, 0.0) # distance to the unfavorable band edge
|
| 734 |
+
base_stop = max(adverse, vol)
|
| 735 |
+
favorable = max(hi if long else -lo, edge) # distance to the favorable band edge (p90/p10)
|
| 736 |
+
sq = edge / vol # edge-to-noise ("forecast Sharpe")
|
| 737 |
+
conf = "High" if sq >= 1.0 else "Medium" if sq >= 0.5 else "Low"
|
| 738 |
+
|
| 739 |
+
def price(pct): # +pct in the trade's favor
|
| 740 |
+
return last * (1 + pct / 100) if long else last * (1 - pct / 100)
|
| 741 |
+
|
| 742 |
+
def adverse_price(entry, pct):
|
| 743 |
+
return entry * (1 - pct / 100) if long else entry * (1 + pct / 100)
|
| 744 |
+
|
| 745 |
+
tiers = []
|
| 746 |
+
for t in TIERS:
|
| 747 |
+
stop_dist = max(t["stop_mult"] * base_stop, MIN_STOP_PCT)
|
| 748 |
+
lev = round(min(max(t["accept_loss"] / stop_dist, 1.0), t["max_lev"]), 1)
|
| 749 |
+
if t["entry"] == "limit":
|
| 750 |
+
pull = 0.25 * vol
|
| 751 |
+
entry = last * (1 - pull / 100) if long else last * (1 + pull / 100)
|
| 752 |
+
entry_kind = f"limit · wait for {pull:.2f}% pullback"
|
| 753 |
+
else:
|
| 754 |
+
entry, pull, entry_kind = last, 0.0, "market"
|
| 755 |
+
sl = adverse_price(entry, stop_dist)
|
| 756 |
+
margin_loss = lev * stop_dist
|
| 757 |
+
liq_dist = (100.0 / lev) * (1 - MAINT_MARGIN)
|
| 758 |
+
liq = adverse_price(entry, liq_dist)
|
| 759 |
+
|
| 760 |
+
if t["name"] == "Conservative":
|
| 761 |
+
tps_pct = [0.7 * edge, edge]
|
| 762 |
+
elif t["name"] == "Balanced":
|
| 763 |
+
tps_pct = [edge, favorable]
|
| 764 |
+
else:
|
| 765 |
+
tps_pct = [favorable, favorable * 1.4]
|
| 766 |
+
targets = []
|
| 767 |
+
for tp in tps_pct:
|
| 768 |
+
# target prices are forecast levels (anchored to last_close); R:R is
|
| 769 |
+
# measured from the actual entry, which may carry a small pullback.
|
| 770 |
+
tp_from_entry = (last * (1 + tp / 100) / entry - 1) * 100 if long else (1 - last * (1 - tp / 100) / entry) * 100
|
| 771 |
+
targets.append({
|
| 772 |
+
"pct": tp, "price": price(tp),
|
| 773 |
+
"rr": tp_from_entry / stop_dist,
|
| 774 |
+
"gain_margin_pct": lev * tp_from_entry,
|
| 775 |
+
})
|
| 776 |
+
|
| 777 |
+
reasoning = [
|
| 778 |
+
f"{t['name']} risk budget: accept ~{t['accept_loss']:.0f}% of margin lost at the stop, "
|
| 779 |
+
f"leverage capped at {t['max_lev']:g}× → sized to {lev:g}× here.",
|
| 780 |
+
f"Stop {stop_dist:.2f}% from entry — {t['stop_mult']:.1f}× the forecast's unfavorable edge "
|
| 781 |
+
f"(p{10 if long else 90} sits {adverse:.2f}% away). "
|
| 782 |
+
+ ("Wide, so ordinary noise won't trip it." if t["stop_mult"] >= 1.2
|
| 783 |
+
else "Tight, so it cuts losers fast but is easier to whipsaw."),
|
| 784 |
+
f"At {lev:g}× a stop-out costs ≈ {margin_loss:.0f}% of margin; liquidation is ~{liq_dist:.1f}% "
|
| 785 |
+
f"away — the stop sits comfortably inside it.",
|
| 786 |
+
f"Targets {tps_pct[0]:+.2f}% / {tps_pct[1]:+.2f}% → R:R {targets[0]['rr']:.2f} / {targets[1]['rr']:.2f}; "
|
| 787 |
+
f"at {lev:g}× that is +{targets[0]['gain_margin_pct']:.0f}% / +{targets[1]['gain_margin_pct']:.0f}% of margin "
|
| 788 |
+
+ ("(base case = model's expected close, stretch = p"
|
| 789 |
+
+ ("90" if long else "10") + " band edge)."
|
| 790 |
+
if t["name"] != "Aggressive" else "(aiming for the p"
|
| 791 |
+
+ ("90" if long else "10") + " edge and a 1.4× extension)."),
|
| 792 |
+
]
|
| 793 |
+
if targets[0]["rr"] < 1:
|
| 794 |
+
reasoning.append("⚠ R:R below 1 on the first target — the forecast move is small versus its "
|
| 795 |
+
"uncertainty, so this is a thin-edge scalp; leverage is doing the heavy lifting.")
|
| 796 |
+
tiers.append({
|
| 797 |
+
"name": t["name"], "direction": "long" if long else "short", "leverage": lev,
|
| 798 |
+
"entry": entry, "entry_kind": entry_kind,
|
| 799 |
+
"stop": sl, "stop_pct": stop_dist, "stop_margin_loss_pct": margin_loss,
|
| 800 |
+
"liq": liq, "liq_pct": liq_dist, "targets": targets, "reasoning": reasoning,
|
| 801 |
+
})
|
| 802 |
+
|
| 803 |
+
headline = (f"{'LONG' if long else 'SHORT'} bias — model-average expected move {mv:+.2f}% over "
|
| 804 |
+
f"{sig['horizon']}×{sig['interval']} bars, p10–p90 band {lo:+.2f}%…{hi:+.2f}% "
|
| 805 |
+
f"(edge/uncertainty {sq:.2f}).")
|
| 806 |
+
return {"direction": "long" if long else "short", "confidence": conf,
|
| 807 |
+
"signal_quality": sq, "disagree": False, "headline": headline, "tiers": tiers}
|
| 808 |
+
|
| 809 |
+
|
| 810 |
+
@app.route("/api/optimize", methods=["POST"])
|
| 811 |
+
def api_optimize():
|
| 812 |
+
body = request.get_json(force=True)
|
| 813 |
+
job = _jobs.get(body.get("job_id"))
|
| 814 |
+
if not job:
|
| 815 |
+
return jsonify({"error": "unknown job"}), 404
|
| 816 |
+
if job["status"] != "done":
|
| 817 |
+
return jsonify({"error": "analysis not finished"}), 400
|
| 818 |
+
interval = str(body.get("interval", ""))
|
| 819 |
+
try:
|
| 820 |
+
horizon = int(body.get("horizon"))
|
| 821 |
+
except (TypeError, ValueError):
|
| 822 |
+
return jsonify({"error": "bad horizon"}), 400
|
| 823 |
+
sig = _signal_for(job, interval, horizon)
|
| 824 |
+
if not sig:
|
| 825 |
+
return jsonify({"error": "no forecast for that interval/horizon"}), 404
|
| 826 |
+
return jsonify({"signal": sig, "optimizer": _optimize_trade(sig)})
|
| 827 |
+
|
| 828 |
+
|
| 829 |
+
def _render_report(p: dict) -> str:
|
| 830 |
+
"""Self-contained HTML snapshot of an analysis (embeds the data + a
|
| 831 |
+
standalone copy of the matrix renderer)."""
|
| 832 |
+
return (
|
| 833 |
+
"<!doctype html><html><head><meta charset='utf-8'>"
|
| 834 |
+
f"<title>Kronos analysis — {p['symbol']}</title>"
|
| 835 |
+
"<style>body{background:#0a0b10;color:#e8eaf2;font:14px/1.5 system-ui,sans-serif;"
|
| 836 |
+
"margin:0;padding:28px}h1{font-size:18px;font-weight:600}h2{font-size:13px;color:#767d96;"
|
| 837 |
+
"text-transform:uppercase;letter-spacing:.1em;margin-top:28px}table{border-collapse:collapse;"
|
| 838 |
+
"margin-top:10px}td,th{border:1px solid #1b1e2a;padding:8px 12px;text-align:center;"
|
| 839 |
+
"font-variant-numeric:tabular-nums}th{color:#767d96;font-weight:500}.cell{font-weight:600}"
|
| 840 |
+
".sub{font-size:11px;color:#9aa0b4;font-weight:400}.meta{color:#767d96;font-size:12px}</style>"
|
| 841 |
+
"</head><body>"
|
| 842 |
+
f"<h1>Kronos portfolio analysis — {p['symbol']}</h1>"
|
| 843 |
+
f"<div class='meta'>{p['provider']} · started {p['started']} · saved {p.get('saved','')}</div>"
|
| 844 |
+
"<div id='app'></div>"
|
| 845 |
+
f"<script>const DATA={json.dumps(p)};</script>"
|
| 846 |
+
"<script>" + _REPORT_JS + "</script>"
|
| 847 |
+
"</body></html>"
|
| 848 |
+
)
|
| 849 |
+
|
| 850 |
+
|
| 851 |
+
if __name__ == "__main__":
|
| 852 |
+
# Local default 127.0.0.1:8765; hosts like Hugging Face Spaces set HOST/PORT.
|
| 853 |
+
host = os.environ.get("HOST", "127.0.0.1")
|
| 854 |
+
port = int(os.environ.get("PORT", "8765"))
|
| 855 |
+
print(f"Kronos forecast UI -> http://{host}:{port}")
|
| 856 |
+
print(f"Portfolio analyzer -> http://{host}:{port}/analyzer")
|
| 857 |
+
app.run(host=host, port=port, debug=False, threaded=True)
|
crypto_ui/index.html
ADDED
|
@@ -0,0 +1,502 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
<!doctype html>
|
| 2 |
+
<html lang="en">
|
| 3 |
+
<head>
|
| 4 |
+
<meta charset="utf-8">
|
| 5 |
+
<meta name="viewport" content="width=device-width, initial-scale=1">
|
| 6 |
+
<title>Kronos — Market Forecast</title>
|
| 7 |
+
<link rel="preconnect" href="https://fonts.googleapis.com">
|
| 8 |
+
<link href="https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600&display=swap" rel="stylesheet">
|
| 9 |
+
<script src="https://unpkg.com/lightweight-charts@4.2.3/dist/lightweight-charts.standalone.production.js"></script>
|
| 10 |
+
<style>
|
| 11 |
+
:root {
|
| 12 |
+
--bg: #0a0b10; --panel: #10121a; --border: #1b1e2a; --text: #e8eaf2; --dim: #767d96;
|
| 13 |
+
--up: #16c784; --down: #ea3943; --accent: #8b5cf6;
|
| 14 |
+
}
|
| 15 |
+
* { box-sizing: border-box; }
|
| 16 |
+
[hidden] { display: none !important; }
|
| 17 |
+
body {
|
| 18 |
+
margin: 0; background: var(--bg); color: var(--text); height: 100vh;
|
| 19 |
+
display: flex; flex-direction: column;
|
| 20 |
+
font: 14px/1.45 'Inter', system-ui, -apple-system, sans-serif;
|
| 21 |
+
-webkit-font-smoothing: antialiased;
|
| 22 |
+
}
|
| 23 |
+
header {
|
| 24 |
+
display: flex; align-items: baseline; justify-content: space-between;
|
| 25 |
+
padding: 18px 28px 14px; border-bottom: 1px solid var(--border);
|
| 26 |
+
}
|
| 27 |
+
.brand { font-size: 13px; font-weight: 600; letter-spacing: .38em; }
|
| 28 |
+
.brand em { font-style: normal; color: var(--accent); }
|
| 29 |
+
.brand span { letter-spacing: .04em; font-weight: 400; color: var(--dim); margin-left: 14px; }
|
| 30 |
+
.model-tag { font-size: 11px; color: var(--dim); letter-spacing: .03em; }
|
| 31 |
+
|
| 32 |
+
.quote { display: flex; align-items: baseline; gap: 14px; padding: 18px 28px 2px; min-height: 46px; }
|
| 33 |
+
.pair { font-size: 13px; color: var(--dim); font-weight: 500; display: flex; align-items: baseline; gap: 10px; }
|
| 34 |
+
.pair b { color: var(--text); font-weight: 600; font-size: 14px; }
|
| 35 |
+
.aname { max-width: 260px; overflow: hidden; text-overflow: ellipsis; white-space: nowrap; }
|
| 36 |
+
.price { font-size: 32px; font-weight: 600; font-variant-numeric: tabular-nums; letter-spacing: -.01em; }
|
| 37 |
+
.chg { font-size: 13px; font-variant-numeric: tabular-nums; }
|
| 38 |
+
.up { color: var(--up); } .down { color: var(--down); }
|
| 39 |
+
.live {
|
| 40 |
+
width: 7px; height: 7px; border-radius: 50%; background: var(--up);
|
| 41 |
+
align-self: center; animation: pulse 2.4s ease-in-out infinite;
|
| 42 |
+
}
|
| 43 |
+
@keyframes pulse { 0%, 100% { opacity: .9; } 50% { opacity: .25; } }
|
| 44 |
+
|
| 45 |
+
.badge {
|
| 46 |
+
font-size: 9px; font-weight: 600; text-transform: uppercase; letter-spacing: .1em;
|
| 47 |
+
padding: 2px 8px; border-radius: 999px; border: 1px solid var(--dim); color: var(--dim);
|
| 48 |
+
align-self: center; white-space: nowrap;
|
| 49 |
+
}
|
| 50 |
+
.badge[data-k="Crypto"] { border-color: #8b5cf6; color: #a78bfa; }
|
| 51 |
+
.badge[data-k="Equity"] { border-color: #16c784; color: #34d399; }
|
| 52 |
+
.badge[data-k="ETF"] { border-color: #2dd4bf; color: #5eead4; }
|
| 53 |
+
.badge[data-k="Forex"] { border-color: #60a5fa; color: #93c5fd; }
|
| 54 |
+
.badge[data-k="Index"] { border-color: #f59e0b; color: #fbbf24; }
|
| 55 |
+
.badge[data-k="Commodity"] { border-color: #fb923c; color: #fdba74; }
|
| 56 |
+
|
| 57 |
+
.controls { display: flex; align-items: center; gap: 16px; flex-wrap: wrap; padding: 10px 28px 0; }
|
| 58 |
+
.controls:last-of-type { padding-bottom: 4px; }
|
| 59 |
+
.group { display: flex; align-items: center; gap: 6px; }
|
| 60 |
+
.group .lbl { font-size: 10px; text-transform: uppercase; letter-spacing: .12em; color: var(--dim); margin-right: 4px; }
|
| 61 |
+
.pill {
|
| 62 |
+
border: 1px solid var(--border); background: transparent; color: var(--dim);
|
| 63 |
+
border-radius: 999px; padding: 5px 13px; font: 500 12px 'Inter', sans-serif;
|
| 64 |
+
cursor: pointer; transition: all .15s;
|
| 65 |
+
}
|
| 66 |
+
.pill:hover { border-color: #2b3044; color: var(--text); }
|
| 67 |
+
.pill.on { background: rgba(139, 92, 246, .13); border-color: rgba(139, 92, 246, .55); color: var(--text); }
|
| 68 |
+
|
| 69 |
+
.search-wrap { position: relative; flex: 1 1 280px; max-width: 460px; }
|
| 70 |
+
#q {
|
| 71 |
+
width: 100%; background: var(--panel); border: 1px solid var(--border); color: var(--text);
|
| 72 |
+
border-radius: 999px; padding: 8px 18px; font: 400 13px 'Inter', sans-serif; outline: none;
|
| 73 |
+
transition: border-color .15s;
|
| 74 |
+
}
|
| 75 |
+
#q:focus { border-color: rgba(139, 92, 246, .55); }
|
| 76 |
+
#q::placeholder { color: #4d5469; }
|
| 77 |
+
.dropdown {
|
| 78 |
+
position: absolute; top: calc(100% + 6px); left: 0; right: 0; z-index: 20;
|
| 79 |
+
background: var(--panel); border: 1px solid var(--border); border-radius: 12px;
|
| 80 |
+
max-height: 322px; overflow-y: auto; box-shadow: 0 14px 40px rgba(0, 0, 0, .5);
|
| 81 |
+
}
|
| 82 |
+
.opt {
|
| 83 |
+
display: flex; align-items: center; gap: 10px; padding: 9px 14px; cursor: pointer;
|
| 84 |
+
font-size: 12px;
|
| 85 |
+
}
|
| 86 |
+
.opt:hover, .opt.sel { background: rgba(139, 92, 246, .10); }
|
| 87 |
+
.opt .sym { font-weight: 600; min-width: 86px; }
|
| 88 |
+
.opt .nm { color: var(--dim); flex: 1; overflow: hidden; text-overflow: ellipsis; white-space: nowrap; }
|
| 89 |
+
.opt .ex { color: #4d5469; font-size: 10px; }
|
| 90 |
+
.none { padding: 12px 14px; color: var(--dim); font-size: 12px; }
|
| 91 |
+
|
| 92 |
+
.staged {
|
| 93 |
+
display: flex; align-items: center; gap: 9px; padding: 5px 7px 5px 13px;
|
| 94 |
+
border: 1px dashed rgba(139, 92, 246, .5); border-radius: 999px; font-size: 12px;
|
| 95 |
+
}
|
| 96 |
+
.staged .x {
|
| 97 |
+
border: none; background: transparent; color: var(--dim); cursor: pointer;
|
| 98 |
+
font-size: 14px; padding: 0 5px; line-height: 1;
|
| 99 |
+
}
|
| 100 |
+
.staged .x:hover { color: var(--text); }
|
| 101 |
+
#load {
|
| 102 |
+
border: 1px solid var(--accent); background: transparent; color: #a78bfa;
|
| 103 |
+
border-radius: 999px; padding: 7px 22px; font: 600 13px 'Inter', sans-serif;
|
| 104 |
+
cursor: pointer; transition: all .15s;
|
| 105 |
+
}
|
| 106 |
+
#load:hover:not(:disabled) { background: rgba(139, 92, 246, .15); }
|
| 107 |
+
#load:disabled { opacity: .35; cursor: default; }
|
| 108 |
+
|
| 109 |
+
#go {
|
| 110 |
+
margin-left: auto; border: none; border-radius: 999px; padding: 8px 26px;
|
| 111 |
+
background: var(--accent); color: #fff; font: 600 13px 'Inter', sans-serif;
|
| 112 |
+
cursor: pointer; transition: filter .15s; min-width: 118px;
|
| 113 |
+
}
|
| 114 |
+
#go:hover:not(:disabled) { filter: brightness(1.12); }
|
| 115 |
+
#go:disabled { opacity: .5; cursor: default; }
|
| 116 |
+
.busy .controls .pill, .busy #load { pointer-events: none; opacity: .6; }
|
| 117 |
+
|
| 118 |
+
.chart-wrap { flex: 1; position: relative; margin: 6px 20px 0; min-height: 300px; }
|
| 119 |
+
#chart { position: absolute; inset: 0; }
|
| 120 |
+
#empty {
|
| 121 |
+
position: absolute; inset: 0; z-index: 5; display: flex; flex-direction: column;
|
| 122 |
+
align-items: center; justify-content: center; gap: 8px; color: var(--dim);
|
| 123 |
+
background: var(--bg); font-size: 13px;
|
| 124 |
+
}
|
| 125 |
+
#empty .big { font-size: 15px; color: var(--text); }
|
| 126 |
+
|
| 127 |
+
.stats { display: flex; flex-wrap: wrap; gap: 8px; padding: 10px 28px 4px; min-height: 38px; }
|
| 128 |
+
.chip {
|
| 129 |
+
border: 1px solid var(--border); border-radius: 8px; padding: 5px 12px;
|
| 130 |
+
font-size: 12px; color: var(--dim); font-variant-numeric: tabular-nums;
|
| 131 |
+
white-space: nowrap;
|
| 132 |
+
}
|
| 133 |
+
.chip b { font-weight: 600; color: var(--text); }
|
| 134 |
+
|
| 135 |
+
footer {
|
| 136 |
+
display: flex; justify-content: space-between; padding: 10px 28px 14px;
|
| 137 |
+
font-size: 11px; color: var(--dim);
|
| 138 |
+
}
|
| 139 |
+
#status.err { color: var(--down); }
|
| 140 |
+
</style>
|
| 141 |
+
</head>
|
| 142 |
+
<body>
|
| 143 |
+
<header>
|
| 144 |
+
<div class="brand">KRONOS<em>.</em><span>market forecast</span></div>
|
| 145 |
+
<div style="display:flex;align-items:baseline;gap:20px">
|
| 146 |
+
<a href="/analyzer" style="color:#a78bfa;text-decoration:none;font-size:12px;letter-spacing:.02em">Portfolio Analyzer →</a>
|
| 147 |
+
<div class="model-tag" id="modelTag">Kronos-small · 24.7M params · CPU</div>
|
| 148 |
+
</div>
|
| 149 |
+
</header>
|
| 150 |
+
|
| 151 |
+
<section class="quote">
|
| 152 |
+
<div class="pair"><b id="assetSym">—</b><span class="aname" id="assetName"></span></div>
|
| 153 |
+
<span class="badge" id="assetClass" hidden></span>
|
| 154 |
+
<div class="price" id="price"></div>
|
| 155 |
+
<div class="chg" id="chg"></div>
|
| 156 |
+
<div class="live" id="live" title="live · auto-refreshes every 30s" hidden></div>
|
| 157 |
+
</section>
|
| 158 |
+
|
| 159 |
+
<section class="controls">
|
| 160 |
+
<div class="search-wrap">
|
| 161 |
+
<input id="q" placeholder="Search any asset — BTC, AAPL, gold, EUR/USD, S&P 500…"
|
| 162 |
+
autocomplete="off" spellcheck="false">
|
| 163 |
+
<div class="dropdown" id="dd" hidden></div>
|
| 164 |
+
</div>
|
| 165 |
+
<div class="staged" id="staged" hidden>
|
| 166 |
+
<span class="badge" id="stagedClass"></span>
|
| 167 |
+
<span id="stagedLabel"></span>
|
| 168 |
+
<button class="x" id="unstage" title="clear selection">×</button>
|
| 169 |
+
</div>
|
| 170 |
+
<button id="load" disabled>Load</button>
|
| 171 |
+
</section>
|
| 172 |
+
|
| 173 |
+
<section class="controls">
|
| 174 |
+
<div class="group" id="intervals"><span class="lbl">Interval</span></div>
|
| 175 |
+
<div class="group" id="horizons"><span class="lbl">Horizon</span></div>
|
| 176 |
+
<div class="group" id="models"><span class="lbl">Model</span></div>
|
| 177 |
+
<button id="go" disabled>Forecast</button>
|
| 178 |
+
</section>
|
| 179 |
+
|
| 180 |
+
<div class="chart-wrap">
|
| 181 |
+
<div id="chart"></div>
|
| 182 |
+
<div id="empty">
|
| 183 |
+
<div class="big">No asset loaded</div>
|
| 184 |
+
<div>search above, pick a result, then press <b>Load</b></div>
|
| 185 |
+
</div>
|
| 186 |
+
</div>
|
| 187 |
+
<section class="stats" id="stats"></section>
|
| 188 |
+
|
| 189 |
+
<footer>
|
| 190 |
+
<span id="status">search an asset to begin</span>
|
| 191 |
+
<span>research demo · not financial advice</span>
|
| 192 |
+
</footer>
|
| 193 |
+
|
| 194 |
+
<script>
|
| 195 |
+
const INTERVALS = ['15m', '1h', '4h', '1d'];
|
| 196 |
+
const HORIZONS = [12, 24, 48];
|
| 197 |
+
const MODEL_PARAMS = { small: '24.7M', base: '102.3M' };
|
| 198 |
+
const REFRESH_MS = 30000;
|
| 199 |
+
const state = {
|
| 200 |
+
asset: null, // {provider, symbol, name, klass, exchange} — loaded & charted
|
| 201 |
+
staged: null, // same shape — picked from search, awaiting Load confirmation
|
| 202 |
+
interval: '1h', horizon: 24, model: 'small', busy: false, candles: [],
|
| 203 |
+
};
|
| 204 |
+
const $ = (s) => document.querySelector(s);
|
| 205 |
+
const assetKey = () => state.asset ? `${state.asset.provider}:${state.asset.symbol}:${state.interval}` : '';
|
| 206 |
+
|
| 207 |
+
/* ---------- chart ---------- */
|
| 208 |
+
const chart = LightweightCharts.createChart($('#chart'), {
|
| 209 |
+
autoSize: true,
|
| 210 |
+
layout: {
|
| 211 |
+
background: { type: 'solid', color: 'transparent' },
|
| 212 |
+
textColor: '#767d96', fontSize: 11,
|
| 213 |
+
fontFamily: "'Inter', system-ui, sans-serif",
|
| 214 |
+
},
|
| 215 |
+
grid: {
|
| 216 |
+
vertLines: { color: 'rgba(255,255,255,.035)' },
|
| 217 |
+
horzLines: { color: 'rgba(255,255,255,.035)' },
|
| 218 |
+
},
|
| 219 |
+
rightPriceScale: { borderVisible: false },
|
| 220 |
+
timeScale: { borderVisible: false, timeVisible: true, secondsVisible: false, rightOffset: 5 },
|
| 221 |
+
crosshair: {
|
| 222 |
+
mode: LightweightCharts.CrosshairMode.Normal,
|
| 223 |
+
vertLine: { color: 'rgba(139,92,246,.35)', labelBackgroundColor: '#8b5cf6' },
|
| 224 |
+
horzLine: { color: 'rgba(139,92,246,.35)', labelBackgroundColor: '#8b5cf6' },
|
| 225 |
+
},
|
| 226 |
+
});
|
| 227 |
+
const candles = chart.addCandlestickSeries({
|
| 228 |
+
upColor: '#16c784', downColor: '#ea3943',
|
| 229 |
+
wickUpColor: '#16c784', wickDownColor: '#ea3943', borderVisible: false,
|
| 230 |
+
});
|
| 231 |
+
const volume = chart.addHistogramSeries({ priceScaleId: 'vol', priceFormat: { type: 'volume' }, lastValueVisible: false, priceLineVisible: false });
|
| 232 |
+
chart.priceScale('vol').applyOptions({ scaleMargins: { top: .85, bottom: 0 }, visible: false });
|
| 233 |
+
const ghost = chart.addCandlestickSeries({
|
| 234 |
+
upColor: 'rgba(139,92,246,.50)', downColor: 'rgba(139,92,246,.18)',
|
| 235 |
+
wickUpColor: 'rgba(139,92,246,.55)', wickDownColor: 'rgba(139,92,246,.55)',
|
| 236 |
+
borderVisible: false, lastValueVisible: false, priceLineVisible: false,
|
| 237 |
+
});
|
| 238 |
+
const bandOpts = {
|
| 239 |
+
color: 'rgba(139,92,246,.45)', lineWidth: 1, lineStyle: LightweightCharts.LineStyle.Dashed,
|
| 240 |
+
lastValueVisible: false, priceLineVisible: false, crosshairMarkerVisible: false,
|
| 241 |
+
};
|
| 242 |
+
const bandHi = chart.addLineSeries(bandOpts);
|
| 243 |
+
const bandLo = chart.addLineSeries(bandOpts);
|
| 244 |
+
|
| 245 |
+
/* ---------- helpers ---------- */
|
| 246 |
+
const fmt = (p) => p >= 1000 ? p.toLocaleString('en-US', { minimumFractionDigits: 2, maximumFractionDigits: 2 })
|
| 247 |
+
: p >= 1 ? p.toLocaleString('en-US', { minimumFractionDigits: 2, maximumFractionDigits: 4 })
|
| 248 |
+
: Number(p.toPrecision(4)).toString();
|
| 249 |
+
const signed = (v) => `${v >= 0 ? '+' : ''}${v.toFixed(2)}%`;
|
| 250 |
+
const esc = (s) => String(s).replace(/[&<>"']/g, c => ({ '&': '&', '<': '<', '>': '>', '"': '"', "'": ''' }[c]));
|
| 251 |
+
|
| 252 |
+
function setStatus(msg, err = false) {
|
| 253 |
+
const el = $('#status');
|
| 254 |
+
el.textContent = msg;
|
| 255 |
+
el.className = err ? 'err' : '';
|
| 256 |
+
}
|
| 257 |
+
|
| 258 |
+
const volBar = (c) => ({
|
| 259 |
+
time: c.time, value: c.volume,
|
| 260 |
+
color: c.close >= c.open ? 'rgba(22,199,132,.16)' : 'rgba(234,57,67,.16)',
|
| 261 |
+
});
|
| 262 |
+
|
| 263 |
+
function updateQuote(rows) {
|
| 264 |
+
const last = rows[rows.length - 1].close;
|
| 265 |
+
$('#price').textContent = fmt(last);
|
| 266 |
+
const dayAgo = rows[rows.length - 1].time - 86400;
|
| 267 |
+
const ref = [...rows].reverse().find(c => c.time <= dayAgo);
|
| 268 |
+
if (ref) {
|
| 269 |
+
const d = (last / ref.close - 1) * 100;
|
| 270 |
+
$('#chg').textContent = `${signed(d)} · 24h`;
|
| 271 |
+
$('#chg').className = `chg ${d >= 0 ? 'up' : 'down'}`;
|
| 272 |
+
} else { $('#chg').textContent = ''; }
|
| 273 |
+
}
|
| 274 |
+
|
| 275 |
+
function setHistory(rows) {
|
| 276 |
+
state.candles = rows;
|
| 277 |
+
candles.setData(rows);
|
| 278 |
+
volume.setData(rows.map(volBar));
|
| 279 |
+
const last = rows[rows.length - 1].close;
|
| 280 |
+
const precision = last >= 100 ? 2 : last >= 1 ? 4 : 6;
|
| 281 |
+
candles.applyOptions({ priceFormat: { type: 'price', precision, minMove: 1 / 10 ** precision } });
|
| 282 |
+
ghost.applyOptions({ priceFormat: { type: 'price', precision, minMove: 1 / 10 ** precision } });
|
| 283 |
+
updateQuote(rows);
|
| 284 |
+
return rows.length;
|
| 285 |
+
}
|
| 286 |
+
|
| 287 |
+
function clearForecast() {
|
| 288 |
+
ghost.setData([]); bandHi.setData([]); bandLo.setData([]);
|
| 289 |
+
$('#stats').innerHTML = '';
|
| 290 |
+
}
|
| 291 |
+
|
| 292 |
+
function setRange(histLen, fcLen) {
|
| 293 |
+
chart.timeScale().setVisibleLogicalRange({ from: histLen - 110, to: histLen + (fcLen || state.horizon) + 4 });
|
| 294 |
+
}
|
| 295 |
+
|
| 296 |
+
function setBusy(b) {
|
| 297 |
+
state.busy = b;
|
| 298 |
+
document.body.classList.toggle('busy', b);
|
| 299 |
+
$('#go').disabled = b || !state.asset;
|
| 300 |
+
$('#go').textContent = b ? 'Forecasting…' : 'Forecast';
|
| 301 |
+
}
|
| 302 |
+
|
| 303 |
+
/* ---------- search & stage (nothing loads until the user confirms) ---------- */
|
| 304 |
+
const qEl = $('#q'), dd = $('#dd');
|
| 305 |
+
let ddItems = [], ddSel = -1, debounce = null, inflight = null;
|
| 306 |
+
|
| 307 |
+
qEl.addEventListener('input', () => {
|
| 308 |
+
clearTimeout(debounce);
|
| 309 |
+
debounce = setTimeout(doSearch, 280);
|
| 310 |
+
});
|
| 311 |
+
qEl.addEventListener('keydown', (e) => {
|
| 312 |
+
if (dd.hidden) return;
|
| 313 |
+
if (e.key === 'ArrowDown') { e.preventDefault(); moveSel(1); }
|
| 314 |
+
else if (e.key === 'ArrowUp') { e.preventDefault(); moveSel(-1); }
|
| 315 |
+
else if (e.key === 'Enter') { e.preventDefault(); if (ddItems.length) stage(ddItems[Math.max(ddSel, 0)]); }
|
| 316 |
+
else if (e.key === 'Escape') hideDD();
|
| 317 |
+
});
|
| 318 |
+
qEl.addEventListener('blur', () => setTimeout(hideDD, 150));
|
| 319 |
+
|
| 320 |
+
async function doSearch() {
|
| 321 |
+
const term = qEl.value.trim();
|
| 322 |
+
if (!term) { hideDD(); return; }
|
| 323 |
+
inflight?.abort();
|
| 324 |
+
inflight = new AbortController();
|
| 325 |
+
try {
|
| 326 |
+
const r = await fetch(`/api/search?q=${encodeURIComponent(term)}`, { signal: inflight.signal });
|
| 327 |
+
const j = await r.json();
|
| 328 |
+
if (term !== qEl.value.trim()) return;
|
| 329 |
+
ddItems = j.results || []; ddSel = -1;
|
| 330 |
+
renderDD();
|
| 331 |
+
} catch (e) { if (e.name !== 'AbortError') hideDD(); }
|
| 332 |
+
}
|
| 333 |
+
|
| 334 |
+
function renderDD() {
|
| 335 |
+
if (!ddItems.length) {
|
| 336 |
+
dd.innerHTML = '<div class="none">no matches</div>';
|
| 337 |
+
dd.hidden = false;
|
| 338 |
+
return;
|
| 339 |
+
}
|
| 340 |
+
dd.innerHTML = ddItems.map((it, i) => `
|
| 341 |
+
<div class="opt${i === ddSel ? ' sel' : ''}" data-i="${i}">
|
| 342 |
+
<span class="badge" data-k="${esc(it.klass)}">${esc(it.klass)}</span>
|
| 343 |
+
<span class="sym">${esc(it.symbol)}</span>
|
| 344 |
+
<span class="nm">${esc(it.name)}</span>
|
| 345 |
+
<span class="ex">${esc(it.exchange)}</span>
|
| 346 |
+
</div>`).join('');
|
| 347 |
+
dd.hidden = false;
|
| 348 |
+
dd.querySelectorAll('.opt').forEach(el =>
|
| 349 |
+
el.addEventListener('mousedown', () => stage(ddItems[+el.dataset.i])));
|
| 350 |
+
}
|
| 351 |
+
|
| 352 |
+
function moveSel(d) {
|
| 353 |
+
ddSel = (ddSel + d + ddItems.length) % ddItems.length;
|
| 354 |
+
renderDD();
|
| 355 |
+
dd.querySelector('.opt.sel')?.scrollIntoView({ block: 'nearest' });
|
| 356 |
+
}
|
| 357 |
+
|
| 358 |
+
function hideDD() { dd.hidden = true; }
|
| 359 |
+
|
| 360 |
+
function stage(it) {
|
| 361 |
+
state.staged = it;
|
| 362 |
+
hideDD();
|
| 363 |
+
qEl.value = '';
|
| 364 |
+
$('#stagedClass').textContent = it.klass;
|
| 365 |
+
$('#stagedClass').dataset.k = it.klass;
|
| 366 |
+
$('#stagedLabel').textContent = `${it.symbol} · ${it.name}`;
|
| 367 |
+
$('#staged').hidden = false;
|
| 368 |
+
$('#load').disabled = false;
|
| 369 |
+
setStatus(`${it.symbol} staged — press Load to fetch data`);
|
| 370 |
+
}
|
| 371 |
+
|
| 372 |
+
$('#unstage').onclick = () => {
|
| 373 |
+
state.staged = null;
|
| 374 |
+
$('#staged').hidden = true;
|
| 375 |
+
$('#load').disabled = true;
|
| 376 |
+
setStatus(state.asset ? `${state.asset.symbol} · ${state.interval} · ready` : 'search an asset to begin');
|
| 377 |
+
};
|
| 378 |
+
|
| 379 |
+
/* ---------- load (the explicit confirmation) ---------- */
|
| 380 |
+
async function loadAsset(it) {
|
| 381 |
+
clearForecast();
|
| 382 |
+
setStatus(`loading ${it.symbol}…`);
|
| 383 |
+
try {
|
| 384 |
+
const r = await fetch(`/api/klines?provider=${it.provider}&symbol=${encodeURIComponent(it.symbol)}&interval=${state.interval}`);
|
| 385 |
+
const j = await r.json();
|
| 386 |
+
if (!r.ok) throw new Error(j.error || r.statusText);
|
| 387 |
+
state.asset = it;
|
| 388 |
+
const n = setHistory(j.candles);
|
| 389 |
+
setRange(n, 0);
|
| 390 |
+
$('#assetSym').textContent = it.symbol;
|
| 391 |
+
$('#assetName').textContent = it.name;
|
| 392 |
+
const b = $('#assetClass');
|
| 393 |
+
b.textContent = it.klass; b.dataset.k = it.klass; b.hidden = false;
|
| 394 |
+
$('#live').hidden = false;
|
| 395 |
+
$('#empty').hidden = true;
|
| 396 |
+
$('#go').disabled = false;
|
| 397 |
+
setStatus(`${it.symbol} · ${state.interval} · ${n} closed bars · ready`);
|
| 398 |
+
return true;
|
| 399 |
+
} catch (e) { setStatus(e.message, true); return false; }
|
| 400 |
+
}
|
| 401 |
+
|
| 402 |
+
$('#load').onclick = async () => {
|
| 403 |
+
if (!state.staged || state.busy) return;
|
| 404 |
+
const it = state.staged;
|
| 405 |
+
$('#load').disabled = true;
|
| 406 |
+
if (await loadAsset(it)) {
|
| 407 |
+
state.staged = null;
|
| 408 |
+
$('#staged').hidden = true;
|
| 409 |
+
} else {
|
| 410 |
+
$('#load').disabled = false;
|
| 411 |
+
}
|
| 412 |
+
};
|
| 413 |
+
|
| 414 |
+
/* ---------- forecast ---------- */
|
| 415 |
+
let baseRan = false;
|
| 416 |
+
async function runForecast() {
|
| 417 |
+
if (state.busy || !state.asset) return;
|
| 418 |
+
setBusy(true);
|
| 419 |
+
clearForecast();
|
| 420 |
+
const note = state.model === 'base' && !baseRan ? ' · first run may download ~400 MB' : '';
|
| 421 |
+
const t0 = performance.now();
|
| 422 |
+
const tick = setInterval(() =>
|
| 423 |
+
setStatus(`sampling 5 paths · ${state.model}${note} · ${((performance.now() - t0) / 1000).toFixed(1)}s`), 100);
|
| 424 |
+
try {
|
| 425 |
+
const r = await fetch('/api/predict', {
|
| 426 |
+
method: 'POST', headers: { 'Content-Type': 'application/json' },
|
| 427 |
+
body: JSON.stringify({
|
| 428 |
+
provider: state.asset.provider, symbol: state.asset.symbol,
|
| 429 |
+
interval: state.interval, horizon: state.horizon, model: state.model,
|
| 430 |
+
}),
|
| 431 |
+
});
|
| 432 |
+
const j = await r.json();
|
| 433 |
+
if (!r.ok) throw new Error(j.error || r.statusText);
|
| 434 |
+
const n = setHistory(j.context);
|
| 435 |
+
ghost.setData(j.forecast.candles);
|
| 436 |
+
bandHi.setData(j.forecast.p90);
|
| 437 |
+
bandLo.setData(j.forecast.p10);
|
| 438 |
+
setRange(n, j.forecast.candles.length);
|
| 439 |
+
const s = j.stats;
|
| 440 |
+
if (s.model === 'base') baseRan = true;
|
| 441 |
+
const cls = s.delta_pct >= 0 ? 'up' : 'down';
|
| 442 |
+
$('#stats').innerHTML = `
|
| 443 |
+
<div class="chip">expected <b class="${cls}">${signed(s.delta_pct)}</b> in ${state.horizon} bars</div>
|
| 444 |
+
<div class="chip">p10–p90 <b>${signed(s.band_lo_pct)} … ${signed(s.band_hi_pct)}</b></div>
|
| 445 |
+
<div class="chip">end close <b>${fmt(s.end_close)}</b></div>
|
| 446 |
+
<div class="chip">${s.model} · ${s.paths} paths · ${s.elapsed_s}s</div>`;
|
| 447 |
+
setStatus(`forecast complete · ${s.elapsed_s}s`);
|
| 448 |
+
} catch (e) { setStatus(e.message, true); }
|
| 449 |
+
finally { clearInterval(tick); setBusy(false); }
|
| 450 |
+
}
|
| 451 |
+
|
| 452 |
+
/* ---------- silent auto-refresh of the loaded asset ---------- */
|
| 453 |
+
async function refresh() {
|
| 454 |
+
if (state.busy || document.hidden || !state.asset || !state.candles.length) return;
|
| 455 |
+
const k = assetKey();
|
| 456 |
+
try {
|
| 457 |
+
const a = state.asset;
|
| 458 |
+
const r = await fetch(`/api/klines?provider=${a.provider}&symbol=${encodeURIComponent(a.symbol)}&interval=${state.interval}`);
|
| 459 |
+
const j = await r.json();
|
| 460 |
+
if (!r.ok || k !== assetKey() || state.busy) return;
|
| 461 |
+
const rows = j.candles;
|
| 462 |
+
const lastT = state.candles[state.candles.length - 1].time;
|
| 463 |
+
rows.filter(c => c.time >= lastT).forEach(c => {
|
| 464 |
+
candles.update(c);
|
| 465 |
+
volume.update(volBar(c));
|
| 466 |
+
});
|
| 467 |
+
state.candles = rows;
|
| 468 |
+
updateQuote(rows);
|
| 469 |
+
} catch { /* transient network errors: try again next tick */ }
|
| 470 |
+
}
|
| 471 |
+
setInterval(refresh, REFRESH_MS);
|
| 472 |
+
|
| 473 |
+
/* ---------- pills ---------- */
|
| 474 |
+
function pills(elId, items, current, fmtFn, onPick) {
|
| 475 |
+
const wrap = $(elId);
|
| 476 |
+
items.forEach(v => {
|
| 477 |
+
const b = document.createElement('button');
|
| 478 |
+
b.className = `pill${v === current ? ' on' : ''}`;
|
| 479 |
+
b.textContent = fmtFn(v);
|
| 480 |
+
b.onclick = () => {
|
| 481 |
+
wrap.querySelectorAll('.pill').forEach(p => p.classList.remove('on'));
|
| 482 |
+
b.classList.add('on');
|
| 483 |
+
onPick(v);
|
| 484 |
+
};
|
| 485 |
+
wrap.appendChild(b);
|
| 486 |
+
});
|
| 487 |
+
}
|
| 488 |
+
pills('#intervals', INTERVALS, state.interval, v => v.toUpperCase(), v => {
|
| 489 |
+
state.interval = v;
|
| 490 |
+
if (state.asset) loadAsset(state.asset);
|
| 491 |
+
});
|
| 492 |
+
pills('#horizons', HORIZONS, state.horizon, v => `${v} bars`, v => { state.horizon = v; });
|
| 493 |
+
pills('#models', Object.keys(MODEL_PARAMS), state.model, v => v[0].toUpperCase() + v.slice(1), v => {
|
| 494 |
+
state.model = v;
|
| 495 |
+
$('#modelTag').textContent = `Kronos-${v} · ${MODEL_PARAMS[v]} params · CPU`;
|
| 496 |
+
});
|
| 497 |
+
$('#go').onclick = runForecast;
|
| 498 |
+
|
| 499 |
+
qEl.focus();
|
| 500 |
+
</script>
|
| 501 |
+
</body>
|
| 502 |
+
</html>
|
crypto_ui/requirements.txt
ADDED
|
@@ -0,0 +1,9 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Runtime dependencies for the Kronos forecast web app (CPU).
|
| 2 |
+
flask==3.1.3
|
| 3 |
+
torch>=2.0.0
|
| 4 |
+
numpy
|
| 5 |
+
pandas==2.2.3
|
| 6 |
+
einops==0.8.1
|
| 7 |
+
huggingface_hub==0.33.1
|
| 8 |
+
safetensors==0.6.2
|
| 9 |
+
requests
|
model/LICENSE
ADDED
|
@@ -0,0 +1,21 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
MIT License
|
| 2 |
+
|
| 3 |
+
Copyright (c) 2025 ShiYu
|
| 4 |
+
|
| 5 |
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
| 6 |
+
of this software and associated documentation files (the "Software"), to deal
|
| 7 |
+
in the Software without restriction, including without limitation the rights
|
| 8 |
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
| 9 |
+
copies of the Software, and to permit persons to whom the Software is
|
| 10 |
+
furnished to do so, subject to the following conditions:
|
| 11 |
+
|
| 12 |
+
The above copyright notice and this permission notice shall be included in all
|
| 13 |
+
copies or substantial portions of the Software.
|
| 14 |
+
|
| 15 |
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
| 16 |
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
| 17 |
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
| 18 |
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
| 19 |
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
| 20 |
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
| 21 |
+
SOFTWARE.
|
model/__init__.py
ADDED
|
@@ -0,0 +1,17 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from .kronos import KronosTokenizer, Kronos, KronosPredictor
|
| 2 |
+
|
| 3 |
+
model_dict = {
|
| 4 |
+
'kronos_tokenizer': KronosTokenizer,
|
| 5 |
+
'kronos': Kronos,
|
| 6 |
+
'kronos_predictor': KronosPredictor
|
| 7 |
+
}
|
| 8 |
+
|
| 9 |
+
|
| 10 |
+
def get_model_class(model_name):
|
| 11 |
+
if model_name in model_dict:
|
| 12 |
+
return model_dict[model_name]
|
| 13 |
+
else:
|
| 14 |
+
print(f"Model {model_name} not found in model_dict")
|
| 15 |
+
raise NotImplementedError
|
| 16 |
+
|
| 17 |
+
|
model/kronos.py
ADDED
|
@@ -0,0 +1,662 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import numpy as np
|
| 2 |
+
import pandas as pd
|
| 3 |
+
import torch
|
| 4 |
+
from huggingface_hub import PyTorchModelHubMixin
|
| 5 |
+
import sys
|
| 6 |
+
|
| 7 |
+
from tqdm import trange
|
| 8 |
+
|
| 9 |
+
sys.path.append("../")
|
| 10 |
+
from model.module import *
|
| 11 |
+
|
| 12 |
+
|
| 13 |
+
class KronosTokenizer(nn.Module, PyTorchModelHubMixin):
|
| 14 |
+
"""
|
| 15 |
+
KronosTokenizer module for tokenizing input data using a hybrid quantization approach.
|
| 16 |
+
|
| 17 |
+
This tokenizer utilizes a combination of encoder and decoder Transformer blocks
|
| 18 |
+
along with the Binary Spherical Quantization (BSQuantizer) to compress and decompress input data.
|
| 19 |
+
|
| 20 |
+
Args:
|
| 21 |
+
d_in (int): Input dimension.
|
| 22 |
+
d_model (int): Model dimension.
|
| 23 |
+
n_heads (int): Number of attention heads.
|
| 24 |
+
ff_dim (int): Feed-forward dimension.
|
| 25 |
+
n_enc_layers (int): Number of encoder layers.
|
| 26 |
+
n_dec_layers (int): Number of decoder layers.
|
| 27 |
+
ffn_dropout_p (float): Dropout probability for feed-forward networks.
|
| 28 |
+
attn_dropout_p (float): Dropout probability for attention mechanisms.
|
| 29 |
+
resid_dropout_p (float): Dropout probability for residual connections.
|
| 30 |
+
s1_bits (int): Number of bits for the pre token in BSQuantizer.
|
| 31 |
+
s2_bits (int): Number of bits for the post token in BSQuantizer.
|
| 32 |
+
beta (float): Beta parameter for BSQuantizer.
|
| 33 |
+
gamma0 (float): Gamma0 parameter for BSQuantizer.
|
| 34 |
+
gamma (float): Gamma parameter for BSQuantizer.
|
| 35 |
+
zeta (float): Zeta parameter for BSQuantizer.
|
| 36 |
+
group_size (int): Group size parameter for BSQuantizer.
|
| 37 |
+
|
| 38 |
+
"""
|
| 39 |
+
|
| 40 |
+
def __init__(self, d_in, d_model, n_heads, ff_dim, n_enc_layers, n_dec_layers, ffn_dropout_p, attn_dropout_p, resid_dropout_p, s1_bits, s2_bits, beta, gamma0, gamma, zeta, group_size):
|
| 41 |
+
|
| 42 |
+
super().__init__()
|
| 43 |
+
self.d_in = d_in
|
| 44 |
+
self.d_model = d_model
|
| 45 |
+
self.n_heads = n_heads
|
| 46 |
+
self.ff_dim = ff_dim
|
| 47 |
+
self.enc_layers = n_enc_layers
|
| 48 |
+
self.dec_layers = n_dec_layers
|
| 49 |
+
self.ffn_dropout_p = ffn_dropout_p
|
| 50 |
+
self.attn_dropout_p = attn_dropout_p
|
| 51 |
+
self.resid_dropout_p = resid_dropout_p
|
| 52 |
+
|
| 53 |
+
self.s1_bits = s1_bits
|
| 54 |
+
self.s2_bits = s2_bits
|
| 55 |
+
self.codebook_dim = s1_bits + s2_bits # Total dimension of the codebook after quantization
|
| 56 |
+
self.embed = nn.Linear(self.d_in, self.d_model)
|
| 57 |
+
self.head = nn.Linear(self.d_model, self.d_in)
|
| 58 |
+
|
| 59 |
+
# Encoder Transformer Blocks
|
| 60 |
+
self.encoder = nn.ModuleList([
|
| 61 |
+
TransformerBlock(self.d_model, self.n_heads, self.ff_dim, self.ffn_dropout_p, self.attn_dropout_p, self.resid_dropout_p)
|
| 62 |
+
for _ in range(self.enc_layers - 1)
|
| 63 |
+
])
|
| 64 |
+
# Decoder Transformer Blocks
|
| 65 |
+
self.decoder = nn.ModuleList([
|
| 66 |
+
TransformerBlock(self.d_model, self.n_heads, self.ff_dim, self.ffn_dropout_p, self.attn_dropout_p, self.resid_dropout_p)
|
| 67 |
+
for _ in range(self.dec_layers - 1)
|
| 68 |
+
])
|
| 69 |
+
self.quant_embed = nn.Linear(in_features=self.d_model, out_features=self.codebook_dim) # Linear layer before quantization
|
| 70 |
+
self.post_quant_embed_pre = nn.Linear(in_features=self.s1_bits, out_features=self.d_model) # Linear layer after quantization (pre part - s1 bits)
|
| 71 |
+
self.post_quant_embed = nn.Linear(in_features=self.codebook_dim, out_features=self.d_model) # Linear layer after quantization (full codebook)
|
| 72 |
+
self.tokenizer = BSQuantizer(self.s1_bits, self.s2_bits, beta, gamma0, gamma, zeta, group_size) # BSQuantizer module
|
| 73 |
+
|
| 74 |
+
def forward(self, x):
|
| 75 |
+
"""
|
| 76 |
+
Forward pass of the KronosTokenizer.
|
| 77 |
+
|
| 78 |
+
Args:
|
| 79 |
+
x (torch.Tensor): Input tensor of shape (batch_size, seq_len, d_in).
|
| 80 |
+
|
| 81 |
+
Returns:
|
| 82 |
+
tuple: A tuple containing:
|
| 83 |
+
- tuple: (z_pre, z) - Reconstructed outputs from decoder with s1_bits and full codebook respectively,
|
| 84 |
+
both of shape (batch_size, seq_len, d_in).
|
| 85 |
+
- torch.Tensor: bsq_loss - Loss from the BSQuantizer.
|
| 86 |
+
- torch.Tensor: quantized - Quantized representation from BSQuantizer.
|
| 87 |
+
- torch.Tensor: z_indices - Indices from the BSQuantizer.
|
| 88 |
+
"""
|
| 89 |
+
z = self.embed(x)
|
| 90 |
+
|
| 91 |
+
for layer in self.encoder:
|
| 92 |
+
z = layer(z)
|
| 93 |
+
|
| 94 |
+
z = self.quant_embed(z) # (B, T, codebook)
|
| 95 |
+
|
| 96 |
+
bsq_loss, quantized, z_indices = self.tokenizer(z)
|
| 97 |
+
|
| 98 |
+
quantized_pre = quantized[:, :, :self.s1_bits] # Extract the first part of quantized representation (s1_bits)
|
| 99 |
+
z_pre = self.post_quant_embed_pre(quantized_pre)
|
| 100 |
+
|
| 101 |
+
z = self.post_quant_embed(quantized)
|
| 102 |
+
|
| 103 |
+
# Decoder layers (for pre part - s1 bits)
|
| 104 |
+
for layer in self.decoder:
|
| 105 |
+
z_pre = layer(z_pre)
|
| 106 |
+
z_pre = self.head(z_pre)
|
| 107 |
+
|
| 108 |
+
# Decoder layers (for full codebook)
|
| 109 |
+
for layer in self.decoder:
|
| 110 |
+
z = layer(z)
|
| 111 |
+
z = self.head(z)
|
| 112 |
+
|
| 113 |
+
return (z_pre, z), bsq_loss, quantized, z_indices
|
| 114 |
+
|
| 115 |
+
def indices_to_bits(self, x, half=False):
|
| 116 |
+
"""
|
| 117 |
+
Converts indices to bit representations and scales them.
|
| 118 |
+
|
| 119 |
+
Args:
|
| 120 |
+
x (torch.Tensor): Indices tensor.
|
| 121 |
+
half (bool, optional): Whether to process only half of the codebook dimension. Defaults to False.
|
| 122 |
+
|
| 123 |
+
Returns:
|
| 124 |
+
torch.Tensor: Bit representation tensor.
|
| 125 |
+
"""
|
| 126 |
+
if half:
|
| 127 |
+
x1 = x[0] # Assuming x is a tuple of indices if half is True
|
| 128 |
+
x2 = x[1]
|
| 129 |
+
mask = 2 ** torch.arange(self.codebook_dim//2, device=x1.device, dtype=torch.long) # Create a mask for bit extraction
|
| 130 |
+
x1 = (x1.unsqueeze(-1) & mask) != 0 # Extract bits for the first half
|
| 131 |
+
x2 = (x2.unsqueeze(-1) & mask) != 0 # Extract bits for the second half
|
| 132 |
+
x = torch.cat([x1, x2], dim=-1) # Concatenate the bit representations
|
| 133 |
+
else:
|
| 134 |
+
mask = 2 ** torch.arange(self.codebook_dim, device=x.device, dtype=torch.long) # Create a mask for bit extraction
|
| 135 |
+
x = (x.unsqueeze(-1) & mask) != 0 # Extract bits
|
| 136 |
+
|
| 137 |
+
x = x.float() * 2 - 1 # Convert boolean to bipolar (-1, 1)
|
| 138 |
+
q_scale = 1. / (self.codebook_dim ** 0.5) # Scaling factor
|
| 139 |
+
x = x * q_scale
|
| 140 |
+
return x
|
| 141 |
+
|
| 142 |
+
def encode(self, x, half=False):
|
| 143 |
+
"""
|
| 144 |
+
Encodes the input data into quantized indices.
|
| 145 |
+
|
| 146 |
+
Args:
|
| 147 |
+
x (torch.Tensor): Input tensor of shape (batch_size, seq_len, d_in).
|
| 148 |
+
half (bool, optional): Whether to use half quantization in BSQuantizer. Defaults to False.
|
| 149 |
+
|
| 150 |
+
Returns:
|
| 151 |
+
torch.Tensor: Quantized indices from BSQuantizer.
|
| 152 |
+
"""
|
| 153 |
+
z = self.embed(x)
|
| 154 |
+
for layer in self.encoder:
|
| 155 |
+
z = layer(z)
|
| 156 |
+
z = self.quant_embed(z)
|
| 157 |
+
|
| 158 |
+
bsq_loss, quantized, z_indices = self.tokenizer(z, half=half, collect_metrics=False)
|
| 159 |
+
return z_indices
|
| 160 |
+
|
| 161 |
+
def decode(self, x, half=False):
|
| 162 |
+
"""
|
| 163 |
+
Decodes quantized indices back to the input data space.
|
| 164 |
+
|
| 165 |
+
Args:
|
| 166 |
+
x (torch.Tensor): Quantized indices tensor.
|
| 167 |
+
half (bool, optional): Whether the indices were generated with half quantization. Defaults to False.
|
| 168 |
+
|
| 169 |
+
Returns:
|
| 170 |
+
torch.Tensor: Reconstructed output tensor of shape (batch_size, seq_len, d_in).
|
| 171 |
+
"""
|
| 172 |
+
quantized = self.indices_to_bits(x, half)
|
| 173 |
+
z = self.post_quant_embed(quantized)
|
| 174 |
+
for layer in self.decoder:
|
| 175 |
+
z = layer(z)
|
| 176 |
+
z = self.head(z)
|
| 177 |
+
return z
|
| 178 |
+
|
| 179 |
+
|
| 180 |
+
class Kronos(nn.Module, PyTorchModelHubMixin):
|
| 181 |
+
"""
|
| 182 |
+
Kronos Model.
|
| 183 |
+
|
| 184 |
+
Args:
|
| 185 |
+
s1_bits (int): Number of bits for pre tokens.
|
| 186 |
+
s2_bits (int): Number of bits for post tokens.
|
| 187 |
+
n_layers (int): Number of Transformer blocks.
|
| 188 |
+
d_model (int): Dimension of the model's embeddings and hidden states.
|
| 189 |
+
n_heads (int): Number of attention heads in the MultiheadAttention layers.
|
| 190 |
+
ff_dim (int): Dimension of the feedforward network in the Transformer blocks.
|
| 191 |
+
ffn_dropout_p (float): Dropout probability for the feedforward network.
|
| 192 |
+
attn_dropout_p (float): Dropout probability for the attention layers.
|
| 193 |
+
resid_dropout_p (float): Dropout probability for residual connections.
|
| 194 |
+
token_dropout_p (float): Dropout probability for token embeddings.
|
| 195 |
+
learn_te (bool): Whether to use learnable temporal embeddings.
|
| 196 |
+
"""
|
| 197 |
+
|
| 198 |
+
def __init__(self, s1_bits, s2_bits, n_layers, d_model, n_heads, ff_dim, ffn_dropout_p, attn_dropout_p, resid_dropout_p, token_dropout_p, learn_te):
|
| 199 |
+
super().__init__()
|
| 200 |
+
self.s1_bits = s1_bits
|
| 201 |
+
self.s2_bits = s2_bits
|
| 202 |
+
self.n_layers = n_layers
|
| 203 |
+
self.d_model = d_model
|
| 204 |
+
self.n_heads = n_heads
|
| 205 |
+
self.learn_te = learn_te
|
| 206 |
+
self.ff_dim = ff_dim
|
| 207 |
+
self.ffn_dropout_p = ffn_dropout_p
|
| 208 |
+
self.attn_dropout_p = attn_dropout_p
|
| 209 |
+
self.resid_dropout_p = resid_dropout_p
|
| 210 |
+
self.token_dropout_p = token_dropout_p
|
| 211 |
+
|
| 212 |
+
self.s1_vocab_size = 2 ** self.s1_bits
|
| 213 |
+
self.token_drop = nn.Dropout(self.token_dropout_p)
|
| 214 |
+
self.embedding = HierarchicalEmbedding(self.s1_bits, self.s2_bits, self.d_model)
|
| 215 |
+
self.time_emb = TemporalEmbedding(self.d_model, self.learn_te)
|
| 216 |
+
self.transformer = nn.ModuleList([
|
| 217 |
+
TransformerBlock(self.d_model, self.n_heads, self.ff_dim, self.ffn_dropout_p, self.attn_dropout_p, self.resid_dropout_p)
|
| 218 |
+
for _ in range(self.n_layers)
|
| 219 |
+
])
|
| 220 |
+
self.norm = RMSNorm(self.d_model)
|
| 221 |
+
self.dep_layer = DependencyAwareLayer(self.d_model)
|
| 222 |
+
self.head = DualHead(self.s1_bits, self.s2_bits, self.d_model)
|
| 223 |
+
self.apply(self._init_weights)
|
| 224 |
+
|
| 225 |
+
def _init_weights(self, module):
|
| 226 |
+
|
| 227 |
+
if isinstance(module, nn.Linear):
|
| 228 |
+
nn.init.xavier_normal_(module.weight)
|
| 229 |
+
if module.bias is not None:
|
| 230 |
+
nn.init.zeros_(module.bias)
|
| 231 |
+
elif isinstance(module, nn.Embedding):
|
| 232 |
+
nn.init.normal_(module.weight, mean=0, std=self.embedding.d_model ** -0.5)
|
| 233 |
+
elif isinstance(module, nn.LayerNorm):
|
| 234 |
+
nn.init.ones_(module.weight)
|
| 235 |
+
nn.init.zeros_(module.bias)
|
| 236 |
+
elif isinstance(module, RMSNorm):
|
| 237 |
+
nn.init.ones_(module.weight)
|
| 238 |
+
|
| 239 |
+
def forward(self, s1_ids, s2_ids, stamp=None, padding_mask=None, use_teacher_forcing=False, s1_targets=None):
|
| 240 |
+
"""
|
| 241 |
+
Args:
|
| 242 |
+
s1_ids (torch.Tensor): Input tensor of s1 token IDs. Shape: [batch_size, seq_len]
|
| 243 |
+
s2_ids (torch.Tensor): Input tensor of s2 token IDs. Shape: [batch_size, seq_len]
|
| 244 |
+
stamp (torch.Tensor, optional): Temporal stamp tensor. Shape: [batch_size, seq_len]. Defaults to None.
|
| 245 |
+
padding_mask (torch.Tensor, optional): Mask for padding tokens. Shape: [batch_size, seq_len]. Defaults to None.
|
| 246 |
+
use_teacher_forcing (bool, optional): Whether to use teacher forcing for s1 decoding. Defaults to False.
|
| 247 |
+
s1_targets (torch.Tensor, optional): Target s1 token IDs for teacher forcing. Shape: [batch_size, seq_len]. Defaults to None.
|
| 248 |
+
|
| 249 |
+
Returns:
|
| 250 |
+
Tuple[torch.Tensor, torch.Tensor]:
|
| 251 |
+
- s1 logits: Logits for s1 token predictions. Shape: [batch_size, seq_len, s1_vocab_size]
|
| 252 |
+
- s2_logits: Logits for s2 token predictions, conditioned on s1. Shape: [batch_size, seq_len, s2_vocab_size]
|
| 253 |
+
"""
|
| 254 |
+
x = self.embedding([s1_ids, s2_ids])
|
| 255 |
+
if stamp is not None:
|
| 256 |
+
time_embedding = self.time_emb(stamp)
|
| 257 |
+
x = x + time_embedding
|
| 258 |
+
x = self.token_drop(x)
|
| 259 |
+
|
| 260 |
+
for layer in self.transformer:
|
| 261 |
+
x = layer(x, key_padding_mask=padding_mask)
|
| 262 |
+
|
| 263 |
+
x = self.norm(x)
|
| 264 |
+
|
| 265 |
+
s1_logits = self.head(x)
|
| 266 |
+
|
| 267 |
+
if use_teacher_forcing:
|
| 268 |
+
sibling_embed = self.embedding.emb_s1(s1_targets)
|
| 269 |
+
else:
|
| 270 |
+
s1_probs = F.softmax(s1_logits.detach(), dim=-1)
|
| 271 |
+
sample_s1_ids = torch.multinomial(s1_probs.view(-1, self.s1_vocab_size), 1).view(s1_ids.shape)
|
| 272 |
+
sibling_embed = self.embedding.emb_s1(sample_s1_ids)
|
| 273 |
+
|
| 274 |
+
x2 = self.dep_layer(x, sibling_embed, key_padding_mask=padding_mask) # Dependency Aware Layer: Condition on s1 embeddings
|
| 275 |
+
s2_logits = self.head.cond_forward(x2)
|
| 276 |
+
return s1_logits, s2_logits
|
| 277 |
+
|
| 278 |
+
def decode_s1(self, s1_ids, s2_ids, stamp=None, padding_mask=None):
|
| 279 |
+
"""
|
| 280 |
+
Decodes only the s1 tokens.
|
| 281 |
+
|
| 282 |
+
This method performs a forward pass to predict only s1 tokens. It returns the s1 logits
|
| 283 |
+
and the context representation from the Transformer, which can be used for subsequent s2 decoding.
|
| 284 |
+
|
| 285 |
+
Args:
|
| 286 |
+
s1_ids (torch.Tensor): Input tensor of s1 token IDs. Shape: [batch_size, seq_len]
|
| 287 |
+
s2_ids (torch.Tensor): Input tensor of s2 token IDs. Shape: [batch_size, seq_len]
|
| 288 |
+
stamp (torch.Tensor, optional): Temporal stamp tensor. Shape: [batch_size, seq_len]. Defaults to None.
|
| 289 |
+
padding_mask (torch.Tensor, optional): Mask for padding tokens. Shape: [batch_size, seq_len]. Defaults to None.
|
| 290 |
+
|
| 291 |
+
Returns:
|
| 292 |
+
Tuple[torch.Tensor, torch.Tensor]:
|
| 293 |
+
- s1 logits: Logits for s1 token predictions. Shape: [batch_size, seq_len, s1_vocab_size]
|
| 294 |
+
- context: Context representation from the Transformer. Shape: [batch_size, seq_len, d_model]
|
| 295 |
+
"""
|
| 296 |
+
x = self.embedding([s1_ids, s2_ids])
|
| 297 |
+
if stamp is not None:
|
| 298 |
+
time_embedding = self.time_emb(stamp)
|
| 299 |
+
x = x + time_embedding
|
| 300 |
+
x = self.token_drop(x)
|
| 301 |
+
|
| 302 |
+
for layer in self.transformer:
|
| 303 |
+
x = layer(x, key_padding_mask=padding_mask)
|
| 304 |
+
|
| 305 |
+
x = self.norm(x)
|
| 306 |
+
|
| 307 |
+
s1_logits = self.head(x)
|
| 308 |
+
return s1_logits, x
|
| 309 |
+
|
| 310 |
+
def decode_s2(self, context, s1_ids, padding_mask=None):
|
| 311 |
+
"""
|
| 312 |
+
Decodes the s2 tokens, conditioned on the context and s1 tokens.
|
| 313 |
+
|
| 314 |
+
This method decodes s2 tokens based on a pre-computed context representation (typically from `decode_s1`)
|
| 315 |
+
and the s1 token IDs. It uses the dependency-aware layer and the conditional s2 head to predict s2 tokens.
|
| 316 |
+
|
| 317 |
+
Args:
|
| 318 |
+
context (torch.Tensor): Context representation from the transformer (output of decode_s1).
|
| 319 |
+
Shape: [batch_size, seq_len, d_model]
|
| 320 |
+
s1_ids (torch.Tensor): Input tensor of s1 token IDs. Shape: [batch_size, seq_len]
|
| 321 |
+
padding_mask (torch.Tensor, optional): Mask for padding tokens. Shape: [batch_size, seq_len]. Defaults to None.
|
| 322 |
+
|
| 323 |
+
Returns:
|
| 324 |
+
torch.Tensor: s2 logits. Shape: [batch_size, seq_len, s2_vocab_size]
|
| 325 |
+
"""
|
| 326 |
+
sibling_embed = self.embedding.emb_s1(s1_ids)
|
| 327 |
+
x2 = self.dep_layer(context, sibling_embed, key_padding_mask=padding_mask)
|
| 328 |
+
return self.head.cond_forward(x2)
|
| 329 |
+
|
| 330 |
+
|
| 331 |
+
def top_k_top_p_filtering(
|
| 332 |
+
logits,
|
| 333 |
+
top_k: int = 0,
|
| 334 |
+
top_p: float = 1.0,
|
| 335 |
+
filter_value: float = -float("Inf"),
|
| 336 |
+
min_tokens_to_keep: int = 1,
|
| 337 |
+
):
|
| 338 |
+
"""Filter a distribution of logits using top-k and/or nucleus (top-p) filtering
|
| 339 |
+
Args:
|
| 340 |
+
logits: logits distribution shape (batch size, vocabulary size)
|
| 341 |
+
if top_k > 0: keep only top k tokens with highest probability (top-k filtering).
|
| 342 |
+
if top_p < 1.0: keep the top tokens with cumulative probability >= top_p (nucleus filtering).
|
| 343 |
+
Nucleus filtering is described in Holtzman et al. (http://arxiv.org/abs/1904.09751)
|
| 344 |
+
Make sure we keep at least min_tokens_to_keep per batch example in the output
|
| 345 |
+
From: https://gist.github.com/thomwolf/1a5a29f6962089e871b94cbd09daf317
|
| 346 |
+
"""
|
| 347 |
+
if top_k > 0:
|
| 348 |
+
top_k = min(max(top_k, min_tokens_to_keep), logits.size(-1)) # Safety check
|
| 349 |
+
# Remove all tokens with a probability less than the last token of the top-k
|
| 350 |
+
indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
|
| 351 |
+
logits[indices_to_remove] = filter_value
|
| 352 |
+
return logits
|
| 353 |
+
|
| 354 |
+
if top_p < 1.0:
|
| 355 |
+
sorted_logits, sorted_indices = torch.sort(logits, descending=True)
|
| 356 |
+
cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
|
| 357 |
+
|
| 358 |
+
# Remove tokens with cumulative probability above the threshold (token with 0 are kept)
|
| 359 |
+
sorted_indices_to_remove = cumulative_probs > top_p
|
| 360 |
+
if min_tokens_to_keep > 1:
|
| 361 |
+
# Keep at least min_tokens_to_keep (set to min_tokens_to_keep-1 because we add the first one below)
|
| 362 |
+
sorted_indices_to_remove[..., :min_tokens_to_keep] = 0
|
| 363 |
+
# Shift the indices to the right to keep also the first token above the threshold
|
| 364 |
+
sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
|
| 365 |
+
sorted_indices_to_remove[..., 0] = 0
|
| 366 |
+
|
| 367 |
+
# scatter sorted tensors to original indexing
|
| 368 |
+
indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
|
| 369 |
+
logits[indices_to_remove] = filter_value
|
| 370 |
+
return logits
|
| 371 |
+
|
| 372 |
+
|
| 373 |
+
def sample_from_logits(logits, temperature=1.0, top_k=None, top_p=None, sample_logits=True):
|
| 374 |
+
logits = logits / temperature
|
| 375 |
+
if top_k is not None or top_p is not None:
|
| 376 |
+
if top_k > 0 or top_p < 1.0:
|
| 377 |
+
logits = top_k_top_p_filtering(logits, top_k=top_k, top_p=top_p)
|
| 378 |
+
|
| 379 |
+
probs = F.softmax(logits, dim=-1)
|
| 380 |
+
|
| 381 |
+
if not sample_logits:
|
| 382 |
+
_, x = torch.topk(probs, k=1, dim=-1)
|
| 383 |
+
else:
|
| 384 |
+
x = torch.multinomial(probs, num_samples=1)
|
| 385 |
+
|
| 386 |
+
return x
|
| 387 |
+
|
| 388 |
+
|
| 389 |
+
def auto_regressive_inference(tokenizer, model, x, x_stamp, y_stamp, max_context, pred_len, clip=5, T=1.0, top_k=0, top_p=0.99, sample_count=5, verbose=False):
|
| 390 |
+
with torch.no_grad():
|
| 391 |
+
x = torch.clip(x, -clip, clip)
|
| 392 |
+
|
| 393 |
+
device = x.device
|
| 394 |
+
x = x.unsqueeze(1).repeat(1, sample_count, 1, 1).reshape(-1, x.size(1), x.size(2)).to(device)
|
| 395 |
+
x_stamp = x_stamp.unsqueeze(1).repeat(1, sample_count, 1, 1).reshape(-1, x_stamp.size(1), x_stamp.size(2)).to(device)
|
| 396 |
+
y_stamp = y_stamp.unsqueeze(1).repeat(1, sample_count, 1, 1).reshape(-1, y_stamp.size(1), y_stamp.size(2)).to(device)
|
| 397 |
+
|
| 398 |
+
x_token = tokenizer.encode(x, half=True)
|
| 399 |
+
|
| 400 |
+
initial_seq_len = x.size(1)
|
| 401 |
+
batch_size = x_token[0].size(0)
|
| 402 |
+
total_seq_len = initial_seq_len + pred_len
|
| 403 |
+
full_stamp = torch.cat([x_stamp, y_stamp], dim=1)
|
| 404 |
+
|
| 405 |
+
generated_pre = x_token[0].new_empty(batch_size, pred_len)
|
| 406 |
+
generated_post = x_token[1].new_empty(batch_size, pred_len)
|
| 407 |
+
|
| 408 |
+
pre_buffer = x_token[0].new_zeros(batch_size, max_context)
|
| 409 |
+
post_buffer = x_token[1].new_zeros(batch_size, max_context)
|
| 410 |
+
buffer_len = min(initial_seq_len, max_context)
|
| 411 |
+
if buffer_len > 0:
|
| 412 |
+
start_idx = max(0, initial_seq_len - max_context)
|
| 413 |
+
pre_buffer[:, :buffer_len] = x_token[0][:, start_idx:start_idx + buffer_len]
|
| 414 |
+
post_buffer[:, :buffer_len] = x_token[1][:, start_idx:start_idx + buffer_len]
|
| 415 |
+
|
| 416 |
+
if verbose:
|
| 417 |
+
ran = trange
|
| 418 |
+
else:
|
| 419 |
+
ran = range
|
| 420 |
+
for i in ran(pred_len):
|
| 421 |
+
current_seq_len = initial_seq_len + i
|
| 422 |
+
window_len = min(current_seq_len, max_context)
|
| 423 |
+
|
| 424 |
+
if current_seq_len <= max_context:
|
| 425 |
+
input_tokens = [
|
| 426 |
+
pre_buffer[:, :window_len],
|
| 427 |
+
post_buffer[:, :window_len]
|
| 428 |
+
]
|
| 429 |
+
else:
|
| 430 |
+
input_tokens = [pre_buffer, post_buffer]
|
| 431 |
+
|
| 432 |
+
context_end = current_seq_len
|
| 433 |
+
context_start = max(0, context_end - max_context)
|
| 434 |
+
current_stamp = full_stamp[:, context_start:context_end, :].contiguous()
|
| 435 |
+
|
| 436 |
+
s1_logits, context = model.decode_s1(input_tokens[0], input_tokens[1], current_stamp)
|
| 437 |
+
s1_logits = s1_logits[:, -1, :]
|
| 438 |
+
sample_pre = sample_from_logits(s1_logits, temperature=T, top_k=top_k, top_p=top_p, sample_logits=True)
|
| 439 |
+
|
| 440 |
+
s2_logits = model.decode_s2(context, sample_pre)
|
| 441 |
+
s2_logits = s2_logits[:, -1, :]
|
| 442 |
+
sample_post = sample_from_logits(s2_logits, temperature=T, top_k=top_k, top_p=top_p, sample_logits=True)
|
| 443 |
+
|
| 444 |
+
generated_pre[:, i] = sample_pre.squeeze(-1)
|
| 445 |
+
generated_post[:, i] = sample_post.squeeze(-1)
|
| 446 |
+
|
| 447 |
+
if current_seq_len < max_context:
|
| 448 |
+
pre_buffer[:, current_seq_len] = sample_pre.squeeze(-1)
|
| 449 |
+
post_buffer[:, current_seq_len] = sample_post.squeeze(-1)
|
| 450 |
+
else:
|
| 451 |
+
pre_buffer.copy_(torch.roll(pre_buffer, shifts=-1, dims=1))
|
| 452 |
+
post_buffer.copy_(torch.roll(post_buffer, shifts=-1, dims=1))
|
| 453 |
+
pre_buffer[:, -1] = sample_pre.squeeze(-1)
|
| 454 |
+
post_buffer[:, -1] = sample_post.squeeze(-1)
|
| 455 |
+
|
| 456 |
+
full_pre = torch.cat([x_token[0], generated_pre], dim=1)
|
| 457 |
+
full_post = torch.cat([x_token[1], generated_post], dim=1)
|
| 458 |
+
|
| 459 |
+
context_start = max(0, total_seq_len - max_context)
|
| 460 |
+
input_tokens = [
|
| 461 |
+
full_pre[:, context_start:total_seq_len].contiguous(),
|
| 462 |
+
full_post[:, context_start:total_seq_len].contiguous()
|
| 463 |
+
]
|
| 464 |
+
z = tokenizer.decode(input_tokens, half=True)
|
| 465 |
+
z = z.reshape(-1, sample_count, z.size(1), z.size(2))
|
| 466 |
+
preds = z.cpu().numpy()
|
| 467 |
+
preds = np.mean(preds, axis=1)
|
| 468 |
+
|
| 469 |
+
return preds
|
| 470 |
+
|
| 471 |
+
|
| 472 |
+
def calc_time_stamps(x_timestamp):
|
| 473 |
+
time_df = pd.DataFrame()
|
| 474 |
+
time_df['minute'] = x_timestamp.dt.minute
|
| 475 |
+
time_df['hour'] = x_timestamp.dt.hour
|
| 476 |
+
time_df['weekday'] = x_timestamp.dt.weekday
|
| 477 |
+
time_df['day'] = x_timestamp.dt.day
|
| 478 |
+
time_df['month'] = x_timestamp.dt.month
|
| 479 |
+
return time_df
|
| 480 |
+
|
| 481 |
+
|
| 482 |
+
class KronosPredictor:
|
| 483 |
+
|
| 484 |
+
def __init__(self, model, tokenizer, device=None, max_context=512, clip=5):
|
| 485 |
+
self.tokenizer = tokenizer
|
| 486 |
+
self.model = model
|
| 487 |
+
self.max_context = max_context
|
| 488 |
+
self.clip = clip
|
| 489 |
+
self.price_cols = ['open', 'high', 'low', 'close']
|
| 490 |
+
self.vol_col = 'volume'
|
| 491 |
+
self.amt_vol = 'amount'
|
| 492 |
+
self.time_cols = ['minute', 'hour', 'weekday', 'day', 'month']
|
| 493 |
+
|
| 494 |
+
# Auto-detect device if not specified
|
| 495 |
+
if device is None:
|
| 496 |
+
if torch.cuda.is_available():
|
| 497 |
+
device = "cuda:0"
|
| 498 |
+
elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
|
| 499 |
+
device = "mps"
|
| 500 |
+
else:
|
| 501 |
+
device = "cpu"
|
| 502 |
+
|
| 503 |
+
self.device = device
|
| 504 |
+
|
| 505 |
+
self.tokenizer = self.tokenizer.to(self.device)
|
| 506 |
+
self.model = self.model.to(self.device)
|
| 507 |
+
|
| 508 |
+
def generate(self, x, x_stamp, y_stamp, pred_len, T, top_k, top_p, sample_count, verbose):
|
| 509 |
+
|
| 510 |
+
x_tensor = torch.from_numpy(np.array(x).astype(np.float32)).to(self.device)
|
| 511 |
+
x_stamp_tensor = torch.from_numpy(np.array(x_stamp).astype(np.float32)).to(self.device)
|
| 512 |
+
y_stamp_tensor = torch.from_numpy(np.array(y_stamp).astype(np.float32)).to(self.device)
|
| 513 |
+
|
| 514 |
+
preds = auto_regressive_inference(self.tokenizer, self.model, x_tensor, x_stamp_tensor, y_stamp_tensor, self.max_context, pred_len,
|
| 515 |
+
self.clip, T, top_k, top_p, sample_count, verbose)
|
| 516 |
+
preds = preds[:, -pred_len:, :]
|
| 517 |
+
return preds
|
| 518 |
+
|
| 519 |
+
def predict(self, df, x_timestamp, y_timestamp, pred_len, T=1.0, top_k=0, top_p=0.9, sample_count=1, verbose=True):
|
| 520 |
+
|
| 521 |
+
if not isinstance(df, pd.DataFrame):
|
| 522 |
+
raise ValueError("Input must be a pandas DataFrame.")
|
| 523 |
+
|
| 524 |
+
if not all(col in df.columns for col in self.price_cols):
|
| 525 |
+
raise ValueError(f"Price columns {self.price_cols} not found in DataFrame.")
|
| 526 |
+
|
| 527 |
+
df = df.copy()
|
| 528 |
+
if self.vol_col not in df.columns:
|
| 529 |
+
df[self.vol_col] = 0.0 # Fill missing volume with zeros
|
| 530 |
+
df[self.amt_vol] = 0.0 # Fill missing amount with zeros
|
| 531 |
+
if self.amt_vol not in df.columns and self.vol_col in df.columns:
|
| 532 |
+
df[self.amt_vol] = df[self.vol_col] * df[self.price_cols].mean(axis=1)
|
| 533 |
+
|
| 534 |
+
if df[self.price_cols + [self.vol_col, self.amt_vol]].isnull().values.any():
|
| 535 |
+
raise ValueError("Input DataFrame contains NaN values in price or volume columns.")
|
| 536 |
+
|
| 537 |
+
x_time_df = calc_time_stamps(x_timestamp)
|
| 538 |
+
y_time_df = calc_time_stamps(y_timestamp)
|
| 539 |
+
|
| 540 |
+
x = df[self.price_cols + [self.vol_col, self.amt_vol]].values.astype(np.float32)
|
| 541 |
+
x_stamp = x_time_df.values.astype(np.float32)
|
| 542 |
+
y_stamp = y_time_df.values.astype(np.float32)
|
| 543 |
+
|
| 544 |
+
x_mean, x_std = np.mean(x, axis=0), np.std(x, axis=0)
|
| 545 |
+
|
| 546 |
+
x = (x - x_mean) / (x_std + 1e-5)
|
| 547 |
+
x = np.clip(x, -self.clip, self.clip)
|
| 548 |
+
|
| 549 |
+
x = x[np.newaxis, :]
|
| 550 |
+
x_stamp = x_stamp[np.newaxis, :]
|
| 551 |
+
y_stamp = y_stamp[np.newaxis, :]
|
| 552 |
+
|
| 553 |
+
preds = self.generate(x, x_stamp, y_stamp, pred_len, T, top_k, top_p, sample_count, verbose)
|
| 554 |
+
|
| 555 |
+
preds = preds.squeeze(0)
|
| 556 |
+
preds = preds * (x_std + 1e-5) + x_mean
|
| 557 |
+
|
| 558 |
+
pred_df = pd.DataFrame(preds, columns=self.price_cols + [self.vol_col, self.amt_vol], index=y_timestamp)
|
| 559 |
+
return pred_df
|
| 560 |
+
|
| 561 |
+
|
| 562 |
+
def predict_batch(self, df_list, x_timestamp_list, y_timestamp_list, pred_len, T=1.0, top_k=0, top_p=0.9, sample_count=1, verbose=True):
|
| 563 |
+
"""
|
| 564 |
+
Perform parallel (batch) prediction on multiple time series. All series must have the same historical length and prediction length (pred_len).
|
| 565 |
+
|
| 566 |
+
Args:
|
| 567 |
+
df_list (List[pd.DataFrame]): List of input DataFrames, each containing price columns and optional volume/amount columns.
|
| 568 |
+
x_timestamp_list (List[pd.DatetimeIndex or Series]): List of timestamps corresponding to historical data, length should match the number of rows in each DataFrame.
|
| 569 |
+
y_timestamp_list (List[pd.DatetimeIndex or Series]): List of future prediction timestamps, length should equal pred_len.
|
| 570 |
+
pred_len (int): Number of prediction steps.
|
| 571 |
+
T (float): Sampling temperature.
|
| 572 |
+
top_k (int): Top-k filtering threshold.
|
| 573 |
+
top_p (float): Top-p (nucleus sampling) threshold.
|
| 574 |
+
sample_count (int): Number of parallel samples per series, automatically averaged internally.
|
| 575 |
+
verbose (bool): Whether to display autoregressive progress.
|
| 576 |
+
|
| 577 |
+
Returns:
|
| 578 |
+
List[pd.DataFrame]: List of prediction results in the same order as input, each DataFrame contains
|
| 579 |
+
`open, high, low, close, volume, amount` columns, indexed by corresponding `y_timestamp`.
|
| 580 |
+
"""
|
| 581 |
+
# Basic validation
|
| 582 |
+
if not isinstance(df_list, (list, tuple)) or not isinstance(x_timestamp_list, (list, tuple)) or not isinstance(y_timestamp_list, (list, tuple)):
|
| 583 |
+
raise ValueError("df_list, x_timestamp_list, y_timestamp_list must be list or tuple types.")
|
| 584 |
+
if not (len(df_list) == len(x_timestamp_list) == len(y_timestamp_list)):
|
| 585 |
+
raise ValueError("df_list, x_timestamp_list, y_timestamp_list must have consistent lengths.")
|
| 586 |
+
|
| 587 |
+
num_series = len(df_list)
|
| 588 |
+
|
| 589 |
+
x_list = []
|
| 590 |
+
x_stamp_list = []
|
| 591 |
+
y_stamp_list = []
|
| 592 |
+
means = []
|
| 593 |
+
stds = []
|
| 594 |
+
seq_lens = []
|
| 595 |
+
y_lens = []
|
| 596 |
+
|
| 597 |
+
for i in range(num_series):
|
| 598 |
+
df = df_list[i]
|
| 599 |
+
if not isinstance(df, pd.DataFrame):
|
| 600 |
+
raise ValueError(f"Input at index {i} is not a pandas DataFrame.")
|
| 601 |
+
if not all(col in df.columns for col in self.price_cols):
|
| 602 |
+
raise ValueError(f"DataFrame at index {i} is missing price columns {self.price_cols}.")
|
| 603 |
+
|
| 604 |
+
df = df.copy()
|
| 605 |
+
if self.vol_col not in df.columns:
|
| 606 |
+
df[self.vol_col] = 0.0
|
| 607 |
+
df[self.amt_vol] = 0.0
|
| 608 |
+
if self.amt_vol not in df.columns and self.vol_col in df.columns:
|
| 609 |
+
df[self.amt_vol] = df[self.vol_col] * df[self.price_cols].mean(axis=1)
|
| 610 |
+
|
| 611 |
+
if df[self.price_cols + [self.vol_col, self.amt_vol]].isnull().values.any():
|
| 612 |
+
raise ValueError(f"DataFrame at index {i} contains NaN values in price or volume columns.")
|
| 613 |
+
|
| 614 |
+
x_timestamp = x_timestamp_list[i]
|
| 615 |
+
y_timestamp = y_timestamp_list[i]
|
| 616 |
+
|
| 617 |
+
x_time_df = calc_time_stamps(x_timestamp)
|
| 618 |
+
y_time_df = calc_time_stamps(y_timestamp)
|
| 619 |
+
|
| 620 |
+
x = df[self.price_cols + [self.vol_col, self.amt_vol]].values.astype(np.float32)
|
| 621 |
+
x_stamp = x_time_df.values.astype(np.float32)
|
| 622 |
+
y_stamp = y_time_df.values.astype(np.float32)
|
| 623 |
+
|
| 624 |
+
if x.shape[0] != x_stamp.shape[0]:
|
| 625 |
+
raise ValueError(f"Inconsistent lengths at index {i}: x has {x.shape[0]} vs x_stamp has {x_stamp.shape[0]}.")
|
| 626 |
+
if y_stamp.shape[0] != pred_len:
|
| 627 |
+
raise ValueError(f"y_timestamp length at index {i} should equal pred_len={pred_len}, got {y_stamp.shape[0]}.")
|
| 628 |
+
|
| 629 |
+
x_mean, x_std = np.mean(x, axis=0), np.std(x, axis=0)
|
| 630 |
+
x_norm = (x - x_mean) / (x_std + 1e-5)
|
| 631 |
+
x_norm = np.clip(x_norm, -self.clip, self.clip)
|
| 632 |
+
|
| 633 |
+
x_list.append(x_norm)
|
| 634 |
+
x_stamp_list.append(x_stamp)
|
| 635 |
+
y_stamp_list.append(y_stamp)
|
| 636 |
+
means.append(x_mean)
|
| 637 |
+
stds.append(x_std)
|
| 638 |
+
|
| 639 |
+
seq_lens.append(x_norm.shape[0])
|
| 640 |
+
y_lens.append(y_stamp.shape[0])
|
| 641 |
+
|
| 642 |
+
# Require all series to have consistent historical and prediction lengths for batch processing
|
| 643 |
+
if len(set(seq_lens)) != 1:
|
| 644 |
+
raise ValueError(f"Parallel prediction requires all series to have consistent historical lengths, got: {seq_lens}")
|
| 645 |
+
if len(set(y_lens)) != 1:
|
| 646 |
+
raise ValueError(f"Parallel prediction requires all series to have consistent prediction lengths, got: {y_lens}")
|
| 647 |
+
|
| 648 |
+
x_batch = np.stack(x_list, axis=0).astype(np.float32) # (B, seq_len, feat)
|
| 649 |
+
x_stamp_batch = np.stack(x_stamp_list, axis=0).astype(np.float32) # (B, seq_len, time_feat)
|
| 650 |
+
y_stamp_batch = np.stack(y_stamp_list, axis=0).astype(np.float32) # (B, pred_len, time_feat)
|
| 651 |
+
|
| 652 |
+
preds = self.generate(x_batch, x_stamp_batch, y_stamp_batch, pred_len, T, top_k, top_p, sample_count, verbose)
|
| 653 |
+
# preds: (B, pred_len, feat)
|
| 654 |
+
|
| 655 |
+
pred_dfs = []
|
| 656 |
+
for i in range(num_series):
|
| 657 |
+
preds_i = preds[i] * (stds[i] + 1e-5) + means[i]
|
| 658 |
+
pred_df = pd.DataFrame(preds_i, columns=self.price_cols + [self.vol_col, self.amt_vol], index=y_timestamp_list[i])
|
| 659 |
+
pred_dfs.append(pred_df)
|
| 660 |
+
|
| 661 |
+
return pred_dfs
|
| 662 |
+
|
model/module.py
ADDED
|
@@ -0,0 +1,570 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import math
|
| 2 |
+
|
| 3 |
+
from einops import rearrange, reduce
|
| 4 |
+
import torch
|
| 5 |
+
import torch.nn as nn
|
| 6 |
+
from torch.autograd import Function
|
| 7 |
+
import torch.nn.functional as F
|
| 8 |
+
|
| 9 |
+
|
| 10 |
+
class DifferentiableEntropyFunction(Function):
|
| 11 |
+
@staticmethod
|
| 12 |
+
def forward(ctx, zq, basis, K, eps):
|
| 13 |
+
zb = (zq + 1) / 2
|
| 14 |
+
zi = ((zb * basis).sum(-1)).to(torch.int64)
|
| 15 |
+
cnt = torch.scatter_reduce(torch.zeros(2 ** K, device=zq.device, dtype=zq.dtype),
|
| 16 |
+
0,
|
| 17 |
+
zi.flatten(),
|
| 18 |
+
torch.ones_like(zi.flatten()).to(zq.dtype),
|
| 19 |
+
'sum')
|
| 20 |
+
prob = (cnt + eps) / (cnt + eps).sum()
|
| 21 |
+
H = -(prob * torch.log(prob)).sum()
|
| 22 |
+
ctx.save_for_backward(zq, zi, prob)
|
| 23 |
+
ctx.K = K
|
| 24 |
+
return H
|
| 25 |
+
|
| 26 |
+
@staticmethod
|
| 27 |
+
def backward(ctx, grad_output):
|
| 28 |
+
zq, zi, prob = ctx.saved_tensors
|
| 29 |
+
grad_array = -grad_output * (torch.log(prob) + 1) / zi.numel() / ctx.K
|
| 30 |
+
reord_grad = grad_array[zi.flatten()].reshape(zi.shape)
|
| 31 |
+
grad_input = reord_grad.unsqueeze(-1) * zq
|
| 32 |
+
return grad_input, None, None, None, None
|
| 33 |
+
|
| 34 |
+
|
| 35 |
+
def codebook_entropy(zq, basis, K, eps=1e-4):
|
| 36 |
+
return DifferentiableEntropyFunction.apply(zq, basis, K, eps)
|
| 37 |
+
|
| 38 |
+
|
| 39 |
+
class BinarySphericalQuantizer(nn.Module):
|
| 40 |
+
def __init__(self, embed_dim, beta, gamma0, gamma, zeta,
|
| 41 |
+
input_format='bchw',
|
| 42 |
+
soft_entropy=True, group_size=9,
|
| 43 |
+
persample_entropy_compute='analytical',
|
| 44 |
+
cb_entropy_compute='group',
|
| 45 |
+
l2_norm=True,
|
| 46 |
+
inv_temperature=1):
|
| 47 |
+
"""
|
| 48 |
+
Paper link: https://arxiv.org/pdf/2406.07548.pdf
|
| 49 |
+
Here we use the official implementation of the BinarySphericalQuantizer.
|
| 50 |
+
"""
|
| 51 |
+
super().__init__()
|
| 52 |
+
self.embed_dim = embed_dim
|
| 53 |
+
self.beta = beta # loss weight for commit loss
|
| 54 |
+
self.gamma0 = gamma0 # loss weight for entropy penalty
|
| 55 |
+
self.gamma = gamma # loss weight for entropy penalty
|
| 56 |
+
self.zeta = zeta # loss weight for entire entropy penalty
|
| 57 |
+
self.input_format = input_format
|
| 58 |
+
assert self.embed_dim % group_size == 0, "embed_dim must be divisible by group_size"
|
| 59 |
+
self.num_groups = self.embed_dim // group_size
|
| 60 |
+
self.group_size = group_size
|
| 61 |
+
assert persample_entropy_compute in ['group', 'analytical'], "persample_entropy_compute must be either 'group' or 'analytical'"
|
| 62 |
+
assert cb_entropy_compute in ['group', 'nce'], "cb_entropy_compute must be either 'group' or 'nce'"
|
| 63 |
+
self.persample_entropy_compute = persample_entropy_compute
|
| 64 |
+
self.cb_entropy_compute = cb_entropy_compute
|
| 65 |
+
self.l2_norm = l2_norm
|
| 66 |
+
self.inv_temperature = inv_temperature
|
| 67 |
+
|
| 68 |
+
self.register_buffer('basis', 2 ** torch.arange(embed_dim - 1, -1, -1))
|
| 69 |
+
self.register_buffer('group_basis', 2 ** torch.arange(group_size - 1, -1, -1))
|
| 70 |
+
|
| 71 |
+
self.num_dimensions = 2 ** embed_dim
|
| 72 |
+
self.bits_per_index = embed_dim
|
| 73 |
+
|
| 74 |
+
# we only need to keep the codebook portion up to the group size
|
| 75 |
+
# because we approximate the H loss with this subcode
|
| 76 |
+
group_codes = torch.arange(2 ** self.group_size)
|
| 77 |
+
group_codebook = self.indexes_to_codes(group_codes).float()[:, -group_size:]
|
| 78 |
+
self.register_buffer('group_codebook', group_codebook, persistent=False)
|
| 79 |
+
|
| 80 |
+
self.soft_entropy = soft_entropy # soft_entropy: Sec 3.2 of https://arxiv.org/pdf/1911.05894.pdf
|
| 81 |
+
|
| 82 |
+
def quantize(self, z):
|
| 83 |
+
assert z.shape[-1] == self.embed_dim, f"Expected {self.embed_dim} dimensions, got {z.shape[-1]}"
|
| 84 |
+
|
| 85 |
+
zhat = torch.where(z > 0,
|
| 86 |
+
torch.tensor(1, dtype=z.dtype, device=z.device),
|
| 87 |
+
torch.tensor(-1, dtype=z.dtype, device=z.device))
|
| 88 |
+
return z + (zhat - z).detach()
|
| 89 |
+
|
| 90 |
+
def forward(self, z, collect_metrics=True):
|
| 91 |
+
# if self.input_format == 'bchw':
|
| 92 |
+
# z = rearrange(z, 'b c h w -> b h w c')
|
| 93 |
+
zq = self.quantize(z)
|
| 94 |
+
|
| 95 |
+
q_scale = 1. / (self.embed_dim ** 0.5) if self.l2_norm else 1.
|
| 96 |
+
|
| 97 |
+
zq = zq * q_scale
|
| 98 |
+
|
| 99 |
+
if not collect_metrics:
|
| 100 |
+
return zq, zq.new_zeros(()), {}
|
| 101 |
+
|
| 102 |
+
indices = self.codes_to_indexes(zq.detach())
|
| 103 |
+
group_indices = self.codes_to_group_indexes(zq.detach())
|
| 104 |
+
if not self.training:
|
| 105 |
+
used_codes = torch.unique(indices, return_counts=False)
|
| 106 |
+
else:
|
| 107 |
+
used_codes = None
|
| 108 |
+
|
| 109 |
+
if self.soft_entropy:
|
| 110 |
+
persample_entropy, cb_entropy, avg_prob = self.soft_entropy_loss(z)
|
| 111 |
+
entropy_penalty = self.gamma0 * persample_entropy - self.gamma * cb_entropy
|
| 112 |
+
else:
|
| 113 |
+
zb_by_sample = ((zq + 1) / 2).reshape(z.shape[0], -1, z.shape[-1]).to(torch.float32)
|
| 114 |
+
persample_entropy = self.get_hard_per_sample_entropy(zb_by_sample)
|
| 115 |
+
cb_entropy = codebook_entropy(zq, self.basis, self.embed_dim)
|
| 116 |
+
entropy_penalty = self.gamma0 * persample_entropy - self.gamma * cb_entropy
|
| 117 |
+
|
| 118 |
+
# commit loss
|
| 119 |
+
commit_loss = self.beta * torch.mean(((zq.detach() - z) ** 2).sum(dim=-1))
|
| 120 |
+
|
| 121 |
+
# if self.input_format == 'bchw':
|
| 122 |
+
# zq = rearrange(zq, 'b h w c -> b c h w')
|
| 123 |
+
|
| 124 |
+
return (
|
| 125 |
+
zq,
|
| 126 |
+
commit_loss + self.zeta * entropy_penalty / self.inv_temperature,
|
| 127 |
+
{"H": cb_entropy, "used_codes": used_codes, "indices": indices, "group_indices": group_indices,
|
| 128 |
+
"avg_prob": avg_prob}
|
| 129 |
+
)
|
| 130 |
+
|
| 131 |
+
def soft_entropy_loss(self, z):
|
| 132 |
+
# if we divide the code in subgroups of size group_size, the codebook will be of size 2 ** group_size
|
| 133 |
+
# the sub-code is the last group_size bits of the full code
|
| 134 |
+
group_code_book = self.group_codebook / (self.embed_dim ** 0.5 if self.l2_norm else 1)
|
| 135 |
+
divided_z = rearrange(z, '... (g c) -> ... g c', c=self.group_size)
|
| 136 |
+
|
| 137 |
+
# we calculate the distance between the divided_z and the codebook for each subgroup
|
| 138 |
+
distance = - 2 * torch.einsum('... g c, d c ->... g d', divided_z, group_code_book)
|
| 139 |
+
prob = (-distance * self.inv_temperature).softmax(dim=-1)
|
| 140 |
+
if self.persample_entropy_compute == 'analytical':
|
| 141 |
+
if self.l2_norm:
|
| 142 |
+
p = torch.sigmoid(-4 * z / (self.embed_dim ** 0.5) * self.inv_temperature)
|
| 143 |
+
else:
|
| 144 |
+
p = torch.sigmoid(-4 * z * self.inv_temperature)
|
| 145 |
+
prob = torch.stack([p, 1 - p], dim=-1)
|
| 146 |
+
per_sample_entropy = self.get_entropy(prob, dim=-1, normalize=False).sum(dim=-1).mean()
|
| 147 |
+
else:
|
| 148 |
+
per_sample_entropy = self.get_entropy(prob, dim=-1, normalize=False).sum(dim=-1).mean()
|
| 149 |
+
|
| 150 |
+
# macro average of the probability of each subgroup
|
| 151 |
+
avg_prob = reduce(prob, '... g d ->g d', 'mean')
|
| 152 |
+
codebook_entropy = self.get_entropy(avg_prob, dim=-1, normalize=False)
|
| 153 |
+
|
| 154 |
+
# the approximation of the entropy is the sum of the entropy of each subgroup
|
| 155 |
+
return per_sample_entropy, codebook_entropy.sum(), avg_prob
|
| 156 |
+
|
| 157 |
+
def get_hard_per_sample_entropy(self, zb_by_sample):
|
| 158 |
+
probs_per_dim = zb_by_sample.sum(1) / zb_by_sample.shape[1]
|
| 159 |
+
persample_entropy = - probs_per_dim * torch.log(probs_per_dim + 1e-8) - (1 - probs_per_dim) * torch.log(1 - probs_per_dim + 1e-8)
|
| 160 |
+
persample_entropy = persample_entropy.sum(-1)
|
| 161 |
+
return persample_entropy.mean()
|
| 162 |
+
|
| 163 |
+
def codes_to_indexes(self, zhat):
|
| 164 |
+
"""Converts a `code` to an index in the codebook.
|
| 165 |
+
Args:
|
| 166 |
+
zhat: A tensor of shape (B, ..., C) containing the codes. must be in {-1, 1}
|
| 167 |
+
"""
|
| 168 |
+
assert zhat.shape[-1] == self.embed_dim, f"Expected {self.embed_dim} dimensions, got {zhat.shape[-1]}"
|
| 169 |
+
return ((zhat + 1) / 2 * self.basis).sum(axis=-1).to(torch.int64)
|
| 170 |
+
|
| 171 |
+
def codes_to_group_indexes(self, zhat):
|
| 172 |
+
"""Converts a `code` to a list of indexes (in groups) in the codebook.
|
| 173 |
+
Args:
|
| 174 |
+
zhat: A tensor of shape (B, ..., C) containing the codes. must be in {-1, 1}
|
| 175 |
+
"""
|
| 176 |
+
zhat_in_group = rearrange(zhat, 'b ... (g c) -> b ... g c', c=self.group_size)
|
| 177 |
+
return ((zhat_in_group + 1) / 2 * self.group_basis).sum(axis=-1).to(torch.int64)
|
| 178 |
+
|
| 179 |
+
def indexes_to_codes(self, indices):
|
| 180 |
+
"""Inverse of `indexes_to_codes`."""
|
| 181 |
+
indices = indices.unsqueeze(-1)
|
| 182 |
+
codes_non_centered = torch.remainder(
|
| 183 |
+
torch.floor_divide(indices, self.basis), 2
|
| 184 |
+
)
|
| 185 |
+
return codes_non_centered * 2 - 1
|
| 186 |
+
|
| 187 |
+
def group_indexes_to_codes(self, group_indices):
|
| 188 |
+
"""Inverse of `group_indexes_to_codes`."""
|
| 189 |
+
group_indices = group_indices.unsqueeze(-1)
|
| 190 |
+
codes_non_centered = torch.remainder(
|
| 191 |
+
torch.floor_divide(group_indices, self.group_basis), 2
|
| 192 |
+
)
|
| 193 |
+
codes_non_centered = rearrange(codes_non_centered, 'b ... g c -> b ... (g c)')
|
| 194 |
+
return codes_non_centered * 2 - 1
|
| 195 |
+
|
| 196 |
+
def get_entropy(self, count, dim=-1, eps=1e-4, normalize=True):
|
| 197 |
+
if normalize:
|
| 198 |
+
probs = (count + eps) / (count + eps).sum(dim=dim, keepdim=True)
|
| 199 |
+
else:
|
| 200 |
+
probs = count
|
| 201 |
+
H = -(probs * torch.log(probs + 1e-8)).sum(dim=dim)
|
| 202 |
+
return H
|
| 203 |
+
|
| 204 |
+
def get_group_codebook_entry(self, group_indices):
|
| 205 |
+
z_q = self.group_indexes_to_codes(group_indices)
|
| 206 |
+
q_scale = 1. / (self.embed_dim ** 0.5) if self.l2_norm else 1.
|
| 207 |
+
z_q = z_q * q_scale
|
| 208 |
+
if self.input_format == 'bchw':
|
| 209 |
+
h, w = int(z_q.shape[1] ** 0.5)
|
| 210 |
+
assert h * w == z_q.shape[1], 'Invalid sequence length'
|
| 211 |
+
z_q = rearrange(z_q, 'b (h w) c -> b c h w', h=h)
|
| 212 |
+
return z_q
|
| 213 |
+
|
| 214 |
+
def get_codebook_entry(self, indices):
|
| 215 |
+
z_q = self.indexes_to_codes(indices)
|
| 216 |
+
q_scale = 1. / (self.embed_dim ** 0.5) if self.l2_norm else 1.
|
| 217 |
+
z_q = z_q * q_scale
|
| 218 |
+
if self.input_format == 'bchw':
|
| 219 |
+
h, w = int(z_q.shape[1] ** 0.5)
|
| 220 |
+
assert h * w == z_q.shape[1], 'Invalid sequence length'
|
| 221 |
+
z_q = rearrange(z_q, 'b (h w) c -> b c h w', h=h)
|
| 222 |
+
return z_q
|
| 223 |
+
|
| 224 |
+
|
| 225 |
+
class BSQuantizer(nn.Module):
|
| 226 |
+
|
| 227 |
+
def __init__(self, s1_bits, s2_bits, beta, gamma0, gamma, zeta, group_size):
|
| 228 |
+
super().__init__()
|
| 229 |
+
self.codebook_dim = s1_bits + s2_bits
|
| 230 |
+
self.s1_bits = s1_bits
|
| 231 |
+
self.s2_bits = s2_bits
|
| 232 |
+
self.bsq = BinarySphericalQuantizer(self.codebook_dim, beta, gamma0, gamma, zeta, group_size=group_size)
|
| 233 |
+
|
| 234 |
+
def bits_to_indices(self, bits):
|
| 235 |
+
bits = (bits >= 0).to(torch.long)
|
| 236 |
+
indices = 2 ** torch.arange(
|
| 237 |
+
0,
|
| 238 |
+
bits.shape[-1],
|
| 239 |
+
1,
|
| 240 |
+
dtype=torch.long,
|
| 241 |
+
device=bits.device,
|
| 242 |
+
)
|
| 243 |
+
return (bits * indices).sum(-1)
|
| 244 |
+
|
| 245 |
+
def forward(self, z, half=False, collect_metrics=True):
|
| 246 |
+
z = F.normalize(z, dim=-1)
|
| 247 |
+
quantized, bsq_loss, metrics = self.bsq(z, collect_metrics=collect_metrics)
|
| 248 |
+
if half:
|
| 249 |
+
q_pre = quantized[:, :, :self.s1_bits]
|
| 250 |
+
q_post = quantized[:, :, self.s1_bits:]
|
| 251 |
+
z_indices = [self.bits_to_indices(q_pre), self.bits_to_indices(q_post)]
|
| 252 |
+
else:
|
| 253 |
+
z_indices = self.bits_to_indices(quantized)
|
| 254 |
+
return bsq_loss, quantized, z_indices
|
| 255 |
+
|
| 256 |
+
|
| 257 |
+
class RMSNorm(torch.nn.Module):
|
| 258 |
+
def __init__(self, dim: int, eps: float = 1e-5):
|
| 259 |
+
super().__init__()
|
| 260 |
+
self.eps = eps
|
| 261 |
+
self.weight = nn.Parameter(torch.ones(dim))
|
| 262 |
+
|
| 263 |
+
def _norm(self, x):
|
| 264 |
+
return x * torch.rsqrt(torch.mean(x * x, dim=-1, keepdim=True) + self.eps)
|
| 265 |
+
|
| 266 |
+
def forward(self, x):
|
| 267 |
+
output = self._norm(x.float()).type_as(x)
|
| 268 |
+
return output * self.weight
|
| 269 |
+
|
| 270 |
+
|
| 271 |
+
class FeedForward(nn.Module):
|
| 272 |
+
def __init__(self, d_model, ff_dim, ffn_dropout_p=0.0):
|
| 273 |
+
super().__init__()
|
| 274 |
+
|
| 275 |
+
self.w1 = nn.Linear(d_model, ff_dim, bias=False)
|
| 276 |
+
self.w3 = nn.Linear(d_model, ff_dim, bias=False)
|
| 277 |
+
self.w2 = nn.Linear(ff_dim, d_model, bias=False)
|
| 278 |
+
self.ffn_dropout = nn.Dropout(ffn_dropout_p)
|
| 279 |
+
|
| 280 |
+
def forward(self, x):
|
| 281 |
+
return self.ffn_dropout(self.w2(F.silu(self.w1(x)) * self.w3(x)))
|
| 282 |
+
|
| 283 |
+
|
| 284 |
+
class RotaryPositionalEmbedding(nn.Module):
|
| 285 |
+
def __init__(self, dim):
|
| 286 |
+
super().__init__()
|
| 287 |
+
inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2).float() / dim))
|
| 288 |
+
self.register_buffer("inv_freq", inv_freq)
|
| 289 |
+
self.seq_len_cached = None
|
| 290 |
+
self.cos_cached = None
|
| 291 |
+
self.sin_cached = None
|
| 292 |
+
|
| 293 |
+
def _update_cos_sin_cache(self, x, seq_len):
|
| 294 |
+
if seq_len != self.seq_len_cached:
|
| 295 |
+
self.seq_len_cached = seq_len
|
| 296 |
+
t = torch.arange(seq_len, device=x.device).type_as(self.inv_freq)
|
| 297 |
+
freqs = torch.einsum('i,j->ij', t, self.inv_freq)
|
| 298 |
+
emb = torch.cat((freqs, freqs), dim=-1).to(x.device)
|
| 299 |
+
self.cos_cached = emb.cos()[None, None, :, :]
|
| 300 |
+
self.sin_cached = emb.sin()[None, None, :, :]
|
| 301 |
+
return self.cos_cached, self.sin_cached
|
| 302 |
+
|
| 303 |
+
def forward(self, q, k):
|
| 304 |
+
cos, sin = self._update_cos_sin_cache(q, q.shape[-2])
|
| 305 |
+
return (
|
| 306 |
+
(q * cos) + (self._rotate_half(q) * sin),
|
| 307 |
+
(k * cos) + (self._rotate_half(k) * sin),
|
| 308 |
+
)
|
| 309 |
+
|
| 310 |
+
def _rotate_half(self, x):
|
| 311 |
+
x1, x2 = x.chunk(2, dim=-1)
|
| 312 |
+
return torch.cat((-x2, x1), dim=-1)
|
| 313 |
+
|
| 314 |
+
|
| 315 |
+
class MultiHeadAttentionWithRoPE(nn.Module):
|
| 316 |
+
def __init__(self, d_model, n_heads, attn_dropout_p=0.0, resid_dropout_p=0.0):
|
| 317 |
+
super().__init__()
|
| 318 |
+
self.d_model = d_model
|
| 319 |
+
self.n_heads = n_heads
|
| 320 |
+
self.head_dim = d_model // n_heads
|
| 321 |
+
|
| 322 |
+
self.q_proj = nn.Linear(d_model, d_model)
|
| 323 |
+
self.k_proj = nn.Linear(d_model, d_model)
|
| 324 |
+
self.v_proj = nn.Linear(d_model, d_model)
|
| 325 |
+
self.out_proj = nn.Linear(d_model, d_model)
|
| 326 |
+
self.rotary = RotaryPositionalEmbedding(self.head_dim)
|
| 327 |
+
self.attn_dropout_p = attn_dropout_p
|
| 328 |
+
self.resid_dropout = nn.Dropout(resid_dropout_p)
|
| 329 |
+
|
| 330 |
+
def forward(self, x, key_padding_mask=None):
|
| 331 |
+
batch_size, seq_len, _ = x.shape
|
| 332 |
+
|
| 333 |
+
q = self.q_proj(x).view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2)
|
| 334 |
+
k = self.k_proj(x).view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2)
|
| 335 |
+
v = self.v_proj(x).view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2)
|
| 336 |
+
|
| 337 |
+
q, k = self.rotary(q, k)
|
| 338 |
+
|
| 339 |
+
if key_padding_mask is not None:
|
| 340 |
+
attn_mask = key_padding_mask.unsqueeze(1).unsqueeze(2) # [batch, 1, 1, seq_len]
|
| 341 |
+
attn_mask = attn_mask.expand(-1, self.n_heads, seq_len, -1) # [batch, n_heads, q_len, k_len]
|
| 342 |
+
else:
|
| 343 |
+
attn_mask = None
|
| 344 |
+
|
| 345 |
+
attn_output = F.scaled_dot_product_attention(
|
| 346 |
+
q, k, v,
|
| 347 |
+
attn_mask=attn_mask,
|
| 348 |
+
dropout_p=self.attn_dropout_p if self.training else 0.0,
|
| 349 |
+
is_causal=True
|
| 350 |
+
)
|
| 351 |
+
|
| 352 |
+
attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model)
|
| 353 |
+
return self.resid_dropout(self.out_proj(attn_output))
|
| 354 |
+
|
| 355 |
+
|
| 356 |
+
class MultiHeadCrossAttentionWithRoPE(nn.Module):
|
| 357 |
+
def __init__(self, d_model, n_heads, attn_dropout_p=0.0, resid_dropout=0.0):
|
| 358 |
+
super().__init__()
|
| 359 |
+
self.d_model = d_model
|
| 360 |
+
self.n_heads = n_heads
|
| 361 |
+
self.head_dim = d_model // n_heads
|
| 362 |
+
|
| 363 |
+
self.q_proj = nn.Linear(d_model, d_model)
|
| 364 |
+
self.k_proj = nn.Linear(d_model, d_model)
|
| 365 |
+
self.v_proj = nn.Linear(d_model, d_model)
|
| 366 |
+
self.out_proj = nn.Linear(d_model, d_model)
|
| 367 |
+
self.rotary = RotaryPositionalEmbedding(self.head_dim)
|
| 368 |
+
self.attn_dropout_p = attn_dropout_p
|
| 369 |
+
self.resid_dropout = nn.Dropout(resid_dropout)
|
| 370 |
+
|
| 371 |
+
def forward(self, query, key, value, key_padding_mask=None):
|
| 372 |
+
batch_size, q_len, _ = query.shape
|
| 373 |
+
_, seq_len, _ = key.shape
|
| 374 |
+
|
| 375 |
+
q = self.q_proj(query).view(batch_size, q_len, self.n_heads, self.head_dim).transpose(1, 2)
|
| 376 |
+
k = self.k_proj(key).view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2)
|
| 377 |
+
v = self.v_proj(value).view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2)
|
| 378 |
+
|
| 379 |
+
q, k = self.rotary(q, k)
|
| 380 |
+
|
| 381 |
+
if key_padding_mask is not None:
|
| 382 |
+
attn_mask = key_padding_mask.unsqueeze(1).unsqueeze(2)
|
| 383 |
+
attn_mask = attn_mask.expand(-1, self.n_heads, q_len, -1)
|
| 384 |
+
else:
|
| 385 |
+
attn_mask = None
|
| 386 |
+
|
| 387 |
+
is_causal_flag = self.training
|
| 388 |
+
|
| 389 |
+
attn_output = F.scaled_dot_product_attention(
|
| 390 |
+
q, k, v,
|
| 391 |
+
attn_mask=attn_mask,
|
| 392 |
+
dropout_p=self.attn_dropout_p if self.training else 0.0,
|
| 393 |
+
is_causal=is_causal_flag
|
| 394 |
+
)
|
| 395 |
+
|
| 396 |
+
attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, q_len, self.d_model)
|
| 397 |
+
return self.resid_dropout(self.out_proj(attn_output))
|
| 398 |
+
|
| 399 |
+
|
| 400 |
+
class HierarchicalEmbedding(nn.Module):
|
| 401 |
+
def __init__(self, s1_bits, s2_bits, d_model=256):
|
| 402 |
+
super().__init__()
|
| 403 |
+
self.s1_bits = s1_bits
|
| 404 |
+
self.s2_bits = s2_bits
|
| 405 |
+
|
| 406 |
+
vocab_s1 = 2 ** s1_bits
|
| 407 |
+
vocab_s2 = 2 ** s2_bits
|
| 408 |
+
|
| 409 |
+
self.emb_s1 = nn.Embedding(vocab_s1, d_model)
|
| 410 |
+
self.emb_s2 = nn.Embedding(vocab_s2, d_model)
|
| 411 |
+
self.d_model = d_model
|
| 412 |
+
self.fusion_proj = nn.Linear(d_model * 2, d_model)
|
| 413 |
+
|
| 414 |
+
nn.init.normal_(self.emb_s1.weight, mean=0, std=d_model ** -0.5)
|
| 415 |
+
nn.init.normal_(self.emb_s2.weight, mean=0, std=d_model ** -0.5)
|
| 416 |
+
|
| 417 |
+
def split_token(self, token_ids: torch.Tensor, s2_bits: int):
|
| 418 |
+
"""Inputs:
|
| 419 |
+
token_ids (torch.Tensor): Composite token IDs of shape [batch_size, seq_len] or [N], each in range [0, 2^(s1_bits + s2_bits) - 1].
|
| 420 |
+
s2_bits (int): Number of low bits used for the fine token (s2).
|
| 421 |
+
"""
|
| 422 |
+
assert isinstance(s2_bits, int) and s2_bits > 0, "s2_bits must be a positive integer"
|
| 423 |
+
|
| 424 |
+
t = token_ids.long()
|
| 425 |
+
mask = (1 << s2_bits) - 1
|
| 426 |
+
s2_ids = t & mask # extract low bits
|
| 427 |
+
s1_ids = t >> s2_bits # extract high bits
|
| 428 |
+
return s1_ids, s2_ids
|
| 429 |
+
|
| 430 |
+
def forward(self, token_ids):
|
| 431 |
+
"""Inputs:
|
| 432 |
+
token_ids:
|
| 433 |
+
- tuple or list: (s1_ids, s2_ids), each of shape [batch_size, seq_len], or
|
| 434 |
+
- torch.Tensor: composite token IDs of shape [batch_size, seq_len], which will be split into (s1_ids, s2_ids) internally.
|
| 435 |
+
Output: [batch_size, seq_len, d_model]
|
| 436 |
+
"""
|
| 437 |
+
if isinstance(token_ids, tuple) or isinstance(token_ids, list):
|
| 438 |
+
s1_ids, s2_ids = token_ids
|
| 439 |
+
else:
|
| 440 |
+
s1_ids, s2_ids = self.split_token(token_ids, self.s2_bits)
|
| 441 |
+
s1_emb = self.emb_s1(s1_ids) * math.sqrt(self.d_model)
|
| 442 |
+
s2_emb = self.emb_s2(s2_ids) * math.sqrt(self.d_model)
|
| 443 |
+
return self.fusion_proj(torch.cat([s1_emb, s2_emb], dim=-1))
|
| 444 |
+
|
| 445 |
+
|
| 446 |
+
class DependencyAwareLayer(nn.Module):
|
| 447 |
+
def __init__(self, d_model, n_heads=4, attn_dropout_p=0.0, resid_dropout=0.0):
|
| 448 |
+
super().__init__()
|
| 449 |
+
self.cross_attn = MultiHeadCrossAttentionWithRoPE(d_model, n_heads, attn_dropout_p, resid_dropout)
|
| 450 |
+
self.norm = RMSNorm(d_model)
|
| 451 |
+
|
| 452 |
+
def forward(self, hidden_states, sibling_embed, key_padding_mask=None):
|
| 453 |
+
"""hidden_states: [batch, seq_len, d_model]
|
| 454 |
+
sibling_embed: Embedding from another subtoken
|
| 455 |
+
"""
|
| 456 |
+
attn_out = self.cross_attn(
|
| 457 |
+
query=sibling_embed,
|
| 458 |
+
key=hidden_states,
|
| 459 |
+
value=hidden_states,
|
| 460 |
+
key_padding_mask=key_padding_mask
|
| 461 |
+
)
|
| 462 |
+
return self.norm(hidden_states + attn_out)
|
| 463 |
+
|
| 464 |
+
|
| 465 |
+
class TransformerBlock(nn.Module):
|
| 466 |
+
def __init__(self, d_model, n_heads, ff_dim=1024, ffn_dropout_p=0.0, attn_dropout_p=0.0, resid_dropout_p=0.0):
|
| 467 |
+
super().__init__()
|
| 468 |
+
self.norm1 = RMSNorm(d_model)
|
| 469 |
+
self.self_attn = MultiHeadAttentionWithRoPE(d_model, n_heads, attn_dropout_p, resid_dropout_p)
|
| 470 |
+
self.norm2 = RMSNorm(d_model)
|
| 471 |
+
self.ffn = FeedForward(d_model, ff_dim, ffn_dropout_p)
|
| 472 |
+
|
| 473 |
+
def forward(self, x, key_padding_mask=None):
|
| 474 |
+
residual = x
|
| 475 |
+
x = self.norm1(x)
|
| 476 |
+
attn_out = self.self_attn(x, key_padding_mask=key_padding_mask)
|
| 477 |
+
x = residual + attn_out
|
| 478 |
+
|
| 479 |
+
residual = x
|
| 480 |
+
x = self.norm2(x)
|
| 481 |
+
ffn_out = self.ffn(x)
|
| 482 |
+
x = residual + ffn_out
|
| 483 |
+
return x
|
| 484 |
+
|
| 485 |
+
|
| 486 |
+
class DualHead(nn.Module):
|
| 487 |
+
def __init__(self, s1_bits, s2_bits, d_model):
|
| 488 |
+
super().__init__()
|
| 489 |
+
self.vocab_s1 = 2 ** s1_bits
|
| 490 |
+
self.vocab_s2 = 2 ** s2_bits
|
| 491 |
+
self.proj_s1 = nn.Linear(d_model, self.vocab_s1)
|
| 492 |
+
self.proj_s2 = nn.Linear(d_model, self.vocab_s2)
|
| 493 |
+
|
| 494 |
+
def compute_loss(self, s1_logits, s2_logits, s1_targets, s2_targets, padding_mask=None):
|
| 495 |
+
if padding_mask is not None:
|
| 496 |
+
valid_mask = (padding_mask == 0)
|
| 497 |
+
s1_logits = s1_logits[valid_mask]
|
| 498 |
+
s2_logits = s2_logits[valid_mask]
|
| 499 |
+
s1_targets = s1_targets[valid_mask]
|
| 500 |
+
s2_targets = s2_targets[valid_mask]
|
| 501 |
+
ce_s1 = F.cross_entropy(s1_logits, s1_targets)
|
| 502 |
+
ce_s2 = F.cross_entropy(s2_logits, s2_targets)
|
| 503 |
+
else:
|
| 504 |
+
ce_s1 = F.cross_entropy(s1_logits.reshape(-1, self.vocab_s1), s1_targets.reshape(-1))
|
| 505 |
+
ce_s2 = F.cross_entropy(s2_logits.reshape(-1, self.vocab_s2), s2_targets.reshape(-1))
|
| 506 |
+
ce_loss = (ce_s1 + ce_s2) / 2
|
| 507 |
+
return ce_loss, ce_s1, ce_s2
|
| 508 |
+
|
| 509 |
+
def forward(self, x):
|
| 510 |
+
return self.proj_s1(x)
|
| 511 |
+
|
| 512 |
+
def cond_forward(self, x2):
|
| 513 |
+
return self.proj_s2(x2)
|
| 514 |
+
|
| 515 |
+
|
| 516 |
+
class FixedEmbedding(nn.Module):
|
| 517 |
+
def __init__(self, c_in, d_model):
|
| 518 |
+
super(FixedEmbedding, self).__init__()
|
| 519 |
+
|
| 520 |
+
w = torch.zeros(c_in, d_model).float()
|
| 521 |
+
w.require_grad = False
|
| 522 |
+
|
| 523 |
+
position = torch.arange(0, c_in).float().unsqueeze(1)
|
| 524 |
+
div_term = (torch.arange(0, d_model, 2).float() * -(math.log(10000.0) / d_model)).exp()
|
| 525 |
+
|
| 526 |
+
w[:, 0::2] = torch.sin(position * div_term)
|
| 527 |
+
w[:, 1::2] = torch.cos(position * div_term)
|
| 528 |
+
|
| 529 |
+
self.emb = nn.Embedding(c_in, d_model)
|
| 530 |
+
self.emb.weight = nn.Parameter(w, requires_grad=False)
|
| 531 |
+
|
| 532 |
+
def forward(self, x):
|
| 533 |
+
return self.emb(x).detach()
|
| 534 |
+
|
| 535 |
+
|
| 536 |
+
class TemporalEmbedding(nn.Module):
|
| 537 |
+
def __init__(self, d_model, learn_pe):
|
| 538 |
+
super(TemporalEmbedding, self).__init__()
|
| 539 |
+
|
| 540 |
+
minute_size = 60
|
| 541 |
+
hour_size = 24
|
| 542 |
+
weekday_size = 7
|
| 543 |
+
day_size = 32
|
| 544 |
+
month_size = 13
|
| 545 |
+
|
| 546 |
+
Embed = FixedEmbedding if not learn_pe else nn.Embedding
|
| 547 |
+
self.minute_embed = Embed(minute_size, d_model)
|
| 548 |
+
self.hour_embed = Embed(hour_size, d_model)
|
| 549 |
+
self.weekday_embed = Embed(weekday_size, d_model)
|
| 550 |
+
self.day_embed = Embed(day_size, d_model)
|
| 551 |
+
self.month_embed = Embed(month_size, d_model)
|
| 552 |
+
|
| 553 |
+
def forward(self, x):
|
| 554 |
+
x = x.long()
|
| 555 |
+
|
| 556 |
+
minute_x = self.minute_embed(x[:, :, 0])
|
| 557 |
+
hour_x = self.hour_embed(x[:, :, 1])
|
| 558 |
+
weekday_x = self.weekday_embed(x[:, :, 2])
|
| 559 |
+
day_x = self.day_embed(x[:, :, 3])
|
| 560 |
+
month_x = self.month_embed(x[:, :, 4])
|
| 561 |
+
|
| 562 |
+
return hour_x + weekday_x + day_x + month_x + minute_x
|
| 563 |
+
|
| 564 |
+
|
| 565 |
+
|
| 566 |
+
|
| 567 |
+
|
| 568 |
+
|
| 569 |
+
|
| 570 |
+
|
roughness_lab/calibrate.py
ADDED
|
@@ -0,0 +1,253 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Texture calibration of Kronos sampling parameters via surface roughness.
|
| 2 |
+
|
| 3 |
+
For each asset and each sampling configuration (temperature T x top_p), the
|
| 4 |
+
script forecasts several historical windows with Kronos-small, computes
|
| 5 |
+
surface-roughness parameters (Ra, Rq, Rz, RSm, Rsk, Rku, Wa) on the forecast
|
| 6 |
+
paths and on the realized continuation, and scores how well the forecast
|
| 7 |
+
*texture* matches reality. Output: results CSV, score heatmaps, a multi-scale
|
| 8 |
+
roughness fingerprint, a decomposition teaching figure, and report.md.
|
| 9 |
+
|
| 10 |
+
Run: python roughness_lab/calibrate.py [--windows 6] [--horizon 64] [--paths 4]
|
| 11 |
+
"""
|
| 12 |
+
import argparse
|
| 13 |
+
import json
|
| 14 |
+
import sys
|
| 15 |
+
import time
|
| 16 |
+
from pathlib import Path
|
| 17 |
+
|
| 18 |
+
import matplotlib
|
| 19 |
+
|
| 20 |
+
matplotlib.use("Agg")
|
| 21 |
+
import matplotlib.pyplot as plt
|
| 22 |
+
import numpy as np
|
| 23 |
+
import pandas as pd
|
| 24 |
+
import torch
|
| 25 |
+
|
| 26 |
+
LAB = Path(__file__).resolve().parent
|
| 27 |
+
sys.path.insert(0, str(LAB))
|
| 28 |
+
sys.path.insert(0, str(LAB.parent / "Kronos"))
|
| 29 |
+
|
| 30 |
+
from roughness import fingerprint, gaussian_filter, roughness_params, texture_error
|
| 31 |
+
from model import Kronos, KronosTokenizer, KronosPredictor
|
| 32 |
+
|
| 33 |
+
GRID_T = [0.7, 1.0, 1.3]
|
| 34 |
+
GRID_P = [0.8, 0.9, 1.0]
|
| 35 |
+
CUTOFF_MATCH = 8 # roughness cutoff (bars) for in-horizon texture params
|
| 36 |
+
CUTOFF_WA = 16 # cutoff for the in-horizon Wa (waviness amplitude)
|
| 37 |
+
FP_CUTOFFS = [8, 24, 72, 168] # multi-scale fingerprint on full history
|
| 38 |
+
CONTEXT = 400
|
| 39 |
+
PARAM_KEYS = ["ra", "rq", "rz", "rsm", "rsk", "rku", "wa"]
|
| 40 |
+
|
| 41 |
+
|
| 42 |
+
def params_for(close: np.ndarray) -> dict:
|
| 43 |
+
p = roughness_params(close, CUTOFF_MATCH).as_dict()
|
| 44 |
+
p["wa"] = roughness_params(close, CUTOFF_WA).wa
|
| 45 |
+
return p
|
| 46 |
+
|
| 47 |
+
|
| 48 |
+
def window_anchors(n: int, n_windows: int, horizon: int) -> np.ndarray:
|
| 49 |
+
lo = max(CONTEXT, int(n * 0.4))
|
| 50 |
+
hi = n - horizon - 1
|
| 51 |
+
return np.linspace(lo, hi, n_windows).astype(int)
|
| 52 |
+
|
| 53 |
+
|
| 54 |
+
def history_fingerprint(close: np.ndarray) -> dict:
|
| 55 |
+
"""Median fingerprint over non-overlapping windows of 5*cutoff bars."""
|
| 56 |
+
out = {}
|
| 57 |
+
for c in FP_CUTOFFS:
|
| 58 |
+
w = 5 * c
|
| 59 |
+
vals = [roughness_params(close[i:i + w], c).as_dict()
|
| 60 |
+
for i in range(0, len(close) - w, w)]
|
| 61 |
+
out[c] = {k: float(np.median([v[k] for v in vals])) for k in PARAM_KEYS}
|
| 62 |
+
return out
|
| 63 |
+
|
| 64 |
+
|
| 65 |
+
def run_asset(name: str, df: pd.DataFrame, predictor, args, results_dir: Path) -> dict:
|
| 66 |
+
n = len(df)
|
| 67 |
+
anchors = window_anchors(n, args.windows, args.horizon)
|
| 68 |
+
print(f"\n=== {name}: {n} bars, {args.windows} windows at {list(anchors)} ===", flush=True)
|
| 69 |
+
|
| 70 |
+
feat_cols = ["open", "high", "low", "close", "volume", "amount"]
|
| 71 |
+
ctx_dfs, x_tss, y_tss, real_params, real_closes = [], [], [], [], []
|
| 72 |
+
for a in anchors:
|
| 73 |
+
ctx = df.iloc[a - CONTEXT:a]
|
| 74 |
+
real = df.iloc[a:a + args.horizon]
|
| 75 |
+
ctx_dfs.append(ctx[feat_cols].reset_index(drop=True))
|
| 76 |
+
x_tss.append(ctx["timestamps"].reset_index(drop=True))
|
| 77 |
+
y_tss.append(real["timestamps"].reset_index(drop=True))
|
| 78 |
+
real_closes.append(real["close"].to_numpy())
|
| 79 |
+
real_params.append(params_for(real["close"].to_numpy()))
|
| 80 |
+
|
| 81 |
+
rows, config_scores, best = [], {}, None
|
| 82 |
+
for T in GRID_T:
|
| 83 |
+
for top_p in GRID_P:
|
| 84 |
+
torch.manual_seed(int(T * 1000) * 7919 + int(top_p * 1000))
|
| 85 |
+
t0 = time.time()
|
| 86 |
+
# one batched call: every window replicated for every path
|
| 87 |
+
preds = predictor.predict_batch(
|
| 88 |
+
df_list=[c for c in ctx_dfs for _ in range(args.paths)],
|
| 89 |
+
x_timestamp_list=[x for x in x_tss for _ in range(args.paths)],
|
| 90 |
+
y_timestamp_list=[y for y in y_tss for _ in range(args.paths)],
|
| 91 |
+
pred_len=args.horizon, T=T, top_p=top_p, sample_count=1, verbose=False,
|
| 92 |
+
)
|
| 93 |
+
dt = time.time() - t0
|
| 94 |
+
|
| 95 |
+
win_errors = []
|
| 96 |
+
for wi in range(args.windows):
|
| 97 |
+
paths = preds[wi * args.paths:(wi + 1) * args.paths]
|
| 98 |
+
path_params = [params_for(p["close"].to_numpy()) for p in paths]
|
| 99 |
+
pred_med = {k: float(np.median([pp[k] for pp in path_params])) for k in PARAM_KEYS}
|
| 100 |
+
err = texture_error(pred_med, real_params[wi])
|
| 101 |
+
win_errors.append(err)
|
| 102 |
+
rows.append({"asset": name, "T": T, "top_p": top_p, "window": wi,
|
| 103 |
+
"anchor": int(anchors[wi]), "texture_error": err,
|
| 104 |
+
**{f"pred_{k}": pred_med[k] for k in PARAM_KEYS},
|
| 105 |
+
**{f"real_{k}": real_params[wi][k] for k in PARAM_KEYS}})
|
| 106 |
+
score = float(np.mean(win_errors))
|
| 107 |
+
config_scores[(T, top_p)] = score
|
| 108 |
+
if best is None or score < best["score"]:
|
| 109 |
+
best = {"T": T, "top_p": top_p, "score": score,
|
| 110 |
+
"paths_close": [p["close"].to_numpy() for p in preds]}
|
| 111 |
+
print(f" T={T:.1f} top_p={top_p:.2f} texture_error={score:.4f} ({dt:.0f}s)", flush=True)
|
| 112 |
+
|
| 113 |
+
pd.DataFrame(rows).to_csv(results_dir / f"sweep_{name}.csv", index=False)
|
| 114 |
+
|
| 115 |
+
# ---- score heatmap ----
|
| 116 |
+
grid = np.array([[config_scores[(T, p)] for p in GRID_P] for T in GRID_T])
|
| 117 |
+
fig, ax = plt.subplots(figsize=(5.2, 4.2))
|
| 118 |
+
im = ax.imshow(grid, cmap="viridis_r")
|
| 119 |
+
ax.set_xticks(range(len(GRID_P)), [f"{p:.2f}" for p in GRID_P])
|
| 120 |
+
ax.set_yticks(range(len(GRID_T)), [f"{t:.1f}" for t in GRID_T])
|
| 121 |
+
ax.set_xlabel("top_p"); ax.set_ylabel("temperature T")
|
| 122 |
+
ax.set_title(f"{name} — texture mismatch (lower = more realistic)")
|
| 123 |
+
for i in range(len(GRID_T)):
|
| 124 |
+
for j in range(len(GRID_P)):
|
| 125 |
+
ax.text(j, i, f"{grid[i, j]:.3f}", ha="center", va="center",
|
| 126 |
+
color="white" if grid[i, j] > grid.mean() else "black", fontsize=9)
|
| 127 |
+
fig.colorbar(im, shrink=0.85)
|
| 128 |
+
fig.tight_layout()
|
| 129 |
+
fig.savefig(results_dir / f"heatmap_{name}.png", dpi=150)
|
| 130 |
+
plt.close(fig)
|
| 131 |
+
|
| 132 |
+
# ---- multi-scale fingerprint: history vs best-config forecasts ----
|
| 133 |
+
hist_fp = history_fingerprint(df["close"].to_numpy())
|
| 134 |
+
fc_ra = {c: float(np.median([roughness_params(pc, c).ra for pc in best["paths_close"]]))
|
| 135 |
+
for c in (CUTOFF_MATCH, CUTOFF_WA)}
|
| 136 |
+
real_ra = {c: float(np.median([roughness_params(rc, c).ra for rc in real_closes]))
|
| 137 |
+
for c in (CUTOFF_MATCH, CUTOFF_WA)}
|
| 138 |
+
fig, ax = plt.subplots(figsize=(6, 4.2))
|
| 139 |
+
ax.plot(FP_CUTOFFS, [hist_fp[c]["ra"] for c in FP_CUTOFFS], "o-", label="history (full)")
|
| 140 |
+
ax.plot(list(fc_ra), list(fc_ra.values()), "s--", label=f"forecast (T={best['T']}, top_p={best['top_p']})")
|
| 141 |
+
ax.plot(list(real_ra), list(real_ra.values()), "^:", label="realized (eval windows)")
|
| 142 |
+
ax.set_xscale("log"); ax.set_yscale("log")
|
| 143 |
+
ax.set_xlabel("cutoff wavelength λc (bars)"); ax.set_ylabel("Ra (%)")
|
| 144 |
+
ax.set_title(f"{name} — multi-scale roughness fingerprint")
|
| 145 |
+
ax.grid(alpha=.3, which="both"); ax.legend(fontsize=9)
|
| 146 |
+
fig.tight_layout()
|
| 147 |
+
fig.savefig(results_dir / f"fingerprint_{name}.png", dpi=150)
|
| 148 |
+
plt.close(fig)
|
| 149 |
+
|
| 150 |
+
# ---- decomposition teaching figure ----
|
| 151 |
+
tail = df.iloc[-800:]
|
| 152 |
+
z = 100 * np.log(tail["close"].to_numpy())
|
| 153 |
+
w = gaussian_filter(z, 24)
|
| 154 |
+
fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(9, 5.6), sharex=True,
|
| 155 |
+
gridspec_kw={"height_ratios": [2, 1]})
|
| 156 |
+
ax1.plot(tail["timestamps"], np.exp(z / 100), lw=.7, label="close")
|
| 157 |
+
ax1.plot(tail["timestamps"], np.exp(w / 100), lw=1.6, label="waviness (λc=24 bars)")
|
| 158 |
+
ax1.set_title(f"{name} — price as a surface profile")
|
| 159 |
+
ax1.legend(fontsize=9); ax1.grid(alpha=.3)
|
| 160 |
+
ax2.plot(tail["timestamps"], z - w, lw=.7, color="#8b5cf6")
|
| 161 |
+
ax2.axhline(0, color="gray", lw=.5)
|
| 162 |
+
ax2.set_ylabel("roughness (%)"); ax2.grid(alpha=.3)
|
| 163 |
+
fig.tight_layout()
|
| 164 |
+
fig.savefig(results_dir / f"decomposition_{name}.png", dpi=150)
|
| 165 |
+
plt.close(fig)
|
| 166 |
+
|
| 167 |
+
med = lambda key, sel: float(np.median([r[key] for r in rows
|
| 168 |
+
if r["T"] == sel[0] and r["top_p"] == sel[1]]))
|
| 169 |
+
return {
|
| 170 |
+
"asset": name,
|
| 171 |
+
"bars": n,
|
| 172 |
+
"best_T": best["T"],
|
| 173 |
+
"best_top_p": best["top_p"],
|
| 174 |
+
"best_score": best["score"],
|
| 175 |
+
"scores": {f"T{T}_p{p}": s for (T, p), s in config_scores.items()},
|
| 176 |
+
"history_fingerprint": hist_fp,
|
| 177 |
+
"best_pred_vs_real": {k: {"pred": med(f"pred_{k}", (best["T"], best["top_p"])),
|
| 178 |
+
"real": med(f"real_{k}", (best["T"], best["top_p"]))}
|
| 179 |
+
for k in PARAM_KEYS},
|
| 180 |
+
}
|
| 181 |
+
|
| 182 |
+
|
| 183 |
+
def write_report(summaries: list, args, results_dir: Path) -> None:
|
| 184 |
+
L = []
|
| 185 |
+
L.append("# Kronos texture calibration via surface-roughness analysis\n")
|
| 186 |
+
L.append("Price is treated as a measured surface profile: an ISO 16610-21 Gaussian filter "
|
| 187 |
+
"splits log-price into **waviness** (trend, cutoff λc) and **roughness** (texture). "
|
| 188 |
+
"Parameters: **Ra/Rq** mean/RMS roughness amplitude (%), **Rz** mean peak-to-valley (%), "
|
| 189 |
+
"**RSm** mean wiggle spacing (bars), **Rsk/Rku** texture skew/kurtosis, **Wa** waviness "
|
| 190 |
+
"amplitude (%). Forecast paths from Kronos-small were scored by how closely their "
|
| 191 |
+
"texture matches the realized continuation (mean |log-ratio| across parameters, "
|
| 192 |
+
"|difference| for Rsk — lower is better).\n")
|
| 193 |
+
L.append(f"Setup: context {CONTEXT} bars, horizon {args.horizon} bars, {args.paths} sampled "
|
| 194 |
+
f"paths x {args.windows} windows per configuration, in-horizon cutoffs "
|
| 195 |
+
f"λc={CUTOFF_MATCH} (roughness) / {CUTOFF_WA} (Wa). Model: Kronos-small (CPU).\n")
|
| 196 |
+
for s in summaries:
|
| 197 |
+
L.append(f"\n## {s['asset']} ({s['bars']} bars)\n")
|
| 198 |
+
L.append(f"**Best sampling configuration: T={s['best_T']}, top_p={s['best_top_p']}** "
|
| 199 |
+
f"(texture error {s['best_score']:.4f})\n")
|
| 200 |
+
L.append("\n| config | texture error |\n|---|---|")
|
| 201 |
+
for k, v in sorted(s["scores"].items(), key=lambda kv: kv[1]):
|
| 202 |
+
L.append(f"| {k} | {v:.4f} |")
|
| 203 |
+
L.append("\n**Best-config forecast texture vs realized** (medians over windows):\n")
|
| 204 |
+
L.append("| param | forecast | realized |\n|---|---|---|")
|
| 205 |
+
for k, pv in s["best_pred_vs_real"].items():
|
| 206 |
+
L.append(f"| {k.upper()} | {pv['pred']:.4g} | {pv['real']:.4g} |")
|
| 207 |
+
L.append("\n**Multi-scale fingerprint of history** (medians, full series):\n")
|
| 208 |
+
L.append("| λc (bars) | " + " | ".join(k.upper() for k in PARAM_KEYS) + " |")
|
| 209 |
+
L.append("|---|" + "---|" * len(PARAM_KEYS))
|
| 210 |
+
for c, vals in s["history_fingerprint"].items():
|
| 211 |
+
L.append(f"| {c} | " + " | ".join(f"{vals[k]:.3g}" for k in PARAM_KEYS) + " |")
|
| 212 |
+
L.append(f"\n\n")
|
| 213 |
+
L.append(f"\n")
|
| 214 |
+
L.append(f"\n")
|
| 215 |
+
L.append("\n## Caveats\n")
|
| 216 |
+
L.append("- Texture match says forecasts *look statistically like* the market, not that they "
|
| 217 |
+
"predict direction; it complements (not replaces) error metrics like MAE.\n"
|
| 218 |
+
"- AAPL bars exist only during trading sessions, so a 'bar' wavelength is trading "
|
| 219 |
+
"time, not wall-clock time.\n"
|
| 220 |
+
"- Windows/paths are modest because everything ran on CPU; the GPU pipeline in "
|
| 221 |
+
"`gpu_finetune/` scales this up and applies the same scoring to checkpoint selection.\n")
|
| 222 |
+
(results_dir / "report.md").write_text("\n".join(L), encoding="utf-8")
|
| 223 |
+
|
| 224 |
+
|
| 225 |
+
def main() -> None:
|
| 226 |
+
ap = argparse.ArgumentParser()
|
| 227 |
+
ap.add_argument("--assets", nargs="+", default=["BTCUSDT_1h", "AAPL_1h"])
|
| 228 |
+
ap.add_argument("--windows", type=int, default=6)
|
| 229 |
+
ap.add_argument("--horizon", type=int, default=64)
|
| 230 |
+
ap.add_argument("--paths", type=int, default=4)
|
| 231 |
+
args = ap.parse_args()
|
| 232 |
+
|
| 233 |
+
results_dir = LAB / "results"
|
| 234 |
+
results_dir.mkdir(exist_ok=True)
|
| 235 |
+
|
| 236 |
+
print("Loading Kronos-small (CPU)...", flush=True)
|
| 237 |
+
tokenizer = KronosTokenizer.from_pretrained("NeoQuasar/Kronos-Tokenizer-base")
|
| 238 |
+
model = Kronos.from_pretrained("NeoQuasar/Kronos-small")
|
| 239 |
+
tokenizer.eval(); model.eval()
|
| 240 |
+
predictor = KronosPredictor(model, tokenizer, device="cpu", max_context=512)
|
| 241 |
+
|
| 242 |
+
summaries = []
|
| 243 |
+
for name in args.assets:
|
| 244 |
+
df = pd.read_csv(LAB / "data" / f"{name}.csv", parse_dates=["timestamps"])
|
| 245 |
+
summaries.append(run_asset(name, df, predictor, args, results_dir))
|
| 246 |
+
(results_dir / "summary.json").write_text(json.dumps(summaries, indent=2), encoding="utf-8")
|
| 247 |
+
|
| 248 |
+
write_report(summaries, args, results_dir)
|
| 249 |
+
print("\nDone. See roughness_lab/results/report.md", flush=True)
|
| 250 |
+
|
| 251 |
+
|
| 252 |
+
if __name__ == "__main__":
|
| 253 |
+
main()
|
roughness_lab/fetch_data.py
ADDED
|
@@ -0,0 +1,102 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Fetch the calibration datasets:
|
| 2 |
+
|
| 3 |
+
* BTCUSDT 1h, ~3 years, paginated from Binance public market data
|
| 4 |
+
* AAPL 1h, ~2 years (Yahoo's hourly history limit), exchange-local time,
|
| 5 |
+
off-grid session-close snapshot bars dropped
|
| 6 |
+
|
| 7 |
+
Saved to roughness_lab/data/<NAME>.csv with columns
|
| 8 |
+
timestamps, open, high, low, close, volume, amount.
|
| 9 |
+
"""
|
| 10 |
+
import time
|
| 11 |
+
from pathlib import Path
|
| 12 |
+
|
| 13 |
+
import pandas as pd
|
| 14 |
+
import requests
|
| 15 |
+
|
| 16 |
+
DATA_DIR = Path(__file__).resolve().parent / "data"
|
| 17 |
+
UA = {"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64)"}
|
| 18 |
+
BINANCE = "https://data-api.binance.vision/api/v3/klines"
|
| 19 |
+
HOUR_MS = 3_600_000
|
| 20 |
+
|
| 21 |
+
|
| 22 |
+
def fetch_btc(years: float = 3.0) -> pd.DataFrame:
|
| 23 |
+
end = int(time.time() * 1000)
|
| 24 |
+
start = end - int(years * 365.25 * 24 * HOUR_MS)
|
| 25 |
+
rows = []
|
| 26 |
+
cursor = start
|
| 27 |
+
while cursor < end:
|
| 28 |
+
r = requests.get(
|
| 29 |
+
BINANCE,
|
| 30 |
+
params={"symbol": "BTCUSDT", "interval": "1h", "startTime": cursor, "limit": 1000},
|
| 31 |
+
timeout=20,
|
| 32 |
+
)
|
| 33 |
+
r.raise_for_status()
|
| 34 |
+
batch = r.json()
|
| 35 |
+
if not batch:
|
| 36 |
+
break
|
| 37 |
+
rows.extend(batch)
|
| 38 |
+
cursor = batch[-1][0] + HOUR_MS
|
| 39 |
+
print(f" BTC: {len(rows)} bars (up to {pd.Timestamp(batch[-1][0], unit='ms')})")
|
| 40 |
+
time.sleep(0.15) # stay well under rate limits
|
| 41 |
+
df = pd.DataFrame(
|
| 42 |
+
{
|
| 43 |
+
"timestamps": pd.to_datetime([b[0] for b in rows], unit="ms"),
|
| 44 |
+
"open": [float(b[1]) for b in rows],
|
| 45 |
+
"high": [float(b[2]) for b in rows],
|
| 46 |
+
"low": [float(b[3]) for b in rows],
|
| 47 |
+
"close": [float(b[4]) for b in rows],
|
| 48 |
+
"volume": [float(b[5]) for b in rows],
|
| 49 |
+
"amount": [float(b[7]) for b in rows],
|
| 50 |
+
}
|
| 51 |
+
)
|
| 52 |
+
df = df.drop_duplicates(subset="timestamps").sort_values("timestamps").reset_index(drop=True)
|
| 53 |
+
return df.iloc[:-1] # drop the still-forming bar
|
| 54 |
+
|
| 55 |
+
|
| 56 |
+
def fetch_aapl() -> pd.DataFrame:
|
| 57 |
+
r = requests.get(
|
| 58 |
+
"https://query1.finance.yahoo.com/v8/finance/chart/AAPL",
|
| 59 |
+
params={"interval": "60m", "range": "730d", "includePrePost": "false"},
|
| 60 |
+
headers=UA, timeout=30,
|
| 61 |
+
)
|
| 62 |
+
r.raise_for_status()
|
| 63 |
+
res = r.json()["chart"]["result"][0]
|
| 64 |
+
ts = res["timestamp"]
|
| 65 |
+
q = res["indicators"]["quote"][0]
|
| 66 |
+
off = int(res["meta"].get("gmtoffset", 0))
|
| 67 |
+
df = pd.DataFrame(
|
| 68 |
+
{
|
| 69 |
+
"timestamps": pd.to_datetime([t + off for t in ts], unit="s"),
|
| 70 |
+
"open": q["open"], "high": q["high"], "low": q["low"],
|
| 71 |
+
"close": q["close"], "volume": q["volume"],
|
| 72 |
+
}
|
| 73 |
+
).dropna(subset=["open", "high", "low", "close"]).reset_index(drop=True)
|
| 74 |
+
df["volume"] = df["volume"].astype(float).fillna(0.0)
|
| 75 |
+
df["amount"] = df["volume"] * df[["open", "high", "low", "close"]].mean(axis=1)
|
| 76 |
+
|
| 77 |
+
# Drop the still-forming bar, then any trailing off-grid close-print bars
|
| 78 |
+
# (Yahoo stamps a 16:00 snapshot after the 15:30 hourly bar).
|
| 79 |
+
now_local = pd.Timestamp.now("UTC").tz_localize(None) + pd.Timedelta(seconds=off)
|
| 80 |
+
if len(df) and df["timestamps"].iloc[-1] + pd.Timedelta(hours=1) > now_local:
|
| 81 |
+
df = df.iloc[:-1]
|
| 82 |
+
while len(df) > 1:
|
| 83 |
+
d = (df["timestamps"].iloc[-1] - df["timestamps"].iloc[-2]).total_seconds()
|
| 84 |
+
same_phase = df["timestamps"].iloc[-1].minute == df["timestamps"].iloc[-2].minute
|
| 85 |
+
if d == 3600 or same_phase:
|
| 86 |
+
break
|
| 87 |
+
df = df.iloc[:-1]
|
| 88 |
+
return df.reset_index(drop=True)
|
| 89 |
+
|
| 90 |
+
|
| 91 |
+
if __name__ == "__main__":
|
| 92 |
+
DATA_DIR.mkdir(exist_ok=True)
|
| 93 |
+
|
| 94 |
+
print("Fetching BTCUSDT 1h (~3y, paginated)...")
|
| 95 |
+
btc = fetch_btc()
|
| 96 |
+
btc.to_csv(DATA_DIR / "BTCUSDT_1h.csv", index=False)
|
| 97 |
+
print(f"BTC saved: {len(btc)} bars, {btc['timestamps'].iloc[0]} .. {btc['timestamps'].iloc[-1]}")
|
| 98 |
+
|
| 99 |
+
print("Fetching AAPL 1h (~2y)...")
|
| 100 |
+
aapl = fetch_aapl()
|
| 101 |
+
aapl.to_csv(DATA_DIR / "AAPL_1h.csv", index=False)
|
| 102 |
+
print(f"AAPL saved: {len(aapl)} bars, {aapl['timestamps'].iloc[0]} .. {aapl['timestamps'].iloc[-1]}")
|
roughness_lab/gpu_finetune/README_GPU.md
ADDED
|
@@ -0,0 +1,86 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Roughness-aware Kronos fine-tuning (GPU-ready)
|
| 2 |
+
|
| 3 |
+
Fine-tunes Kronos on your data with **surface-roughness checkpoint selection**:
|
| 4 |
+
every epoch, sampled forecasts of held-out windows are scored against the
|
| 5 |
+
realized continuation using surface-texture parameters (Ra, Rq, Rz, RSm, Rsk,
|
| 6 |
+
Rku, Wa — see `../roughness.py`), and the checkpoint with the most realistic
|
| 7 |
+
*texture* is kept alongside the usual lowest-val-loss one.
|
| 8 |
+
|
| 9 |
+
## Files
|
| 10 |
+
|
| 11 |
+
| file | purpose |
|
| 12 |
+
|---|---|
|
| 13 |
+
| `train_rough.py` | orchestrator: upstream tokenizer phase + roughness-aware predictor phase |
|
| 14 |
+
| `evaluate_texture.py` | score any checkpoint (or the pretrained baseline) post-hoc |
|
| 15 |
+
| `config_btc_1h.yaml` | BTC/USDT 1h experiment (edit absolute paths first) |
|
| 16 |
+
| `config_aapl_1h.yaml` | AAPL 1h experiment (edit absolute paths first) |
|
| 17 |
+
| `config_smoke_cpu.yaml` | minutes-long CPU smoke test of the whole pipeline |
|
| 18 |
+
|
| 19 |
+
Everything reuses the upstream `Kronos/finetune_csv` code (dataset, loaders,
|
| 20 |
+
tokenizer trainer, config loader) — no files in the clone are modified.
|
| 21 |
+
|
| 22 |
+
## Setup on the GPU box
|
| 23 |
+
|
| 24 |
+
```bash
|
| 25 |
+
git clone https://github.com/shiyu-coder/Kronos
|
| 26 |
+
# copy the roughness_lab/ folder (this folder + roughness.py + data/) next to the clone:
|
| 27 |
+
# <root>/Kronos
|
| 28 |
+
# <root>/roughness_lab/...
|
| 29 |
+
pip install -r Kronos/requirements.txt
|
| 30 |
+
pip install pyyaml tabulate # config parsing + markdown tables
|
| 31 |
+
```
|
| 32 |
+
|
| 33 |
+
Then edit the two `/ABSOLUTE/PATH/...` entries in the config you want to run.
|
| 34 |
+
|
| 35 |
+
## Run
|
| 36 |
+
|
| 37 |
+
```bash
|
| 38 |
+
cd roughness_lab/gpu_finetune
|
| 39 |
+
|
| 40 |
+
# single GPU (tokenizer phase + predictor phase, in order)
|
| 41 |
+
python train_rough.py --config config_btc_1h.yaml
|
| 42 |
+
|
| 43 |
+
# multiple GPUs: tokenizer phase is single-process, predictor phase is DDP —
|
| 44 |
+
# 1) run with experiment.train_basemodel: false (tokenizer only)
|
| 45 |
+
# 2) run with experiment.train_tokenizer: false under torchrun:
|
| 46 |
+
torchrun --standalone --nproc_per_node=4 train_rough.py --config config_btc_1h.yaml
|
| 47 |
+
```
|
| 48 |
+
|
| 49 |
+
Before running, set `roughness.eval_T` / `eval_top_p` in the config to the
|
| 50 |
+
best sampling configuration reported by the calibration study
|
| 51 |
+
(`roughness_lab/results/report.md`) so checkpoint selection happens at the
|
| 52 |
+
operating point you will actually use.
|
| 53 |
+
|
| 54 |
+
## Outputs (under `finetuned/<exp_name>/basemodel/`)
|
| 55 |
+
|
| 56 |
+
- `best_model/` — lowest validation loss (upstream behavior)
|
| 57 |
+
- `best_texture/` — lowest texture error (this pipeline's addition)
|
| 58 |
+
- `epoch_NN/` — every epoch (`roughness.save_every_epoch: true`)
|
| 59 |
+
- `metrics_log.csv` — per-epoch train loss, val loss, texture error
|
| 60 |
+
- `ranking.md` — summary table + which epoch won each criterion
|
| 61 |
+
|
| 62 |
+
## Compare checkpoints (including the pretrained baseline)
|
| 63 |
+
|
| 64 |
+
```bash
|
| 65 |
+
python evaluate_texture.py --csv ../data/BTCUSDT_1h.csv --device cuda:0 # baseline
|
| 66 |
+
python evaluate_texture.py --csv ../data/BTCUSDT_1h.csv --device cuda:0 \
|
| 67 |
+
--model finetuned/btc_1h_rough/basemodel/best_texture \
|
| 68 |
+
--tokenizer finetuned/btc_1h_rough/tokenizer/best_model
|
| 69 |
+
```
|
| 70 |
+
|
| 71 |
+
## Expected runtime (estimates)
|
| 72 |
+
|
| 73 |
+
- BTC 1h config (~22k train windows, batch 64, 512 ctx): roughly 3–6 min/epoch
|
| 74 |
+
predictor on an A100-class GPU, tokenizer phase faster; texture eval adds
|
| 75 |
+
seconds per epoch on GPU. Whole BTC experiment ≈ 1–2 h single GPU.
|
| 76 |
+
- The CPU smoke test (`config_smoke_cpu.yaml`) runs the full plumbing in
|
| 77 |
+
minutes and was verified on this machine — see the bottom of this file.
|
| 78 |
+
|
| 79 |
+
## Why texture selection?
|
| 80 |
+
|
| 81 |
+
Validation token-loss rewards average correctness; it can prefer checkpoints
|
| 82 |
+
whose samples are too smooth (volatility-damped). The texture criterion keeps
|
| 83 |
+
the checkpoint whose *generated* price paths statistically resemble real
|
| 84 |
+
market texture (noise amplitude Ra/Rq, swing size Rz, wiggle spacing RSm,
|
| 85 |
+
tail shape Rku) at your chosen sampling settings. Direction accuracy and
|
| 86 |
+
texture realism are complementary — `metrics_log.csv` lets you see both.
|
roughness_lab/gpu_finetune/config_aapl_1h.yaml
ADDED
|
@@ -0,0 +1,65 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Roughness-aware fine-tune of Kronos-small on AAPL 1h (~3y).
|
| 2 |
+
# EDIT THE ABSOLUTE PATHS below for your machine (same convention as the
|
| 3 |
+
# upstream finetune_csv template), then see README_GPU.md for commands.
|
| 4 |
+
|
| 5 |
+
data:
|
| 6 |
+
data_path: "/ABSOLUTE/PATH/TO/Kronos-small/roughness_lab/data/AAPL_1h.csv"
|
| 7 |
+
lookback_window: 512
|
| 8 |
+
predict_window: 64
|
| 9 |
+
max_context: 512
|
| 10 |
+
clip: 5.0
|
| 11 |
+
train_ratio: 0.85
|
| 12 |
+
val_ratio: 0.15
|
| 13 |
+
test_ratio: 0.0
|
| 14 |
+
|
| 15 |
+
training:
|
| 16 |
+
tokenizer_epochs: 8
|
| 17 |
+
basemodel_epochs: 15
|
| 18 |
+
batch_size: 64
|
| 19 |
+
log_interval: 50
|
| 20 |
+
num_workers: 4
|
| 21 |
+
seed: 42
|
| 22 |
+
tokenizer_learning_rate: 0.0001
|
| 23 |
+
predictor_learning_rate: 0.00002
|
| 24 |
+
adam_beta1: 0.9
|
| 25 |
+
adam_beta2: 0.95
|
| 26 |
+
adam_weight_decay: 0.1
|
| 27 |
+
accumulation_steps: 1
|
| 28 |
+
|
| 29 |
+
model_paths:
|
| 30 |
+
# Hugging Face hub names work directly; local dirs work too.
|
| 31 |
+
pretrained_tokenizer: "NeoQuasar/Kronos-Tokenizer-base"
|
| 32 |
+
pretrained_predictor: "NeoQuasar/Kronos-small"
|
| 33 |
+
exp_name: "aapl_1h_rough"
|
| 34 |
+
base_path: "/ABSOLUTE/PATH/TO/Kronos-small/roughness_lab/gpu_finetune/finetuned/"
|
| 35 |
+
base_save_path: ""
|
| 36 |
+
finetuned_tokenizer: ""
|
| 37 |
+
tokenizer_save_name: "tokenizer"
|
| 38 |
+
basemodel_save_name: "basemodel"
|
| 39 |
+
|
| 40 |
+
experiment:
|
| 41 |
+
name: "kronos_rough_aapl"
|
| 42 |
+
description: "Kronos-small on AAPL 1h with surface-roughness checkpoint selection"
|
| 43 |
+
use_comet: false
|
| 44 |
+
train_tokenizer: true
|
| 45 |
+
train_basemodel: true
|
| 46 |
+
skip_existing: false
|
| 47 |
+
|
| 48 |
+
device:
|
| 49 |
+
use_cuda: true
|
| 50 |
+
device_id: 0
|
| 51 |
+
|
| 52 |
+
# Texture evaluation during training (read by train_rough.py only).
|
| 53 |
+
# Set eval_T / eval_top_p to the best configuration found by
|
| 54 |
+
# roughness_lab/calibrate.py (see roughness_lab/results/report.md).
|
| 55 |
+
roughness:
|
| 56 |
+
save_every_epoch: true
|
| 57 |
+
eval_every: 1
|
| 58 |
+
eval_windows: 6
|
| 59 |
+
eval_paths: 6
|
| 60 |
+
eval_horizon: 64
|
| 61 |
+
eval_T: 1.0
|
| 62 |
+
eval_top_p: 0.9
|
| 63 |
+
cutoff: 8
|
| 64 |
+
wa_cutoff: 16
|
| 65 |
+
|
roughness_lab/gpu_finetune/config_btc_1h.yaml
ADDED
|
@@ -0,0 +1,64 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Roughness-aware fine-tune of Kronos-small on BTCUSDT 1h (~3y).
|
| 2 |
+
# EDIT THE ABSOLUTE PATHS below for your machine (same convention as the
|
| 3 |
+
# upstream finetune_csv template), then see README_GPU.md for commands.
|
| 4 |
+
|
| 5 |
+
data:
|
| 6 |
+
data_path: "/ABSOLUTE/PATH/TO/Kronos-small/roughness_lab/data/BTCUSDT_1h.csv"
|
| 7 |
+
lookback_window: 512
|
| 8 |
+
predict_window: 64
|
| 9 |
+
max_context: 512
|
| 10 |
+
clip: 5.0
|
| 11 |
+
train_ratio: 0.85
|
| 12 |
+
val_ratio: 0.15
|
| 13 |
+
test_ratio: 0.0
|
| 14 |
+
|
| 15 |
+
training:
|
| 16 |
+
tokenizer_epochs: 8
|
| 17 |
+
basemodel_epochs: 15
|
| 18 |
+
batch_size: 64
|
| 19 |
+
log_interval: 50
|
| 20 |
+
num_workers: 4
|
| 21 |
+
seed: 42
|
| 22 |
+
tokenizer_learning_rate: 0.0001
|
| 23 |
+
predictor_learning_rate: 0.00002
|
| 24 |
+
adam_beta1: 0.9
|
| 25 |
+
adam_beta2: 0.95
|
| 26 |
+
adam_weight_decay: 0.1
|
| 27 |
+
accumulation_steps: 1
|
| 28 |
+
|
| 29 |
+
model_paths:
|
| 30 |
+
# Hugging Face hub names work directly; local dirs work too.
|
| 31 |
+
pretrained_tokenizer: "NeoQuasar/Kronos-Tokenizer-base"
|
| 32 |
+
pretrained_predictor: "NeoQuasar/Kronos-small"
|
| 33 |
+
exp_name: "btc_1h_rough"
|
| 34 |
+
base_path: "/ABSOLUTE/PATH/TO/Kronos-small/roughness_lab/gpu_finetune/finetuned/"
|
| 35 |
+
base_save_path: ""
|
| 36 |
+
finetuned_tokenizer: ""
|
| 37 |
+
tokenizer_save_name: "tokenizer"
|
| 38 |
+
basemodel_save_name: "basemodel"
|
| 39 |
+
|
| 40 |
+
experiment:
|
| 41 |
+
name: "kronos_rough_btc"
|
| 42 |
+
description: "Kronos-small on BTC 1h with surface-roughness checkpoint selection"
|
| 43 |
+
use_comet: false
|
| 44 |
+
train_tokenizer: true
|
| 45 |
+
train_basemodel: true
|
| 46 |
+
skip_existing: false
|
| 47 |
+
|
| 48 |
+
device:
|
| 49 |
+
use_cuda: true
|
| 50 |
+
device_id: 0
|
| 51 |
+
|
| 52 |
+
# Texture evaluation during training (read by train_rough.py only).
|
| 53 |
+
# Set eval_T / eval_top_p to the best configuration found by
|
| 54 |
+
# roughness_lab/calibrate.py (see roughness_lab/results/report.md).
|
| 55 |
+
roughness:
|
| 56 |
+
save_every_epoch: true
|
| 57 |
+
eval_every: 1
|
| 58 |
+
eval_windows: 6
|
| 59 |
+
eval_paths: 6
|
| 60 |
+
eval_horizon: 64
|
| 61 |
+
eval_T: 1.0
|
| 62 |
+
eval_top_p: 0.9
|
| 63 |
+
cutoff: 8
|
| 64 |
+
wa_cutoff: 16
|
roughness_lab/gpu_finetune/config_smoke_cpu.yaml
ADDED
|
@@ -0,0 +1,60 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Tiny CPU smoke test: proves the whole pipeline executes end-to-end.
|
| 2 |
+
# Runs in minutes on a laptop; numbers are meaningless, plumbing is real.
|
| 3 |
+
|
| 4 |
+
data:
|
| 5 |
+
data_path: "C:/Users/ademo/Downloads/Kronos-small/roughness_lab/data/BTC_smoke.csv"
|
| 6 |
+
lookback_window: 256
|
| 7 |
+
predict_window: 32
|
| 8 |
+
max_context: 256
|
| 9 |
+
clip: 5.0
|
| 10 |
+
# val slice must exceed lookback+predict+1 bars for the upstream dataset
|
| 11 |
+
train_ratio: 0.8
|
| 12 |
+
val_ratio: 0.2
|
| 13 |
+
test_ratio: 0.0
|
| 14 |
+
|
| 15 |
+
training:
|
| 16 |
+
tokenizer_epochs: 1
|
| 17 |
+
basemodel_epochs: 2
|
| 18 |
+
batch_size: 6
|
| 19 |
+
log_interval: 25
|
| 20 |
+
num_workers: 0
|
| 21 |
+
seed: 42
|
| 22 |
+
tokenizer_learning_rate: 0.0001
|
| 23 |
+
predictor_learning_rate: 0.00002
|
| 24 |
+
adam_beta1: 0.9
|
| 25 |
+
adam_beta2: 0.95
|
| 26 |
+
adam_weight_decay: 0.1
|
| 27 |
+
accumulation_steps: 1
|
| 28 |
+
|
| 29 |
+
model_paths:
|
| 30 |
+
pretrained_tokenizer: "NeoQuasar/Kronos-Tokenizer-base"
|
| 31 |
+
pretrained_predictor: "NeoQuasar/Kronos-small"
|
| 32 |
+
exp_name: "smoke_cpu"
|
| 33 |
+
base_path: "C:/Users/ademo/Downloads/Kronos-small/roughness_lab/gpu_finetune/finetuned/"
|
| 34 |
+
base_save_path: ""
|
| 35 |
+
finetuned_tokenizer: ""
|
| 36 |
+
tokenizer_save_name: "tokenizer"
|
| 37 |
+
basemodel_save_name: "basemodel"
|
| 38 |
+
|
| 39 |
+
experiment:
|
| 40 |
+
name: "kronos_rough_smoke"
|
| 41 |
+
description: "CPU smoke test of the roughness-aware pipeline"
|
| 42 |
+
use_comet: false
|
| 43 |
+
train_tokenizer: true
|
| 44 |
+
train_basemodel: true
|
| 45 |
+
skip_existing: false
|
| 46 |
+
|
| 47 |
+
device:
|
| 48 |
+
use_cuda: false
|
| 49 |
+
device_id: 0
|
| 50 |
+
|
| 51 |
+
roughness:
|
| 52 |
+
save_every_epoch: true
|
| 53 |
+
eval_every: 1
|
| 54 |
+
eval_windows: 2
|
| 55 |
+
eval_paths: 2
|
| 56 |
+
eval_horizon: 32
|
| 57 |
+
eval_T: 1.0
|
| 58 |
+
eval_top_p: 0.9
|
| 59 |
+
cutoff: 8
|
| 60 |
+
wa_cutoff: 16
|
roughness_lab/gpu_finetune/evaluate_texture.py
ADDED
|
@@ -0,0 +1,58 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Standalone texture evaluation of any Kronos checkpoint.
|
| 2 |
+
|
| 3 |
+
Scores how realistically a (tokenizer, predictor) pair reproduces the
|
| 4 |
+
surface-roughness texture of held-out data. Use it to compare the pretrained
|
| 5 |
+
baseline against fine-tuned checkpoints (best_model, best_texture, epoch_NN).
|
| 6 |
+
|
| 7 |
+
Example:
|
| 8 |
+
python evaluate_texture.py --model finetuned/btc_1h_rough/basemodel/best_texture \
|
| 9 |
+
--tokenizer finetuned/btc_1h_rough/tokenizer/best_model \
|
| 10 |
+
--csv ../data/BTCUSDT_1h.csv --device cuda:0
|
| 11 |
+
"""
|
| 12 |
+
import argparse
|
| 13 |
+
from types import SimpleNamespace
|
| 14 |
+
from pathlib import Path
|
| 15 |
+
import sys
|
| 16 |
+
|
| 17 |
+
HERE = Path(__file__).resolve().parent
|
| 18 |
+
sys.path.insert(0, str(HERE))
|
| 19 |
+
|
| 20 |
+
from train_rough import TextureEvaluator, PARAM_KEYS # noqa: E402
|
| 21 |
+
from model import Kronos, KronosTokenizer # noqa: E402
|
| 22 |
+
|
| 23 |
+
|
| 24 |
+
def main() -> None:
|
| 25 |
+
ap = argparse.ArgumentParser()
|
| 26 |
+
ap.add_argument("--model", default="NeoQuasar/Kronos-small")
|
| 27 |
+
ap.add_argument("--tokenizer", default="NeoQuasar/Kronos-Tokenizer-base")
|
| 28 |
+
ap.add_argument("--csv", required=True)
|
| 29 |
+
ap.add_argument("--device", default="cpu")
|
| 30 |
+
ap.add_argument("--context", type=int, default=512)
|
| 31 |
+
ap.add_argument("--horizon", type=int, default=64)
|
| 32 |
+
ap.add_argument("--windows", type=int, default=6)
|
| 33 |
+
ap.add_argument("--paths", type=int, default=6)
|
| 34 |
+
ap.add_argument("--T", type=float, default=1.0)
|
| 35 |
+
ap.add_argument("--top_p", type=float, default=0.9)
|
| 36 |
+
ap.add_argument("--train-ratio", type=float, default=0.85,
|
| 37 |
+
help="windows are drawn after this fraction (the val region)")
|
| 38 |
+
args = ap.parse_args()
|
| 39 |
+
|
| 40 |
+
config = SimpleNamespace(
|
| 41 |
+
data_path=args.csv, train_ratio=args.train_ratio,
|
| 42 |
+
val_ratio=1.0 - args.train_ratio, lookback_window=args.context,
|
| 43 |
+
max_context=args.context,
|
| 44 |
+
)
|
| 45 |
+
rough_cfg = {"eval_horizon": args.horizon, "eval_windows": args.windows,
|
| 46 |
+
"eval_paths": args.paths, "eval_T": args.T, "eval_top_p": args.top_p}
|
| 47 |
+
|
| 48 |
+
tokenizer = KronosTokenizer.from_pretrained(args.tokenizer).eval()
|
| 49 |
+
model = Kronos.from_pretrained(args.model).eval()
|
| 50 |
+
evaluator = TextureEvaluator(config, rough_cfg, args.device)
|
| 51 |
+
score = evaluator.score(model, tokenizer)
|
| 52 |
+
print(f"\nmodel: {args.model}\ntokenizer: {args.tokenizer}")
|
| 53 |
+
print(f"texture_error = {score:.4f} (windows={args.windows}, paths={args.paths}, "
|
| 54 |
+
f"T={args.T}, top_p={args.top_p}; lower is better)")
|
| 55 |
+
|
| 56 |
+
|
| 57 |
+
if __name__ == "__main__":
|
| 58 |
+
main()
|
roughness_lab/gpu_finetune/train_rough.py
ADDED
|
@@ -0,0 +1,274 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Roughness-aware Kronos fine-tuning (GPU-ready).
|
| 2 |
+
|
| 3 |
+
Wraps the repo's finetune_csv pipeline:
|
| 4 |
+
* phase 1: tokenizer fine-tune — runs the repo's finetune_tokenizer.py
|
| 5 |
+
unchanged (skipped if experiment.train_tokenizer is false);
|
| 6 |
+
* phase 2: predictor fine-tune — same objective/optimizer/schedule as the
|
| 7 |
+
repo's finetune_base_model.py, plus per-epoch checkpoints and a
|
| 8 |
+
surface-roughness texture evaluation (Ra/Rq/Rz/RSm/Rsk/Rku/Wa of sampled
|
| 9 |
+
forecasts vs realized validation windows). Two "best" checkpoints are
|
| 10 |
+
kept: best_model (lowest val loss, as upstream) and best_texture (lowest
|
| 11 |
+
texture error). metrics_log.csv + ranking.md let you compare.
|
| 12 |
+
|
| 13 |
+
Single GPU / CPU: python train_rough.py --config config_btc_1h.yaml
|
| 14 |
+
Multi-GPU (DDP): torchrun --standalone --nproc_per_node=N train_rough.py --config ...
|
| 15 |
+
"""
|
| 16 |
+
import argparse
|
| 17 |
+
import os
|
| 18 |
+
import sys
|
| 19 |
+
import time
|
| 20 |
+
from pathlib import Path
|
| 21 |
+
|
| 22 |
+
import numpy as np
|
| 23 |
+
import pandas as pd
|
| 24 |
+
import torch
|
| 25 |
+
import torch.distributed as dist
|
| 26 |
+
from torch.nn.parallel import DistributedDataParallel as DDP
|
| 27 |
+
|
| 28 |
+
HERE = Path(__file__).resolve().parent
|
| 29 |
+
LAB = HERE.parent
|
| 30 |
+
REPO = LAB.parent / "Kronos"
|
| 31 |
+
FTCSV = REPO / "finetune_csv"
|
| 32 |
+
for p in (str(LAB), str(REPO), str(FTCSV)):
|
| 33 |
+
sys.path.insert(0, p)
|
| 34 |
+
|
| 35 |
+
from roughness import roughness_params, texture_error # noqa: E402
|
| 36 |
+
from model import Kronos, KronosTokenizer, KronosPredictor # noqa: E402
|
| 37 |
+
from config_loader import ConfigLoader, CustomFinetuneConfig # noqa: E402
|
| 38 |
+
from finetune_base_model import create_dataloaders, setup_logging # noqa: E402
|
| 39 |
+
from finetune_tokenizer import set_seed, train_tokenizer # noqa: E402
|
| 40 |
+
from finetune_tokenizer import setup_logging as setup_tok_logging # noqa: E402
|
| 41 |
+
|
| 42 |
+
PARAM_KEYS = ["ra", "rq", "rz", "rsm", "rsk", "rku", "wa"]
|
| 43 |
+
|
| 44 |
+
|
| 45 |
+
def params_for(close: np.ndarray, cutoff: int, wa_cutoff: int) -> dict:
|
| 46 |
+
p = roughness_params(close, cutoff).as_dict()
|
| 47 |
+
p["wa"] = roughness_params(close, wa_cutoff).wa
|
| 48 |
+
return p
|
| 49 |
+
|
| 50 |
+
|
| 51 |
+
class TextureEvaluator:
|
| 52 |
+
"""Forecasts held-out validation windows and scores texture realism."""
|
| 53 |
+
|
| 54 |
+
def __init__(self, config, rough_cfg: dict, device):
|
| 55 |
+
df = pd.read_csv(config.data_path, parse_dates=["timestamps"])
|
| 56 |
+
df = df.sort_values("timestamps").reset_index(drop=True)
|
| 57 |
+
n = len(df)
|
| 58 |
+
val_start = int(n * config.train_ratio)
|
| 59 |
+
val_end = int(n * (config.train_ratio + config.val_ratio))
|
| 60 |
+
self.df = df
|
| 61 |
+
self.context = config.lookback_window
|
| 62 |
+
self.horizon = int(rough_cfg.get("eval_horizon", 64))
|
| 63 |
+
self.paths = int(rough_cfg.get("eval_paths", 4))
|
| 64 |
+
self.T = float(rough_cfg.get("eval_T", 1.0))
|
| 65 |
+
self.top_p = float(rough_cfg.get("eval_top_p", 0.9))
|
| 66 |
+
self.cutoff = int(rough_cfg.get("cutoff", 8))
|
| 67 |
+
self.wa_cutoff = int(rough_cfg.get("wa_cutoff", 16))
|
| 68 |
+
self.max_context = config.max_context
|
| 69 |
+
self.device = device
|
| 70 |
+
k = int(rough_cfg.get("eval_windows", 4))
|
| 71 |
+
lo = max(self.context, val_start + self.context)
|
| 72 |
+
hi = val_end - self.horizon - 1
|
| 73 |
+
if hi <= lo: # validation slice too small: fall back to series tail
|
| 74 |
+
lo, hi = max(self.context, n // 2), n - self.horizon - 1
|
| 75 |
+
self.anchors = np.linspace(lo, hi, k).astype(int)
|
| 76 |
+
feat = ["open", "high", "low", "close", "volume", "amount"]
|
| 77 |
+
self.ctx_dfs, self.x_tss, self.y_tss, self.real_params = [], [], [], []
|
| 78 |
+
for a in self.anchors:
|
| 79 |
+
ctx = df.iloc[a - self.context:a]
|
| 80 |
+
real = df.iloc[a:a + self.horizon]
|
| 81 |
+
self.ctx_dfs.append(ctx[feat].reset_index(drop=True))
|
| 82 |
+
self.x_tss.append(ctx["timestamps"].reset_index(drop=True))
|
| 83 |
+
self.y_tss.append(real["timestamps"].reset_index(drop=True))
|
| 84 |
+
self.real_params.append(params_for(real["close"].to_numpy(), self.cutoff, self.wa_cutoff))
|
| 85 |
+
|
| 86 |
+
@torch.no_grad()
|
| 87 |
+
def score(self, model, tokenizer) -> float:
|
| 88 |
+
predictor = KronosPredictor(model, tokenizer, device=self.device, max_context=self.max_context)
|
| 89 |
+
preds = predictor.predict_batch(
|
| 90 |
+
df_list=[c for c in self.ctx_dfs for _ in range(self.paths)],
|
| 91 |
+
x_timestamp_list=[x for x in self.x_tss for _ in range(self.paths)],
|
| 92 |
+
y_timestamp_list=[y for y in self.y_tss for _ in range(self.paths)],
|
| 93 |
+
pred_len=self.horizon, T=self.T, top_p=self.top_p, sample_count=1, verbose=False,
|
| 94 |
+
)
|
| 95 |
+
errs = []
|
| 96 |
+
for wi in range(len(self.anchors)):
|
| 97 |
+
pp = [params_for(p["close"].to_numpy(), self.cutoff, self.wa_cutoff)
|
| 98 |
+
for p in preds[wi * self.paths:(wi + 1) * self.paths]]
|
| 99 |
+
med = {k: float(np.median([q[k] for q in pp])) for k in PARAM_KEYS}
|
| 100 |
+
errs.append(texture_error(med, self.real_params[wi]))
|
| 101 |
+
return float(np.mean(errs))
|
| 102 |
+
|
| 103 |
+
|
| 104 |
+
def run_tokenizer_phase(config) -> None:
|
| 105 |
+
"""Upstream tokenizer fine-tune, invoked as a function.
|
| 106 |
+
|
| 107 |
+
Deliberately not via `python finetune_tokenizer.py --config ...`: that
|
| 108 |
+
script's main() has a scoping bug (`import json, os` inside a branch
|
| 109 |
+
shadows the module-level `os`) that crashes the pretrained-tokenizer
|
| 110 |
+
path. Calling train_tokenizer() directly sidesteps it without touching
|
| 111 |
+
the upstream clone."""
|
| 112 |
+
print("\n=== Phase 1: tokenizer fine-tune (upstream trainer) ===", flush=True)
|
| 113 |
+
device = torch.device("cuda" if config.use_cuda and torch.cuda.is_available() else "cpu")
|
| 114 |
+
os.makedirs(config.tokenizer_save_path, exist_ok=True)
|
| 115 |
+
logger = setup_tok_logging(config.exp_name, os.path.join(config.base_save_path, "logs"), 0)
|
| 116 |
+
set_seed(config.seed)
|
| 117 |
+
print(f"Loading pretrained tokenizer: {config.pretrained_tokenizer_path}")
|
| 118 |
+
tokenizer = KronosTokenizer.from_pretrained(config.pretrained_tokenizer_path).to(device)
|
| 119 |
+
best = train_tokenizer(tokenizer, device, config, config.tokenizer_save_path, logger)
|
| 120 |
+
print(f"Tokenizer phase done (best val loss {best:.4f}) -> {config.tokenizer_save_path}", flush=True)
|
| 121 |
+
|
| 122 |
+
|
| 123 |
+
def run_predictor_phase(config_path: str, config, rough_cfg: dict) -> None:
|
| 124 |
+
use_ddp_env = int(os.environ.get("WORLD_SIZE", "1")) > 1
|
| 125 |
+
rank = int(os.environ.get("RANK", "0"))
|
| 126 |
+
if use_ddp_env and torch.cuda.is_available() and not dist.is_initialized():
|
| 127 |
+
dist.init_process_group(backend=os.environ.get("DIST_BACKEND", "nccl"))
|
| 128 |
+
use_ddp = dist.is_available() and dist.is_initialized()
|
| 129 |
+
|
| 130 |
+
if config.use_cuda and torch.cuda.is_available():
|
| 131 |
+
local_rank = int(os.environ.get("LOCAL_RANK", str(config.device_id)))
|
| 132 |
+
torch.cuda.set_device(local_rank)
|
| 133 |
+
device = torch.device(f"cuda:{local_rank}")
|
| 134 |
+
else:
|
| 135 |
+
device = torch.device("cpu")
|
| 136 |
+
print(f"\n=== Phase 2: predictor fine-tune (device={device}, ddp={use_ddp}) ===", flush=True)
|
| 137 |
+
|
| 138 |
+
set_seed(config.seed)
|
| 139 |
+
save_dir = Path(config.basemodel_save_path)
|
| 140 |
+
save_dir.mkdir(parents=True, exist_ok=True)
|
| 141 |
+
logger = setup_logging(config.exp_name, str(Path(config.base_save_path) / "logs"), rank)
|
| 142 |
+
|
| 143 |
+
# tokenizer: prefer the just-finetuned one, fall back to pretrained
|
| 144 |
+
tok_best = Path(config.tokenizer_save_path) / "best_model"
|
| 145 |
+
tok_src = str(tok_best) if tok_best.exists() else config.pretrained_tokenizer_path
|
| 146 |
+
print(f"Tokenizer: {tok_src}")
|
| 147 |
+
tokenizer = KronosTokenizer.from_pretrained(tok_src).to(device).eval()
|
| 148 |
+
model = Kronos.from_pretrained(config.pretrained_predictor_path).to(device)
|
| 149 |
+
|
| 150 |
+
evaluator = TextureEvaluator(config, rough_cfg, device) if rank == 0 else None
|
| 151 |
+
eval_every = int(rough_cfg.get("eval_every", 1))
|
| 152 |
+
save_epochs = bool(rough_cfg.get("save_every_epoch", True))
|
| 153 |
+
|
| 154 |
+
train_loader, val_loader, train_ds, val_ds, train_sampler, _ = create_dataloaders(config)
|
| 155 |
+
optimizer = torch.optim.AdamW(model.parameters(), lr=config.predictor_learning_rate,
|
| 156 |
+
betas=(config.adam_beta1, config.adam_beta2),
|
| 157 |
+
weight_decay=config.adam_weight_decay)
|
| 158 |
+
scheduler = torch.optim.lr_scheduler.OneCycleLR(
|
| 159 |
+
optimizer, max_lr=config.predictor_learning_rate,
|
| 160 |
+
steps_per_epoch=len(train_loader), epochs=config.basemodel_epochs,
|
| 161 |
+
pct_start=0.03, div_factor=10)
|
| 162 |
+
if use_ddp:
|
| 163 |
+
lr_ = int(os.environ.get("LOCAL_RANK", "0"))
|
| 164 |
+
model = DDP(model, device_ids=[lr_], output_device=lr_)
|
| 165 |
+
raw = lambda: model.module if use_ddp else model
|
| 166 |
+
|
| 167 |
+
history, best_val, best_tex = [], float("inf"), float("inf")
|
| 168 |
+
for epoch in range(config.basemodel_epochs):
|
| 169 |
+
t0 = time.time()
|
| 170 |
+
model.train()
|
| 171 |
+
train_ds.set_epoch_seed(epoch * 10000)
|
| 172 |
+
val_ds.set_epoch_seed(0)
|
| 173 |
+
if train_sampler is not None:
|
| 174 |
+
train_sampler.set_epoch(epoch)
|
| 175 |
+
|
| 176 |
+
tr_loss, tr_n = 0.0, 0
|
| 177 |
+
for bi, (bx, bs) in enumerate(train_loader):
|
| 178 |
+
bx, bs = bx.to(device, non_blocking=True), bs.to(device, non_blocking=True)
|
| 179 |
+
with torch.no_grad():
|
| 180 |
+
t0_, t1_ = tokenizer.encode(bx, half=True)
|
| 181 |
+
logits = raw()(t0_[:, :-1], t1_[:, :-1], bs[:, :-1, :])
|
| 182 |
+
loss, _, _ = raw().head.compute_loss(logits[0], logits[1], t0_[:, 1:], t1_[:, 1:])
|
| 183 |
+
optimizer.zero_grad()
|
| 184 |
+
loss.backward()
|
| 185 |
+
torch.nn.utils.clip_grad_norm_(raw().parameters(), max_norm=3.0)
|
| 186 |
+
optimizer.step()
|
| 187 |
+
scheduler.step()
|
| 188 |
+
tr_loss += loss.item(); tr_n += 1
|
| 189 |
+
if (bi + 1) % config.log_interval == 0 and rank == 0:
|
| 190 |
+
print(f"[epoch {epoch+1}/{config.basemodel_epochs} step {bi+1}/{len(train_loader)}] "
|
| 191 |
+
f"loss {loss.item():.4f}", flush=True)
|
| 192 |
+
|
| 193 |
+
model.eval()
|
| 194 |
+
va_loss, va_n = 0.0, 0
|
| 195 |
+
with torch.no_grad():
|
| 196 |
+
for bx, bs in val_loader:
|
| 197 |
+
bx, bs = bx.to(device, non_blocking=True), bs.to(device, non_blocking=True)
|
| 198 |
+
t0_, t1_ = tokenizer.encode(bx, half=True)
|
| 199 |
+
logits = raw()(t0_[:, :-1], t1_[:, :-1], bs[:, :-1, :])
|
| 200 |
+
loss, _, _ = raw().head.compute_loss(logits[0], logits[1], t0_[:, 1:], t1_[:, 1:])
|
| 201 |
+
va_loss += loss.item(); va_n += 1
|
| 202 |
+
if use_ddp:
|
| 203 |
+
agg = torch.tensor([tr_loss, tr_n, va_loss, va_n], dtype=torch.float64, device=device)
|
| 204 |
+
dist.all_reduce(agg, op=dist.ReduceOp.SUM)
|
| 205 |
+
tr_loss, tr_n, va_loss, va_n = agg.tolist()
|
| 206 |
+
avg_tr = tr_loss / max(tr_n, 1)
|
| 207 |
+
avg_va = va_loss / max(va_n, 1)
|
| 208 |
+
|
| 209 |
+
tex = float("nan")
|
| 210 |
+
if rank == 0 and (epoch + 1) % eval_every == 0:
|
| 211 |
+
tex = evaluator.score(raw(), tokenizer)
|
| 212 |
+
|
| 213 |
+
if rank == 0:
|
| 214 |
+
dt = time.time() - t0
|
| 215 |
+
print(f"--- epoch {epoch+1}: train {avg_tr:.4f} val {avg_va:.4f} "
|
| 216 |
+
f"texture_error {tex:.4f} ({dt:.0f}s) ---", flush=True)
|
| 217 |
+
logger.info(f"epoch {epoch+1}: train={avg_tr:.4f} val={avg_va:.4f} texture={tex:.4f}")
|
| 218 |
+
history.append({"epoch": epoch + 1, "train_loss": avg_tr,
|
| 219 |
+
"val_loss": avg_va, "texture_error": tex})
|
| 220 |
+
pd.DataFrame(history).to_csv(save_dir / "metrics_log.csv", index=False)
|
| 221 |
+
if save_epochs:
|
| 222 |
+
raw().save_pretrained(str(save_dir / f"epoch_{epoch+1:02d}"))
|
| 223 |
+
if avg_va < best_val:
|
| 224 |
+
best_val = avg_va
|
| 225 |
+
raw().save_pretrained(str(save_dir / "best_model"))
|
| 226 |
+
if np.isfinite(tex) and tex < best_tex:
|
| 227 |
+
best_tex = tex
|
| 228 |
+
raw().save_pretrained(str(save_dir / "best_texture"))
|
| 229 |
+
|
| 230 |
+
if rank == 0:
|
| 231 |
+
hist = pd.DataFrame(history)
|
| 232 |
+
try:
|
| 233 |
+
table = hist.to_markdown(index=False) # needs the optional 'tabulate' package
|
| 234 |
+
except ImportError:
|
| 235 |
+
table = hist.to_string(index=False)
|
| 236 |
+
lines = ["# Checkpoint ranking\n",
|
| 237 |
+
f"Best val loss: epoch {int(hist.loc[hist.val_loss.idxmin(), 'epoch'])} "
|
| 238 |
+
f"({hist.val_loss.min():.4f}) -> `best_model/`",
|
| 239 |
+
f"Best texture: epoch {int(hist.loc[hist.texture_error.idxmin(), 'epoch'])} "
|
| 240 |
+
f"({hist.texture_error.min():.4f}) -> `best_texture/`\n",
|
| 241 |
+
table]
|
| 242 |
+
(save_dir / "ranking.md").write_text("\n".join(lines), encoding="utf-8")
|
| 243 |
+
print(f"\nDone. Checkpoints + metrics_log.csv + ranking.md in {save_dir}", flush=True)
|
| 244 |
+
if use_ddp:
|
| 245 |
+
dist.destroy_process_group()
|
| 246 |
+
|
| 247 |
+
|
| 248 |
+
def main() -> None:
|
| 249 |
+
ap = argparse.ArgumentParser()
|
| 250 |
+
ap.add_argument("--config", required=True)
|
| 251 |
+
args = ap.parse_args()
|
| 252 |
+
|
| 253 |
+
config = CustomFinetuneConfig(args.config)
|
| 254 |
+
rough_cfg = ConfigLoader(args.config).config.get("roughness", {})
|
| 255 |
+
config.print_config_summary()
|
| 256 |
+
print(f"Roughness settings: {rough_cfg}")
|
| 257 |
+
|
| 258 |
+
world_size = int(os.environ.get("WORLD_SIZE", "1"))
|
| 259 |
+
if world_size > 1 and config.train_tokenizer:
|
| 260 |
+
raise SystemExit(
|
| 261 |
+
"Under torchrun, run the tokenizer phase first as a single process\n"
|
| 262 |
+
" python train_rough.py --config <cfg with train_basemodel: false>\n"
|
| 263 |
+
"then launch torchrun with experiment.train_tokenizer: false."
|
| 264 |
+
)
|
| 265 |
+
if config.train_tokenizer:
|
| 266 |
+
run_tokenizer_phase(config)
|
| 267 |
+
else:
|
| 268 |
+
print("experiment.train_tokenizer = false -> skipping tokenizer phase")
|
| 269 |
+
if config.train_basemodel:
|
| 270 |
+
run_predictor_phase(args.config, config, rough_cfg)
|
| 271 |
+
|
| 272 |
+
|
| 273 |
+
if __name__ == "__main__":
|
| 274 |
+
main()
|
roughness_lab/roughness.py
ADDED
|
@@ -0,0 +1,187 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Surface-roughness analysis for price series.
|
| 2 |
+
|
| 3 |
+
Treats log-price as a measured surface profile (ISO 4287 / ISO 21920 spirit):
|
| 4 |
+
a Gaussian profile filter (ISO 16610-21) splits the profile into a smooth
|
| 5 |
+
*waviness* component (trend) and a *roughness* residual (texture), and the
|
| 6 |
+
standard amplitude/spacing/shape parameters are computed on those components.
|
| 7 |
+
|
| 8 |
+
All profile values are log-price multiplied by 100, so every amplitude
|
| 9 |
+
parameter reads directly in percent: Ra = 0.35 means the price wiggles an
|
| 10 |
+
average of 0.35% around its local trend. The cutoff wavelength lambda_c is
|
| 11 |
+
expressed in bars.
|
| 12 |
+
"""
|
| 13 |
+
from dataclasses import dataclass
|
| 14 |
+
|
| 15 |
+
import numpy as np
|
| 16 |
+
|
| 17 |
+
# ISO 16610-21 Gaussian weighting constant: 50% transmission at lambda_c.
|
| 18 |
+
_ALPHA = float(np.sqrt(np.log(2.0) / np.pi))
|
| 19 |
+
|
| 20 |
+
|
| 21 |
+
def gaussian_filter(profile: np.ndarray, cutoff: int) -> np.ndarray:
|
| 22 |
+
"""Return the waviness (low-pass mean line) of a profile.
|
| 23 |
+
|
| 24 |
+
Implements the ISO 16610-21 Gaussian profile filter by direct
|
| 25 |
+
convolution, with reflected ends (the standard's end-effect zone is
|
| 26 |
+
handled by reflection rather than truncation).
|
| 27 |
+
"""
|
| 28 |
+
z = np.asarray(profile, dtype=np.float64)
|
| 29 |
+
if cutoff < 2 or len(z) < 4:
|
| 30 |
+
return z.copy()
|
| 31 |
+
half = int(cutoff) # kernel support +/- lambda_c (weight ~1e-7 at the edge)
|
| 32 |
+
x = np.arange(-half, half + 1, dtype=np.float64)
|
| 33 |
+
s = np.exp(-np.pi * (x / (_ALPHA * cutoff)) ** 2)
|
| 34 |
+
s /= s.sum()
|
| 35 |
+
padded = np.pad(z, half, mode="reflect")
|
| 36 |
+
return np.convolve(padded, s, mode="valid")
|
| 37 |
+
|
| 38 |
+
|
| 39 |
+
@dataclass
|
| 40 |
+
class RoughnessParams:
|
| 41 |
+
ra: float # arithmetic mean deviation of roughness profile (%)
|
| 42 |
+
rq: float # RMS deviation (%)
|
| 43 |
+
rz: float # mean peak-to-valley over sampling lengths (%)
|
| 44 |
+
rsm: float # mean spacing of profile elements (bars)
|
| 45 |
+
rsk: float # skewness of roughness profile (dimensionless)
|
| 46 |
+
rku: float # kurtosis of roughness profile (dimensionless, Pearson)
|
| 47 |
+
wa: float # arithmetic mean deviation of form-removed waviness (%)
|
| 48 |
+
|
| 49 |
+
def as_dict(self) -> dict:
|
| 50 |
+
return {k: float(v) for k, v in self.__dict__.items()}
|
| 51 |
+
|
| 52 |
+
|
| 53 |
+
def _rsm(r: np.ndarray, rq: float) -> float:
|
| 54 |
+
"""Mean spacing between profile elements: distance between successive
|
| 55 |
+
upward mean-line crossings, with a +/-10%-of-Rq hysteresis band so
|
| 56 |
+
micro-wiggles do not register as elements (ISO height discrimination)."""
|
| 57 |
+
if rq <= 0:
|
| 58 |
+
return float("nan")
|
| 59 |
+
band = 0.1 * rq
|
| 60 |
+
crossings = []
|
| 61 |
+
armed = r[0] < -band
|
| 62 |
+
for i in range(1, len(r)):
|
| 63 |
+
if r[i] < -band:
|
| 64 |
+
armed = True
|
| 65 |
+
elif armed and r[i] > band:
|
| 66 |
+
crossings.append(i)
|
| 67 |
+
armed = False
|
| 68 |
+
if len(crossings) < 2:
|
| 69 |
+
return float(len(r)) # fewer than two elements: spacing ~ window size
|
| 70 |
+
return float(np.mean(np.diff(crossings)))
|
| 71 |
+
|
| 72 |
+
|
| 73 |
+
def _rz(r: np.ndarray, cutoff: int) -> float:
|
| 74 |
+
"""Mean peak-to-valley height over consecutive sampling lengths of
|
| 75 |
+
lambda_c bars (ISO evaluates five; we use as many whole ones as fit)."""
|
| 76 |
+
n_seg = max(1, len(r) // cutoff)
|
| 77 |
+
heights = []
|
| 78 |
+
for k in range(n_seg):
|
| 79 |
+
seg = r[k * cutoff:(k + 1) * cutoff]
|
| 80 |
+
if len(seg) >= 2:
|
| 81 |
+
heights.append(seg.max() - seg.min())
|
| 82 |
+
return float(np.mean(heights)) if heights else float("nan")
|
| 83 |
+
|
| 84 |
+
|
| 85 |
+
def roughness_params(close: np.ndarray, cutoff: int) -> RoughnessParams:
|
| 86 |
+
"""Compute the parameter set for a close-price window at one cutoff."""
|
| 87 |
+
z = 100.0 * np.log(np.asarray(close, dtype=np.float64)) # percent units
|
| 88 |
+
w = gaussian_filter(z, cutoff)
|
| 89 |
+
r = z - w
|
| 90 |
+
|
| 91 |
+
ra = float(np.mean(np.abs(r)))
|
| 92 |
+
rq = float(np.sqrt(np.mean(r ** 2)))
|
| 93 |
+
if rq > 0:
|
| 94 |
+
rsk = float(np.mean(r ** 3) / rq ** 3)
|
| 95 |
+
rku = float(np.mean(r ** 4) / rq ** 4)
|
| 96 |
+
else:
|
| 97 |
+
rsk, rku = 0.0, 0.0
|
| 98 |
+
|
| 99 |
+
# Waviness amplitude after form removal (least-squares line = "form")
|
| 100 |
+
x = np.arange(len(w), dtype=np.float64)
|
| 101 |
+
coef = np.polyfit(x, w, 1)
|
| 102 |
+
wa = float(np.mean(np.abs(w - np.polyval(coef, x))))
|
| 103 |
+
|
| 104 |
+
return RoughnessParams(
|
| 105 |
+
ra=ra, rq=rq, rz=_rz(r, cutoff), rsm=_rsm(r, rq),
|
| 106 |
+
rsk=rsk, rku=rku, wa=wa,
|
| 107 |
+
)
|
| 108 |
+
|
| 109 |
+
|
| 110 |
+
def fingerprint(close: np.ndarray, cutoffs: list) -> dict:
|
| 111 |
+
"""Multi-scale texture fingerprint: parameters across cutoff wavelengths."""
|
| 112 |
+
return {c: roughness_params(close, c).as_dict() for c in cutoffs}
|
| 113 |
+
|
| 114 |
+
|
| 115 |
+
# ---------------------------------------------------------------------------
|
| 116 |
+
# Texture-match scoring (forecast realism)
|
| 117 |
+
# ---------------------------------------------------------------------------
|
| 118 |
+
# Parameters compared via |log ratio| (scale-free); skewness via |difference|
|
| 119 |
+
# because its sign legitimately straddles zero.
|
| 120 |
+
_LOG_RATIO_KEYS = ("ra", "rq", "rz", "rsm", "rku", "wa")
|
| 121 |
+
_DIFF_KEYS = ("rsk",)
|
| 122 |
+
|
| 123 |
+
|
| 124 |
+
def texture_error(pred: dict, real: dict) -> float:
|
| 125 |
+
"""Mean texture mismatch between two parameter dicts (lower is better)."""
|
| 126 |
+
errs = []
|
| 127 |
+
for k in _LOG_RATIO_KEYS:
|
| 128 |
+
p, r = pred.get(k), real.get(k)
|
| 129 |
+
if p and r and p > 0 and r > 0 and np.isfinite(p) and np.isfinite(r):
|
| 130 |
+
errs.append(abs(np.log(p / r)))
|
| 131 |
+
for k in _DIFF_KEYS:
|
| 132 |
+
p, r = pred.get(k), real.get(k)
|
| 133 |
+
if p is not None and r is not None and np.isfinite(p) and np.isfinite(r):
|
| 134 |
+
errs.append(abs(p - r))
|
| 135 |
+
return float(np.mean(errs)) if errs else float("nan")
|
| 136 |
+
|
| 137 |
+
|
| 138 |
+
if __name__ == "__main__":
|
| 139 |
+
# ------------------------------------------------------------------
|
| 140 |
+
# Self-test on synthetic profiles with known properties.
|
| 141 |
+
# ------------------------------------------------------------------
|
| 142 |
+
rng = np.random.default_rng(7)
|
| 143 |
+
n = 4000
|
| 144 |
+
|
| 145 |
+
# 1) Pure sine "price": period 32 bars, log-amplitude 1% -> after a
|
| 146 |
+
# cutoff well above the period, roughness keeps the sine:
|
| 147 |
+
# Ra = 2A/pi, Rq = A/sqrt(2), Rz ~ 2A, RSm = period, Rku = 1.5.
|
| 148 |
+
A = 1.0 # percent
|
| 149 |
+
t = np.arange(n)
|
| 150 |
+
sine_price = np.exp(A / 100.0 * np.sin(2 * np.pi * t / 32))
|
| 151 |
+
p = roughness_params(sine_price, cutoff=128)
|
| 152 |
+
print("sine: Ra=%.3f (exp %.3f) Rq=%.3f (exp %.3f) Rz=%.3f (exp ~%.1f) "
|
| 153 |
+
"RSm=%.1f (exp 32) Rku=%.2f (exp 1.50)"
|
| 154 |
+
% (p.ra, 2 * A / np.pi, p.rq, A / np.sqrt(2), p.rz, 2 * A, p.rsm, p.rku))
|
| 155 |
+
assert abs(p.ra - 2 * A / np.pi) < 0.02
|
| 156 |
+
assert abs(p.rq - A / np.sqrt(2)) < 0.02
|
| 157 |
+
assert abs(p.rsm - 32) < 1.0
|
| 158 |
+
assert abs(p.rku - 1.5) < 0.05
|
| 159 |
+
|
| 160 |
+
# 2) Gaussian white noise on log-price, sigma=0.5%: Rq ~ sigma (part of
|
| 161 |
+
# the variance moves into waviness, so slightly below), Rku ~ 3.
|
| 162 |
+
sigma = 0.5
|
| 163 |
+
noise_price = np.exp(sigma / 100.0 * rng.standard_normal(n))
|
| 164 |
+
p = roughness_params(noise_price, cutoff=64)
|
| 165 |
+
print("noise: Rq=%.3f (exp ~%.2f) Rku=%.2f (exp ~3) Rsk=%.2f (exp ~0)"
|
| 166 |
+
% (p.rq, sigma, p.rku, p.rsk))
|
| 167 |
+
assert abs(p.rq - sigma) < 0.05
|
| 168 |
+
assert abs(p.rku - 3.0) < 0.3
|
| 169 |
+
assert abs(p.rsk) < 0.2
|
| 170 |
+
|
| 171 |
+
# 3) Sine + trend: waviness should absorb the trend; form removal makes
|
| 172 |
+
# Wa reflect only long undulations, and a long-period sine (256 bars)
|
| 173 |
+
# lands in waviness at cutoff 64 (transmission to roughness ~0).
|
| 174 |
+
slow = np.exp(0.05 * t / n + 2.0 / 100.0 * np.sin(2 * np.pi * t / 256))
|
| 175 |
+
p = roughness_params(slow, cutoff=64)
|
| 176 |
+
print("slow: Ra=%.4f (exp ~0) Wa=%.3f (exp ~%.3f)"
|
| 177 |
+
% (p.ra, p.wa, 2 * 2.0 / np.pi))
|
| 178 |
+
assert p.ra < 0.1
|
| 179 |
+
assert abs(p.wa - 2 * 2.0 / np.pi) < 0.15
|
| 180 |
+
|
| 181 |
+
# 4) texture_error: identical dicts -> 0; doubled Ra -> ln2 contribution.
|
| 182 |
+
d = p.as_dict()
|
| 183 |
+
assert texture_error(d, d) == 0.0
|
| 184 |
+
d2 = dict(d, ra=d["ra"] * 2)
|
| 185 |
+
assert texture_error(d2, d) > 0
|
| 186 |
+
print("texture_error self-test OK")
|
| 187 |
+
print("ALL SELF-TESTS PASSED")
|
run_prediction.py
ADDED
|
@@ -0,0 +1,116 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Run a Kronos-small forecast end-to-end on the K-line data bundled with the repo.
|
| 2 |
+
|
| 3 |
+
Loads NeoQuasar/Kronos-Tokenizer-base + NeoQuasar/Kronos-small from Hugging Face,
|
| 4 |
+
predicts 120 five-minute bars from a 400-bar context, then compares the forecast
|
| 5 |
+
against the held-out ground truth and saves a plot + CSV.
|
| 6 |
+
"""
|
| 7 |
+
import random
|
| 8 |
+
import sys
|
| 9 |
+
from pathlib import Path
|
| 10 |
+
|
| 11 |
+
import matplotlib
|
| 12 |
+
|
| 13 |
+
matplotlib.use("Agg") # headless: save to file instead of opening a window
|
| 14 |
+
import matplotlib.pyplot as plt
|
| 15 |
+
import numpy as np
|
| 16 |
+
import pandas as pd
|
| 17 |
+
import torch
|
| 18 |
+
|
| 19 |
+
REPO_ROOT = Path(__file__).resolve().parent / "Kronos"
|
| 20 |
+
sys.path.insert(0, str(REPO_ROOT))
|
| 21 |
+
|
| 22 |
+
from model import Kronos, KronosTokenizer, KronosPredictor
|
| 23 |
+
|
| 24 |
+
DATA_PATH = REPO_ROOT / "tests" / "data" / "regression_input.csv"
|
| 25 |
+
OUT_DIR = Path(__file__).resolve().parent / "output"
|
| 26 |
+
LOOKBACK = 400
|
| 27 |
+
PRED_LEN = 120
|
| 28 |
+
SEED = 123
|
| 29 |
+
|
| 30 |
+
|
| 31 |
+
def set_seed(seed: int) -> None:
|
| 32 |
+
random.seed(seed)
|
| 33 |
+
np.random.seed(seed)
|
| 34 |
+
torch.manual_seed(seed)
|
| 35 |
+
|
| 36 |
+
|
| 37 |
+
def main() -> None:
|
| 38 |
+
set_seed(SEED)
|
| 39 |
+
OUT_DIR.mkdir(exist_ok=True)
|
| 40 |
+
|
| 41 |
+
print("Loading tokenizer and model from Hugging Face Hub...")
|
| 42 |
+
tokenizer = KronosTokenizer.from_pretrained("NeoQuasar/Kronos-Tokenizer-base")
|
| 43 |
+
model = Kronos.from_pretrained("NeoQuasar/Kronos-small")
|
| 44 |
+
tokenizer.eval()
|
| 45 |
+
model.eval()
|
| 46 |
+
n_params = sum(p.numel() for p in model.parameters())
|
| 47 |
+
print(f"Model loaded: Kronos-small ({n_params / 1e6:.1f}M params)")
|
| 48 |
+
|
| 49 |
+
predictor = KronosPredictor(model, tokenizer, device="cpu", max_context=512)
|
| 50 |
+
|
| 51 |
+
df = pd.read_csv(DATA_PATH, parse_dates=["timestamps"])
|
| 52 |
+
print(f"Data: {DATA_PATH.name}, {len(df)} rows, "
|
| 53 |
+
f"{df['timestamps'].iloc[0]} .. {df['timestamps'].iloc[-1]}")
|
| 54 |
+
|
| 55 |
+
x_df = df.loc[:LOOKBACK - 1, ["open", "high", "low", "close", "volume", "amount"]]
|
| 56 |
+
x_timestamp = df.loc[:LOOKBACK - 1, "timestamps"]
|
| 57 |
+
y_timestamp = df.loc[LOOKBACK:LOOKBACK + PRED_LEN - 1, "timestamps"]
|
| 58 |
+
|
| 59 |
+
print(f"Forecasting {PRED_LEN} bars from a {LOOKBACK}-bar context (CPU)...")
|
| 60 |
+
pred_df = predictor.predict(
|
| 61 |
+
df=x_df,
|
| 62 |
+
x_timestamp=x_timestamp,
|
| 63 |
+
y_timestamp=y_timestamp,
|
| 64 |
+
pred_len=PRED_LEN,
|
| 65 |
+
T=1.0,
|
| 66 |
+
top_p=0.9,
|
| 67 |
+
sample_count=1,
|
| 68 |
+
verbose=True,
|
| 69 |
+
)
|
| 70 |
+
|
| 71 |
+
print("\nForecasted Data Head:")
|
| 72 |
+
print(pred_df.head())
|
| 73 |
+
|
| 74 |
+
# Compare against held-out ground truth
|
| 75 |
+
truth_df = df.loc[LOOKBACK:LOOKBACK + PRED_LEN - 1].set_index("timestamps")
|
| 76 |
+
price_cols = ["open", "high", "low", "close"]
|
| 77 |
+
mae = np.mean(np.abs(pred_df[price_cols].values - truth_df[price_cols].values))
|
| 78 |
+
mape = np.mean(
|
| 79 |
+
np.abs(pred_df[price_cols].values - truth_df[price_cols].values)
|
| 80 |
+
/ truth_df[price_cols].values
|
| 81 |
+
) * 100
|
| 82 |
+
print(f"\nPrice MAE vs ground truth: {mae:.4f}")
|
| 83 |
+
print(f"Price MAPE vs ground truth: {mape:.2f}%")
|
| 84 |
+
|
| 85 |
+
pred_csv = OUT_DIR / "kronos_small_forecast.csv"
|
| 86 |
+
pred_df.to_csv(pred_csv, index_label="timestamps")
|
| 87 |
+
|
| 88 |
+
# Plot: history + forecast vs ground truth
|
| 89 |
+
hist = df.loc[:LOOKBACK - 1].set_index("timestamps")
|
| 90 |
+
fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 7), sharex=True)
|
| 91 |
+
|
| 92 |
+
ax1.plot(hist.index, hist["close"], color="gray", linewidth=1, label="History")
|
| 93 |
+
ax1.plot(truth_df.index, truth_df["close"], color="blue", linewidth=1.5, label="Ground Truth")
|
| 94 |
+
ax1.plot(pred_df.index, pred_df["close"], color="red", linewidth=1.5, label="Kronos-small Forecast")
|
| 95 |
+
ax1.set_ylabel("Close Price")
|
| 96 |
+
ax1.legend(loc="best")
|
| 97 |
+
ax1.grid(True, alpha=0.4)
|
| 98 |
+
ax1.set_title(f"Kronos-small: {PRED_LEN}-step forecast ({LOOKBACK}-bar context)")
|
| 99 |
+
|
| 100 |
+
ax2.plot(hist.index, hist["volume"], color="gray", linewidth=1, label="History")
|
| 101 |
+
ax2.plot(truth_df.index, truth_df["volume"], color="blue", linewidth=1.5, label="Ground Truth")
|
| 102 |
+
ax2.plot(pred_df.index, pred_df["volume"], color="red", linewidth=1.5, label="Kronos-small Forecast")
|
| 103 |
+
ax2.set_ylabel("Volume")
|
| 104 |
+
ax2.legend(loc="best")
|
| 105 |
+
ax2.grid(True, alpha=0.4)
|
| 106 |
+
|
| 107 |
+
plt.tight_layout()
|
| 108 |
+
plot_path = OUT_DIR / "kronos_small_forecast.png"
|
| 109 |
+
plt.savefig(plot_path, dpi=150)
|
| 110 |
+
|
| 111 |
+
print(f"\nSaved forecast CSV to {pred_csv}")
|
| 112 |
+
print(f"Saved plot to {plot_path}")
|
| 113 |
+
|
| 114 |
+
|
| 115 |
+
if __name__ == "__main__":
|
| 116 |
+
main()
|