简介
Kronos 是清华大学信息科学研究院开源的 K 线分析模型，基于全球 45 个交易所的海量数据训练，支持沪深市场复权数据。该模型提供多个版本以适配不同的算力和应用需求。
| 模型名称 | Tokenizer | 上下文长度 | 参数量 | 开源情况 |
|---|---|---|---|---|
| Kronos-mini | Tokenizer-2k | 2048 | 4.1M | NeoQuasar/Kronos-mini |
| Kronos-small | Tokenizer-base | 512 | 24.7M | NeoQuasar/Kronos-small |
| Kronos-base | Tokenizer-base | 512 | 102.3M | NeoQuasar/Kronos-base |
| Kronos-large | Tokenizer-base | 512 | 499.2M | 暂未开放 |
安装与使用
1. 安装依赖
确保已安装 Python 3.10+，并运行：
pip install -r requirements.txt
2. 加载模型与 Tokenizer
from model import Kronos, KronosTokenizer, KronosPredictor
tokenizer = KronosTokenizer.from_pretrained("NeoQuasar/Kronos-Tokenizer-base")
model = Kronos.from_pretrained("NeoQuasar/Kronos-small")
3. 初始化预测器
predictor = KronosPredictor(model, tokenizer, device="cuda:0", max_context=512)
4. 准备输入数据
输入需要包含历史的 K 线数据：
import pandas as pd
df = pd.read_csv("./data/XSHG_5min_600977.csv")
df['timestamps'] = pd.to_datetime(df['timestamps'])
lookback = 400
pred_len = 120
x_df = df.loc[:lookback-1, ['open', 'high', 'low', 'close', , ]]
x_timestamp = df.loc[:lookback-, ]
y_timestamp = df.loc[lookback:lookback+pred_len-, ]
