Add US stock 5-min prediction script using yfinance#310
Open
akommos88-source wants to merge 2 commits into
Open
Add US stock 5-min prediction script using yfinance#310akommos88-source wants to merge 2 commits into
akommos88-source wants to merge 2 commits into
Conversation
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>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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
yfinance, strips pre/post market hours (09:30–16:00 ET), and selects the most recentlookback + pred_lencandlesKronosTokenizer+Kronosfrom HuggingFace Hub on first run (cached after); supports local paths via--tokenizer/--modelflagsKronosPredictor.predict()with configurable temperature, top-p, and Monte Carlo sample count (averaged for smoother output)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.8Notes for reviewer
amountcolumn required byKronosPredictoris not provided by yfinance; the predictor handles this gracefully (fills with zeros) when the column is absent — no changes to model code needed.finetune_csv/.