Skip to content

Add US stock 5-min prediction script using yfinance#310

Open
akommos88-source wants to merge 2 commits into
shiyu-coder:masterfrom
akommos88-source:feature/us-stock-prediction-script
Open

Add US stock 5-min prediction script using yfinance#310
akommos88-source wants to merge 2 commits into
shiyu-coder:masterfrom
akommos88-source:feature/us-stock-prediction-script

Conversation

@akommos88-source
Copy link
Copy Markdown

What changed

Adds predict_us_stock.py — a self-contained script for running Kronos inference on US equities using live 5-minute OHLCV data fetched from Yahoo Finance.

Why

The existing examples in examples/ rely on pre-downloaded CSV files from Chinese markets (XSHG). This script makes it easy for users with no existing data pipeline to run Kronos on any US stock ticker out of the box.

What it does

  1. Fetches data — pulls up to 30 days of 5-min bars via yfinance, strips pre/post market hours (09:30–16:00 ET), and selects the most recent lookback + pred_len candles
  2. Loads models — downloads KronosTokenizer + Kronos from HuggingFace Hub on first run (cached after); supports local paths via --tokenizer / --model flags
  3. Runs inference — uses KronosPredictor.predict() with configurable temperature, top-p, and Monte Carlo sample count (averaged for smoother output)
  4. Outputs results — prints predicted OHLCV table + directional summary (UP/DOWN % vs last historical close), saves a two-panel close + volume PNG chart

Usage

pip install yfinance  # only additional dependency beyond requirements.txt
python predict_us_stock.py --ticker AAPL --lookback 200 --pred_len 20
python predict_us_stock.py --ticker TSLA --lookback 300 --pred_len 10 --samples 10 --temp 0.8

Notes for reviewer

  • The amount column required by KronosPredictor is not provided by yfinance; the predictor handles this gracefully (fills with zeros) when the column is absent — no changes to model code needed.
  • No existing files were modified; this is a pure addition.
  • The pretrained models were trained primarily on global/Asian exchange data, so directional accuracy on US stocks will be limited without fine-tuning via finetune_csv/.

Adam El Kommos and others added 2 commits May 30, 2026 14:44
Adds predict_us_stock.py — a self-contained script for running Kronos
inference on US equities using 5-minute OHLCV data fetched live from
Yahoo Finance via yfinance.

Features:
- Fetches 5-min bars automatically for any ticker (yf.download)
- Strips pre/post market hours (09:30–16:00 ET)
- Loads KronosTokenizer + Kronos from HuggingFace Hub (or local path)
- Auto-detects CUDA / MPS / CPU device
- Monte Carlo sampling with configurable temperature, top-p, sample count
- Prints directional summary (UP/DOWN %, predicted close)
- Saves a two-panel close+volume PNG chart

Usage:
  python predict_us_stock.py --ticker AAPL --lookback 200 --pred_len 20

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Three new files under finetune_csv/:

- fetch_us_data.py: fetches 5-min OHLCV from yfinance for any US tickers,
  strips pre/post market hours, derives the 'amount' column, and saves a
  CSV in the format expected by CustomKlineDataset

- configs/config_us_5min.yaml: ready-to-use YAML config tuned for US 5-min
  data (lower LRs to avoid catastrophic forgetting, 200-bar lookback,
  20-bar prediction horizon, HuggingFace Hub model IDs)

- run_us_finetune.py: end-to-end runner — fetches data, patches the config
  with the resolved data path, runs finetune_tokenizer.py then
  finetune_base_model.py in sequence; individual stages can be skipped via
  --skip_tokenizer / --skip_predictor flags

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant