Deploy Kronos BTC Prediction API

- FastAPI server with /predict and /signal endpoints
- Kronos model (iter5, 62.4% win rate, 3.25 Sharpe ratio)
- Python client SDK
- Docker configuration for HuggingFace Spaces

DEPLOYMENT.md

@@ -0,0 +1,217 @@
+# HuggingFace Space 部署指南
+
+本指南介绍如何将 Kronos BTC 预测 API 部署到 HuggingFace Spaces。
+
+## 准备工作
+
+### 1. 创建 HuggingFace 账户
+
+如果还没有账户，请访问 https://huggingface.co/join 注册。
+
+### 2. 安装 HuggingFace CLI
+
+```bash
+pip install huggingface_hub
+huggingface-cli login
+```
+
+## 方法一：通过 Git 部署 (推荐)
+
+### 1. 创建新 Space
+
+访问 https://huggingface.co/new-space 创建新 Space：
+
+- **Space name**: `kronos-btc-predictor` (或任意名称)
+- **License**: MIT
+- **SDK**: Docker
+- **Hardware**: CPU basic (免费)
+
+### 2. 克隆 Space 仓库
+
+```bash
+git clone https://huggingface.co/spaces/YOUR_USERNAME/kronos-btc-predictor
+cd kronos-btc-predictor
+```
+
+### 3. 复制文件
+
+```bash
+# 复制所有文件到 Space 仓库
+cp -r /path/to/hf_space/* .
+
+# 文件结构应该是:
+# ├── app.py
+# ├── requirements.txt
+# ├── README.md
+# ├── client.py
+# ├── Dockerfile         # 需要创建
+# ├── model/
+# │   ├── __init__.py
+# │   ├── kronos.py
+# │   └── module.py
+# └── models/
+#     ├── tokenizer/
+#     │   ├── config.json
+#     │   └── model.safetensors
+#     └── predictor/
+#         ├── config.json
+#         └── model.safetensors
+```
+
+### 4. 创建 Dockerfile
+
+```dockerfile
+FROM python:3.10-slim
+
+WORKDIR /app
+
+# 安装依赖
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+# 复制应用代码
+COPY . .
+
+# 暴露端口
+EXPOSE 7860
+
+# 启动服务
+CMD ["python", "app.py"]
+```
+
+### 5. 推送到 HuggingFace
+
+```bash
+git add .
+git commit -m "Initial deployment"
+git push
+```
+
+### 6. 等待构建
+
+Space 会自动构建和部署。你可以在 Space 页面查看构建日志。
+
+构建完成后，API 将在以下地址可用：
+```
+https://YOUR_USERNAME-kronos-btc-predictor.hf.space
+```
+
+## 方法二：通过 Web 界面上传
+
+### 1. 创建 Space
+
+访问 https://huggingface.co/new-space：
+- SDK: Docker
+- Hardware: CPU basic
+
+### 2. 上传文件
+
+在 Space 页面点击 "Files" 标签，然后 "Add file" -> "Upload files"：
+
+逐个上传以下文件：
+- `app.py`
+- `requirements.txt`
+- `Dockerfile`
+- `model/__init__.py`
+- `model/kronos.py`
+- `model/module.py`
+- `models/tokenizer/config.json`
+- `models/tokenizer/model.safetensors`
+- `models/predictor/config.json`
+- `models/predictor/model.safetensors`
+
+## 验证部署
+
+### 1. 健康检查
+
+```bash
+curl https://YOUR_USERNAME-kronos-btc-predictor.hf.space/health
+```
+
+预期响应：
+```json
+{
+  "status": "healthy",
+  "model_loaded": true,
+  "model_version": "iter5 (converged)",
+  "device": "cpu"
+}
+```
+
+### 2. API 文档
+
+访问 Swagger UI：
+```
+https://YOUR_USERNAME-kronos-btc-predictor.hf.space/docs
+```
+
+### 3. 测试预测
+
+```python
+from client import KronosClient
+
+client = KronosClient("https://YOUR_USERNAME-kronos-btc-predictor.hf.space")
+health = client.health()
+print(f"Status: {health.status}")
+```
+
+## 配置自定义域名
+
+1. 在 Space 设置中找到 "Custom domain"
+2. 输入你的域名 (如 `api.yourdomain.com`)
+3. 配置 DNS CNAME 记录指向 HuggingFace
+
+## 注意事项
+
+### 免费版限制
+
+- **CPU**: 2 vCPU
+- **内存**: 16GB RAM
+- **存储**: 50GB
+- **请求**: 无硬性限制，但有速率控制
+- **冷启动**: 不活动时会休眠，首次请求需等待约 30-60 秒
+
+### 性能优化
+
+1. **减少 n_paths**: 使用 10-20 个路径而不是 30-100
+2. **减少 pred_len**: 使用 12-24 而不是 72
+3. **预热**: 定期发送健康检查请求防止休眠
+
+### 安全建议
+
+1. 不要在代码中硬编码 API 密钥
+2. 使用 HuggingFace Secrets 存储敏感信息
+3. 考虑添加请求速率限制
+
+## 升级到 Pro
+
+如果需要更好的性能，可以升级到 HuggingFace Pro：
+
+- **CPU upgrade**: 更快的 CPU
+- **GPU**: T4 GPU (付费)
+- **永不休眠**: 始终保持运行
+
+访问 https://huggingface.co/pricing 了解详情。
+
+## 故障排除
+
+### 构建失败
+
+1. 检查 `requirements.txt` 中的版本兼容性
+2. 确保所有文件都已上传
+3. 查看构建日志中的错误信息
+
+### 模型加载失败
+
+1. 确认 `models/` 目录结构正确
+2. 检查 `config.json` 和 `model.safetensors` 文件
+
+### 请求超时
+
+1. 减少 `n_paths` 和 `pred_len` 参数
+2. 检查输入数据大小
+3. 考虑升级到更好的硬件
+
+## 联系支持
+
+如有问题，请在项目仓库提交 Issue。

Dockerfile

@@ -0,0 +1,40 @@
+# Kronos BTC Prediction API - Docker Image
+# Optimized for HuggingFace Spaces
+
+FROM python:3.10-slim
+
+# Set working directory
+WORKDIR /app
+
+# Install system dependencies
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    build-essential \
+    && rm -rf /var/lib/apt/lists/*
+
+# Copy requirements first for better caching
+COPY requirements.txt .
+
+# Install Python dependencies
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Copy application code
+COPY . .
+
+# Create non-root user for security
+RUN useradd -m -u 1000 user
+USER user
+
+# Set environment variables
+ENV HOME=/home/user \
+    PATH=/home/user/.local/bin:$PATH \
+    PYTHONUNBUFFERED=1
+
+# Expose port (HuggingFace Spaces uses 7860)
+EXPOSE 7860
+
+# Health check
+HEALTHCHECK --interval=30s --timeout=10s --start-period=60s --retries=3 \
+    CMD python -c "import httpx; httpx.get('http://localhost:7860/health', timeout=5)" || exit 1
+
+# Start the application
+CMD ["python", "app.py"]

README.md

@@ -1,12 +1,50 @@
 ---
-title: Tslm
-emoji: 🦀
-colorFrom: indigo
-colorTo: yellow
+title: TSLM - Time Series Language Model
+emoji: 📈
+colorFrom: blue
+colorTo: green
 sdk: docker
 pinned: false
-license: apache-2.0
-short_description: time series language modle
+license: mit
+short_description: BTC price prediction API based on Kronos model
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Kronos BTC Prediction API
+
+基于 Kronos 时序预测模型的 BTC 价格预测 API 服务。
+
+## 模型信息
+
+| 指标 | 值 |
+|------|-----|
+| 迭代版本 | iter5 (收敛) |
+| 胜率 | 62.4% |
+| 夏普比率 | 3.25 |
+| 盈亏比 | 1.47 |
+| 模型大小 | ~32MB |
+
+## API 端点
+
+- `GET /health` - 健康检查
+- `POST /predict` - 价格预测
+- `POST /signal` - 交易信号
+
+## 快速开始
+
+```python
+from client import KronosClient
+
+client = KronosClient("https://xianqiu-tslm.hf.space")
+
+# 价格预测
+prediction = client.predict(ohlcv_data, pred_len=24)
+print(f"上涨概率: {prediction.upside_probability:.1%}")
+
+# 交易信号
+signal = client.get_signal(ohlcv_data)
+print(f"信号: {signal.signal}")
+```
+
+## API 文档
+
+访问 `/docs` 查看完整的 Swagger API 文档。

app.py

@@ -0,0 +1,607 @@
+"""
+Kronos BTC Prediction API - HuggingFace Space
+基于 Kronos 模型的 BTC 价格预测 API 服务
+
+模型信息:
+- 迭代版本: iter5 (收敛)
+- 胜率: 62.4%
+- 夏普比率: 3.25
+- 盈亏比: 1.47
+
+API 端点:
+- GET /health - 健康检查
+- POST /predict - 预测 BTC 价格走势
+- POST /signal - 生成交易信号
+"""
+
+import os
+import logging
+from datetime import datetime, timedelta
+from typing import List, Optional
+from enum import Enum
+
+import numpy as np
+import pandas as pd
+import torch
+from fastapi import FastAPI, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel, Field
+
+from model import KronosTokenizer, Kronos, calc_time_stamps
+
+# 配置日志
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# 创建 FastAPI 应用
+app = FastAPI(
+    title="Kronos BTC Prediction API",
+    description="基于 Kronos 时序预测模型的 BTC 价格预测服务",
+    version="1.0.0",
+    docs_url="/docs",
+    redoc_url="/redoc"
+)
+
+# CORS 配置
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+
+# ==================== 数据模型 ====================
+
+class OHLCVData(BaseModel):
+    """OHLCV K线数据"""
+    timestamp: str = Field(..., description="时间戳 (ISO格式或毫秒)")
+    open: float = Field(..., description="开盘价")
+    high: float = Field(..., description="最高价")
+    low: float = Field(..., description="最低价")
+    close: float = Field(..., description="收盘价")
+    volume: float = Field(0.0, description="成交量")
+    amount: Optional[float] = Field(None, description="成交额 (可选)")
+
+
+class PredictRequest(BaseModel):
+    """预测请求"""
+    data: List[OHLCVData] = Field(..., description="历史 OHLCV 数据 (至少 100 条)")
+    pred_len: int = Field(24, ge=1, le=72, description="预测长度 (小时)")
+    n_paths: int = Field(30, ge=10, le=100, description="Monte Carlo 路径数")
+    temperature: float = Field(1.0, ge=0.1, le=2.0, description="采样温度")
+    top_p: float = Field(0.9, ge=0.5, le=1.0, description="Top-p 采样")
+
+
+class PredictResponse(BaseModel):
+    """预测响应"""
+    current_price: float
+    mean_forecast: float
+    min_forecast: float
+    max_forecast: float
+    upside_probability: float
+    expected_return: float
+    volatility_amplification: float
+    confidence: float
+    forecast_prices: List[float]
+    timestamp: str
+
+
+class SignalType(str, Enum):
+    """交易信号类型"""
+    STRONG_BUY = "STRONG_BUY"
+    BUY = "BUY"
+    HOLD = "HOLD"
+    SELL = "SELL"
+    STRONG_SELL = "STRONG_SELL"
+
+
+class SignalRequest(BaseModel):
+    """交易信号请求"""
+    data: List[OHLCVData] = Field(..., description="历史 OHLCV 数据")
+    buy_threshold: float = Field(0.58, ge=0.5, le=0.9, description="买入阈值")
+    sell_threshold: float = Field(0.42, ge=0.1, le=0.5, description="卖出阈值")
+    stop_loss: float = Field(0.03, ge=0.01, le=0.1, description="止损比例")
+    take_profit: float = Field(0.08, ge=0.02, le=0.2, description="止盈比例")
+    n_paths: int = Field(30, ge=10, le=100, description="Monte Carlo 路径数")
+
+
+class SignalResponse(BaseModel):
+    """交易信号响应"""
+    signal: SignalType
+    confidence: float
+    current_price: float
+    target_price: float
+    stop_loss_price: float
+    take_profit_price: float
+    upside_probability: float
+    expected_return: float
+    suggested_position_size: float
+    reason: str
+    timestamp: str
+
+
+class HealthResponse(BaseModel):
+    """健康检查响应"""
+    model_config = {"protected_namespaces": ()}
+    
+    status: str
+    model_loaded: bool
+    model_version: str
+    device: str
+    timestamp: str
+
+
+# ==================== 模型加载 ====================
+
+class KronosPredictor:
+    """Kronos 预测器"""
+    
+    def __init__(self):
+        self.tokenizer = None
+        self.model = None
+        self.device = "cpu"  # HuggingFace Space 免费版只有 CPU
+        self.max_context = 2048
+        self.clip = 5.0
+        self.loaded = False
+    
+    def load_models(self):
+        """加载模型"""
+        if self.loaded:
+            return
+        
+        logger.info("Loading Kronos models...")
+        
+        # 模型路径
+        tokenizer_path = os.path.join(os.path.dirname(__file__), "models", "tokenizer")
+        predictor_path = os.path.join(os.path.dirname(__file__), "models", "predictor")
+        
+        # 加载 tokenizer
+        logger.info(f"Loading tokenizer from {tokenizer_path}")
+        self.tokenizer = KronosTokenizer.from_pretrained(tokenizer_path)
+        self.tokenizer.to(self.device)
+        self.tokenizer.eval()
+        
+        # 加载 predictor
+        logger.info(f"Loading predictor from {predictor_path}")
+        self.model = Kronos.from_pretrained(predictor_path)
+        self.model.to(self.device)
+        self.model.eval()
+        
+        self.loaded = True
+        logger.info("Models loaded successfully!")
+    
+    def _sample_from_logits(self, logits, T, top_p):
+        """从 logits 采样"""
+        logits = logits / T
+        
+        if top_p < 1.0:
+            sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+            cumulative_probs = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1)
+            sorted_indices_to_remove = cumulative_probs > top_p
+            sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
+            sorted_indices_to_remove[..., 0] = 0
+            indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
+            logits[indices_to_remove] = float('-inf')
+        
+        probs = torch.softmax(logits, dim=-1)
+        return torch.multinomial(probs, num_samples=1)
+    
+    @torch.no_grad()
+    def predict(self, df: pd.DataFrame, pred_len: int = 24, n_paths: int = 30,
+                T: float = 1.0, top_p: float = 0.9) -> dict:
+        """
+        执行预测
+        
+        Args:
+            df: OHLCV DataFrame
+            pred_len: 预测长度
+            n_paths: Monte Carlo 路径数
+            T: 温度
+            top_p: Top-p 采样
+        
+        Returns:
+            预测结果字典
+        """
+        if not self.loaded:
+            self.load_models()
+        
+        df = df.copy()
+        
+        # 确保有 amount 列
+        if 'amount' not in df.columns:
+            df['amount'] = df['volume'] * df['close']
+        
+        # 时间戳处理
+        x_timestamp = pd.Series(df['timestamp'])
+        last_time = df['timestamp'].iloc[-1]
+        y_timestamp = pd.Series(pd.date_range(
+            start=last_time + timedelta(hours=1),
+            periods=pred_len,
+            freq='1h'
+        ))
+        
+        x_time_df = calc_time_stamps(x_timestamp)
+        y_time_df = calc_time_stamps(y_timestamp)
+        
+        # 特征
+        price_cols = ['open', 'high', 'low', 'close']
+        x = df[price_cols + ['volume', 'amount']].values.astype(np.float32)
+        x_stamp = x_time_df.values.astype(np.float32)
+        y_stamp = y_time_df.values.astype(np.float32)
+        
+        # 归一化
+        x_mean = np.mean(x, axis=0)
+        x_std = np.std(x, axis=0)
+        x_norm = (x - x_mean) / (x_std + 1e-5)
+        x_norm = np.clip(x_norm, -self.clip, self.clip)
+        
+        # 添加 batch 维度
+        x_norm = x_norm[np.newaxis, :]
+        x_stamp = x_stamp[np.newaxis, :]
+        y_stamp = y_stamp[np.newaxis, :]
+        
+        # 转换为 tensor
+        x_tensor = torch.from_numpy(x_norm).to(self.device)
+        x_stamp_tensor = torch.from_numpy(x_stamp).to(self.device)
+        y_stamp_tensor = torch.from_numpy(y_stamp).to(self.device)
+        
+        # Monte Carlo 采样
+        all_paths = self._generate_mc_paths(
+            x_tensor, x_stamp_tensor, y_stamp_tensor,
+            pred_len, T, top_p, n_paths
+        )
+        
+        # 反归一化
+        all_paths = all_paths[0]  # (n_paths, pred_len, 6)
+        all_paths = all_paths * (x_std + 1e-5) + x_mean
+        
+        # 计算指标
+        current_price = float(df['close'].iloc[-1])
+        close_paths = all_paths[:, :, 3]  # close price
+        
+        final_prices = close_paths[:, -1]
+        mean_forecast = float(np.mean(final_prices))
+        upside_prob = float(np.mean(final_prices > current_price))
+        
+        # 波动率放大
+        hist_returns = df['close'].pct_change().dropna().values[-24:]
+        hist_vol = np.std(hist_returns) if len(hist_returns) > 0 else 0.01
+        
+        pred_vols = []
+        for path in close_paths:
+            returns = np.diff(path) / path[:-1]
+            pred_vols.append(np.std(returns))
+        vol_amp = float(np.mean(np.array(pred_vols) > hist_vol))
+        
+        # 置信度
+        confidence = 1.0 - np.std(final_prices) / (np.mean(final_prices) + 1e-8)
+        confidence = float(max(0.0, min(1.0, confidence)))
+        
+        # 预期收益
+        expected_return = (mean_forecast - current_price) / current_price
+        
+        # 平均预测路径
+        mean_path = np.mean(close_paths, axis=0).tolist()
+        
+        return {
+            "current_price": current_price,
+            "mean_forecast": mean_forecast,
+            "min_forecast": float(np.min(final_prices)),
+            "max_forecast": float(np.max(final_prices)),
+            "upside_probability": upside_prob,
+            "expected_return": float(expected_return),
+            "volatility_amplification": vol_amp,
+            "confidence": confidence,
+            "forecast_prices": mean_path,
+            "timestamp": datetime.utcnow().isoformat()
+        }
+    
+    def _generate_mc_paths(self, x, x_stamp, y_stamp, pred_len, T, top_p, n_paths):
+        """生成 Monte Carlo 路径"""
+        x = torch.clip(x, -self.clip, self.clip)
+        
+        # 扩展为多个路径
+        x = x.unsqueeze(1).repeat(1, n_paths, 1, 1).reshape(-1, x.size(1), x.size(2))
+        x_stamp = x_stamp.unsqueeze(1).repeat(1, n_paths, 1, 1).reshape(-1, x_stamp.size(1), x_stamp.size(2))
+        y_stamp = y_stamp.unsqueeze(1).repeat(1, n_paths, 1, 1).reshape(-1, y_stamp.size(1), y_stamp.size(2))
+        
+        # 编码
+        x_token = self.tokenizer.encode(x, half=True)
+        
+        initial_seq_len = x.size(1)
+        batch_size = x_token[0].size(0)
+        total_seq_len = initial_seq_len + pred_len
+        full_stamp = torch.cat([x_stamp, y_stamp], dim=1)
+        
+        generated_pre = x_token[0].new_empty(batch_size, pred_len)
+        generated_post = x_token[1].new_empty(batch_size, pred_len)
+        
+        pre_buffer = x_token[0].new_zeros(batch_size, self.max_context)
+        post_buffer = x_token[1].new_zeros(batch_size, self.max_context)
+        buffer_len = min(initial_seq_len, self.max_context)
+        
+        if buffer_len > 0:
+            start_idx = max(0, initial_seq_len - self.max_context)
+            pre_buffer[:, :buffer_len] = x_token[0][:, start_idx:start_idx + buffer_len]
+            post_buffer[:, :buffer_len] = x_token[1][:, start_idx:start_idx + buffer_len]
+        
+        # 自回归生成
+        for i in range(pred_len):
+            current_seq_len = initial_seq_len + i
+            window_len = min(current_seq_len, self.max_context)
+            
+            if current_seq_len <= self.max_context:
+                input_tokens = [pre_buffer[:, :window_len], post_buffer[:, :window_len]]
+            else:
+                input_tokens = [pre_buffer, post_buffer]
+            
+            context_end = current_seq_len
+            context_start = max(0, context_end - self.max_context)
+            current_stamp = full_stamp[:, context_start:context_end, :].contiguous()
+            
+            s1_logits, context = self.model.decode_s1(input_tokens[0], input_tokens[1], current_stamp)
+            s1_logits = s1_logits[:, -1, :]
+            sample_pre = self._sample_from_logits(s1_logits, T, top_p)
+            
+            s2_logits = self.model.decode_s2(context, sample_pre)
+            s2_logits = s2_logits[:, -1, :]
+            sample_post = self._sample_from_logits(s2_logits, T, top_p)
+            
+            generated_pre[:, i] = sample_pre.squeeze(-1)
+            generated_post[:, i] = sample_post.squeeze(-1)
+            
+            if current_seq_len < self.max_context:
+                pre_buffer[:, current_seq_len] = sample_pre.squeeze(-1)
+                post_buffer[:, current_seq_len] = sample_post.squeeze(-1)
+            else:
+                pre_buffer = torch.roll(pre_buffer, shifts=-1, dims=1)
+                post_buffer = torch.roll(post_buffer, shifts=-1, dims=1)
+                pre_buffer[:, -1] = sample_pre.squeeze(-1)
+                post_buffer[:, -1] = sample_post.squeeze(-1)
+        
+        # 解码
+        full_pre = torch.cat([x_token[0], generated_pre], dim=1)
+        full_post = torch.cat([x_token[1], generated_post], dim=1)
+        
+        context_start = max(0, total_seq_len - self.max_context)
+        input_tokens = [
+            full_pre[:, context_start:total_seq_len].contiguous(),
+            full_post[:, context_start:total_seq_len].contiguous()
+        ]
+        z = self.tokenizer.decode(input_tokens, half=True)
+        
+        # 重塑
+        z = z.reshape(-1, n_paths, z.size(1), z.size(2))
+        preds = z.cpu().numpy()
+        
+        return preds[:, :, -pred_len:, :]
+    
+    def generate_signal(self, df: pd.DataFrame, buy_threshold: float = 0.58,
+                        sell_threshold: float = 0.42, stop_loss: float = 0.03,
+                        take_profit: float = 0.08, n_paths: int = 30) -> dict:
+        """
+        生成交易信号
+        """
+        # 获取预测
+        pred = self.predict(df, pred_len=24, n_paths=n_paths)
+        
+        current_price = pred["current_price"]
+        upside_prob = pred["upside_probability"]
+        confidence = pred["confidence"]
+        vol_amp = pred["volatility_amplification"]
+        
+        # 生成信号
+        if confidence < 0.3:
+            signal = SignalType.HOLD
+            reason = f"Low confidence ({confidence:.2f})"
+        elif upside_prob > 0.7 and vol_amp < 0.5:
+            signal = SignalType.STRONG_BUY
+            reason = f"Strong upside ({upside_prob:.1%}) with low volatility"
+        elif upside_prob > buy_threshold:
+            signal = SignalType.BUY
+            reason = f"Upside probability {upside_prob:.1%} > {buy_threshold:.0%}"
+        elif upside_prob < 0.3 and vol_amp < 0.5:
+            signal = SignalType.STRONG_SELL
+            reason = f"Strong downside ({1-upside_prob:.1%}) with low volatility"
+        elif upside_prob < sell_threshold:
+            signal = SignalType.SELL
+            reason = f"Upside probability {upside_prob:.1%} < {sell_threshold:.0%}"
+        else:
+            signal = SignalType.HOLD
+            reason = f"Neutral zone (upside: {upside_prob:.1%})"
+        
+        # 计算价格目标
+        if signal in [SignalType.BUY, SignalType.STRONG_BUY]:
+            target_price = pred["mean_forecast"]
+            stop_loss_price = current_price * (1 - stop_loss)
+            take_profit_price = current_price * (1 + take_profit)
+        elif signal in [SignalType.SELL, SignalType.STRONG_SELL]:
+            target_price = pred["mean_forecast"]
+            stop_loss_price = current_price * (1 + stop_loss)
+            take_profit_price = current_price * (1 - take_profit)
+        else:
+            target_price = current_price
+            stop_loss_price = current_price
+            take_profit_price = current_price
+        
+        # 动态仓位 (Kelly Criterion 简化版)
+        if signal != SignalType.HOLD:
+            win_prob = upside_prob if upside_prob > 0.5 else (1 - upside_prob)
+            odds = take_profit / stop_loss
+            kelly = (win_prob * odds - (1 - win_prob)) / odds
+            kelly = kelly * 0.5  # 半 Kelly
+            direction_strength = abs(upside_prob - 0.5) * 2
+            position_size = 0.1 * (1 + direction_strength) * confidence
+            position_size = max(0.02, min(0.3, position_size))
+            if kelly > 0:
+                position_size = min(position_size, kelly)
+            else:
+                position_size = 0.02
+        else:
+            position_size = 0.0
+        
+        return {
+            "signal": signal,
+            "confidence": confidence,
+            "current_price": current_price,
+            "target_price": target_price,
+            "stop_loss_price": stop_loss_price,
+            "take_profit_price": take_profit_price,
+            "upside_probability": upside_prob,
+            "expected_return": pred["expected_return"],
+            "suggested_position_size": position_size,
+            "reason": reason,
+            "timestamp": datetime.utcnow().isoformat()
+        }
+
+
+# 全局预测器实例
+predictor = KronosPredictor()
+
+
+# ==================== API 端点 ====================
+
+@app.on_event("startup")
+async def startup_event():
+    """启动时加载模型"""
+    logger.info("Starting Kronos BTC Prediction API...")
+    try:
+        predictor.load_models()
+    except Exception as e:
+        logger.error(f"Failed to load models: {e}")
+
+
+@app.get("/", response_model=HealthResponse)
+async def root():
+    """根路径 - 返回健康状态"""
+    return await health_check()
+
+
+@app.get("/health", response_model=HealthResponse)
+async def health_check():
+    """健康检查"""
+    return HealthResponse(
+        status="healthy" if predictor.loaded else "loading",
+        model_loaded=predictor.loaded,
+        model_version="iter5 (converged)",
+        device=predictor.device,
+        timestamp=datetime.utcnow().isoformat()
+    )
+
+
+@app.post("/predict", response_model=PredictResponse)
+async def predict(request: PredictRequest):
+    """
+    预测 BTC 价格走势
+    
+    输入至少 100 条历史 OHLCV 数据，返回未来价格预测。
+    """
+    if not predictor.loaded:
+        raise HTTPException(status_code=503, detail="Model not loaded yet")
+    
+    if len(request.data) < 100:
+        raise HTTPException(status_code=400, detail="At least 100 data points required")
+    
+    try:
+        # 转换数据
+        records = []
+        for d in request.data:
+            try:
+                # 处理时间戳
+                if d.timestamp.isdigit():
+                    ts = pd.Timestamp(int(d.timestamp), unit='ms')
+                else:
+                    ts = pd.Timestamp(d.timestamp)
+                
+                records.append({
+                    'timestamp': ts,
+                    'open': d.open,
+                    'high': d.high,
+                    'low': d.low,
+                    'close': d.close,
+                    'volume': d.volume,
+                    'amount': d.amount if d.amount else d.volume * d.close
+                })
+            except Exception as e:
+                raise HTTPException(status_code=400, detail=f"Invalid data format: {e}")
+        
+        df = pd.DataFrame(records)
+        df = df.sort_values('timestamp').reset_index(drop=True)
+        
+        # 执行预测
+        result = predictor.predict(
+            df,
+            pred_len=request.pred_len,
+            n_paths=request.n_paths,
+            T=request.temperature,
+            top_p=request.top_p
+        )
+        
+        return PredictResponse(**result)
+    
+    except Exception as e:
+        logger.error(f"Prediction error: {e}")
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+@app.post("/signal", response_model=SignalResponse)
+async def get_signal(request: SignalRequest):
+    """
+    生成交易信号
+    
+    基于历史数据和策略参数，生成买入/卖出/持有信号。
+    """
+    if not predictor.loaded:
+        raise HTTPException(status_code=503, detail="Model not loaded yet")
+    
+    if len(request.data) < 100:
+        raise HTTPException(status_code=400, detail="At least 100 data points required")
+    
+    try:
+        # 转换数据
+        records = []
+        for d in request.data:
+            try:
+                if d.timestamp.isdigit():
+                    ts = pd.Timestamp(int(d.timestamp), unit='ms')
+                else:
+                    ts = pd.Timestamp(d.timestamp)
+                
+                records.append({
+                    'timestamp': ts,
+                    'open': d.open,
+                    'high': d.high,
+                    'low': d.low,
+                    'close': d.close,
+                    'volume': d.volume,
+                    'amount': d.amount if d.amount else d.volume * d.close
+                })
+            except Exception as e:
+                raise HTTPException(status_code=400, detail=f"Invalid data format: {e}")
+        
+        df = pd.DataFrame(records)
+        df = df.sort_values('timestamp').reset_index(drop=True)
+        
+        # 生成信号
+        result = predictor.generate_signal(
+            df,
+            buy_threshold=request.buy_threshold,
+            sell_threshold=request.sell_threshold,
+            stop_loss=request.stop_loss,
+            take_profit=request.take_profit,
+            n_paths=request.n_paths
+        )
+        
+        return SignalResponse(**result)
+    
+    except Exception as e:
+        logger.error(f"Signal generation error: {e}")
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+# HuggingFace Spaces 入口
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)

client.py

@@ -0,0 +1,643 @@
+"""
+Kronos BTC Prediction API Client SDK
+
+用于连接 Kronos BTC 预测 API 的 Python 客户端。
+
+使用示例:
+    from client import KronosClient
+    
+    client = KronosClient("https://your-space.hf.space")
+    
+    # 健康检查
+    health = client.health()
+    
+    # 价格预测
+    prediction = client.predict(ohlcv_data, pred_len=24)
+    
+    # 交易信号
+    signal = client.get_signal(ohlcv_data)
+"""
+
+import time
+from datetime import datetime
+from typing import List, Dict, Any, Optional, Union
+from dataclasses import dataclass
+from enum import Enum
+
+import httpx
+import pandas as pd
+
+
+class SignalType(str, Enum):
+    """交易信号类型"""
+    STRONG_BUY = "STRONG_BUY"
+    BUY = "BUY"
+    HOLD = "HOLD"
+    SELL = "SELL"
+    STRONG_SELL = "STRONG_SELL"
+
+
+@dataclass
+class OHLCVData:
+    """OHLCV K线数据"""
+    timestamp: str
+    open: float
+    high: float
+    low: float
+    close: float
+    volume: float
+    amount: Optional[float] = None
+    
+    def to_dict(self) -> Dict[str, Any]:
+        return {
+            "timestamp": self.timestamp,
+            "open": self.open,
+            "high": self.high,
+            "low": self.low,
+            "close": self.close,
+            "volume": self.volume,
+            "amount": self.amount
+        }
+
+
+@dataclass
+class PredictResult:
+    """预测结果"""
+    current_price: float
+    mean_forecast: float
+    min_forecast: float
+    max_forecast: float
+    upside_probability: float
+    expected_return: float
+    volatility_amplification: float
+    confidence: float
+    forecast_prices: List[float]
+    timestamp: str
+    
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "PredictResult":
+        return cls(**data)
+    
+    def __repr__(self) -> str:
+        return (
+            f"PredictResult(\n"
+            f"  current_price={self.current_price:.2f},\n"
+            f"  mean_forecast={self.mean_forecast:.2f},\n"
+            f"  upside_probability={self.upside_probability:.1%},\n"
+            f"  expected_return={self.expected_return:.2%},\n"
+            f"  confidence={self.confidence:.1%}\n"
+            f")"
+        )
+
+
+@dataclass
+class SignalResult:
+    """交易信号结果"""
+    signal: SignalType
+    confidence: float
+    current_price: float
+    target_price: float
+    stop_loss_price: float
+    take_profit_price: float
+    upside_probability: float
+    expected_return: float
+    suggested_position_size: float
+    reason: str
+    timestamp: str
+    
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "SignalResult":
+        data["signal"] = SignalType(data["signal"])
+        return cls(**data)
+    
+    def __repr__(self) -> str:
+        return (
+            f"SignalResult(\n"
+            f"  signal={self.signal.value},\n"
+            f"  confidence={self.confidence:.1%},\n"
+            f"  current_price={self.current_price:.2f},\n"
+            f"  target_price={self.target_price:.2f},\n"
+            f"  stop_loss={self.stop_loss_price:.2f},\n"
+            f"  take_profit={self.take_profit_price:.2f},\n"
+            f"  position_size={self.suggested_position_size:.1%},\n"
+            f"  reason='{self.reason}'\n"
+            f")"
+        )
+
+
+@dataclass
+class HealthResult:
+    """健康检查结果"""
+    status: str
+    model_loaded: bool
+    model_version: str
+    device: str
+    timestamp: str
+    
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "HealthResult":
+        return cls(**data)
+
+
+class KronosClientError(Exception):
+    """Kronos 客户端错误"""
+    pass
+
+
+class KronosClient:
+    """
+    Kronos BTC 预测 API 客户端
+    
+    Args:
+        base_url: API 服务地址 (如 "https://your-space.hf.space")
+        timeout: 请求超时时间 (秒)
+        max_retries: 最大重试次数
+    
+    Examples:
+        >>> client = KronosClient("https://your-space.hf.space")
+        >>> health = client.health()
+        >>> print(f"Status: {health.status}")
+    """
+    
+    def __init__(
+        self,
+        base_url: str,
+        timeout: float = 60.0,
+        max_retries: int = 3
+    ):
+        self.base_url = base_url.rstrip("/")
+        self.timeout = timeout
+        self.max_retries = max_retries
+        self._client = httpx.Client(timeout=timeout)
+    
+    def __enter__(self):
+        return self
+    
+    def __exit__(self, *args):
+        self.close()
+    
+    def close(self):
+        """关闭客户端"""
+        self._client.close()
+    
+    def _request(
+        self,
+        method: str,
+        endpoint: str,
+        json: Optional[Dict] = None,
+        retry_count: int = 0
+    ) -> Dict[str, Any]:
+        """发送 HTTP 请求"""
+        url = f"{self.base_url}{endpoint}"
+        
+        try:
+            response = self._client.request(method, url, json=json)
+            
+            if response.status_code == 503:
+                # 模型未加载，等待重试
+                if retry_count < self.max_retries:
+                    time.sleep(5)
+                    return self._request(method, endpoint, json, retry_count + 1)
+                raise KronosClientError("Model not loaded after retries")
+            
+            response.raise_for_status()
+            return response.json()
+        
+        except httpx.ConnectError as e:
+            raise KronosClientError(f"Connection failed: {e}")
+        except httpx.TimeoutException as e:
+            raise KronosClientError(f"Request timeout: {e}")
+        except httpx.HTTPStatusError as e:
+            raise KronosClientError(f"HTTP error {e.response.status_code}: {e.response.text}")
+    
+    def health(self) -> HealthResult:
+        """
+        健康检查
+        
+        Returns:
+            HealthResult: 健康状态
+        """
+        data = self._request("GET", "/health")
+        return HealthResult.from_dict(data)
+    
+    def predict(
+        self,
+        data: Union[List[Dict], List[OHLCVData], pd.DataFrame],
+        pred_len: int = 24,
+        n_paths: int = 30,
+        temperature: float = 1.0,
+        top_p: float = 0.9
+    ) -> PredictResult:
+        """
+        预测 BTC 价格走势
+        
+        Args:
+            data: OHLCV 数据 (至少 100 条)
+                - List[Dict]: 字典列表
+                - List[OHLCVData]: OHLCVData 对象列表
+                - pd.DataFrame: DataFrame (需包含 timestamp, open, high, low, close, volume 列)
+            pred_len: 预测长度 (1-72 小时)
+            n_paths: Monte Carlo 路径数 (10-100)
+            temperature: 采样温度 (0.1-2.0)
+            top_p: Top-p 采样 (0.5-1.0)
+        
+        Returns:
+            PredictResult: 预测结果
+        
+        Examples:
+            >>> result = client.predict(df, pred_len=24)
+            >>> print(f"上涨概率: {result.upside_probability:.1%}")
+        """
+        ohlcv_list = self._convert_data(data)
+        
+        if len(ohlcv_list) < 100:
+            raise KronosClientError(f"At least 100 data points required, got {len(ohlcv_list)}")
+        
+        request_data = {
+            "data": ohlcv_list,
+            "pred_len": pred_len,
+            "n_paths": n_paths,
+            "temperature": temperature,
+            "top_p": top_p
+        }
+        
+        response = self._request("POST", "/predict", json=request_data)
+        return PredictResult.from_dict(response)
+    
+    def get_signal(
+        self,
+        data: Union[List[Dict], List[OHLCVData], pd.DataFrame],
+        buy_threshold: float = 0.58,
+        sell_threshold: float = 0.42,
+        stop_loss: float = 0.03,
+        take_profit: float = 0.08,
+        n_paths: int = 30
+    ) -> SignalResult:
+        """
+        获取交易信号
+        
+        Args:
+            data: OHLCV 数据 (至少 100 条)
+            buy_threshold: 买入阈值 (0.5-0.9)
+            sell_threshold: 卖出阈值 (0.1-0.5)
+            stop_loss: 止损比例 (0.01-0.1)
+            take_profit: 止盈比例 (0.02-0.2)
+            n_paths: Monte Carlo 路径数 (10-100)
+        
+        Returns:
+            SignalResult: 交易信号
+        
+        Examples:
+            >>> signal = client.get_signal(df)
+            >>> if signal.signal == SignalType.BUY:
+            ...     print(f"买入! 目标价: {signal.target_price:.2f}")
+        """
+        ohlcv_list = self._convert_data(data)
+        
+        if len(ohlcv_list) < 100:
+            raise KronosClientError(f"At least 100 data points required, got {len(ohlcv_list)}")
+        
+        request_data = {
+            "data": ohlcv_list,
+            "buy_threshold": buy_threshold,
+            "sell_threshold": sell_threshold,
+            "stop_loss": stop_loss,
+            "take_profit": take_profit,
+            "n_paths": n_paths
+        }
+        
+        response = self._request("POST", "/signal", json=request_data)
+        return SignalResult.from_dict(response)
+    
+    def _convert_data(
+        self,
+        data: Union[List[Dict], List[OHLCVData], pd.DataFrame]
+    ) -> List[Dict[str, Any]]:
+        """转换数据格式"""
+        if isinstance(data, pd.DataFrame):
+            return self._dataframe_to_list(data)
+        elif isinstance(data, list):
+            if len(data) == 0:
+                return []
+            if isinstance(data[0], OHLCVData):
+                return [d.to_dict() for d in data]
+            elif isinstance(data[0], dict):
+                return data
+        
+        raise KronosClientError(f"Unsupported data type: {type(data)}")
+    
+    def _dataframe_to_list(self, df: pd.DataFrame) -> List[Dict[str, Any]]:
+        """将 DataFrame 转换为列表"""
+        required_cols = ["open", "high", "low", "close", "volume"]
+        for col in required_cols:
+            if col not in df.columns:
+                raise KronosClientError(f"Missing required column: {col}")
+        
+        result = []
+        for _, row in df.iterrows():
+            # 处理时间戳
+            if "timestamp" in df.columns:
+                ts = row["timestamp"]
+                if isinstance(ts, pd.Timestamp):
+                    ts = ts.isoformat()
+                elif isinstance(ts, datetime):
+                    ts = ts.isoformat()
+                else:
+                    ts = str(ts)
+            else:
+                ts = datetime.utcnow().isoformat()
+            
+            result.append({
+                "timestamp": ts,
+                "open": float(row["open"]),
+                "high": float(row["high"]),
+                "low": float(row["low"]),
+                "close": float(row["close"]),
+                "volume": float(row["volume"]),
+                "amount": float(row["amount"]) if "amount" in df.columns else None
+            })
+        
+        return result
+
+
+class AsyncKronosClient:
+    """
+    异步 Kronos BTC 预测 API 客户端
+    
+    用于异步场景 (如 asyncio 应用)。
+    
+    Examples:
+        >>> async with AsyncKronosClient("https://your-space.hf.space") as client:
+        ...     health = await client.health()
+        ...     prediction = await client.predict(df)
+    """
+    
+    def __init__(
+        self,
+        base_url: str,
+        timeout: float = 60.0,
+        max_retries: int = 3
+    ):
+        self.base_url = base_url.rstrip("/")
+        self.timeout = timeout
+        self.max_retries = max_retries
+        self._client: Optional[httpx.AsyncClient] = None
+    
+    async def __aenter__(self):
+        self._client = httpx.AsyncClient(timeout=self.timeout)
+        return self
+    
+    async def __aexit__(self, *args):
+        await self.close()
+    
+    async def close(self):
+        """关闭客户端"""
+        if self._client:
+            await self._client.aclose()
+            self._client = None
+    
+    async def _request(
+        self,
+        method: str,
+        endpoint: str,
+        json: Optional[Dict] = None,
+        retry_count: int = 0
+    ) -> Dict[str, Any]:
+        """发送 HTTP 请求"""
+        if not self._client:
+            raise KronosClientError("Client not initialized. Use 'async with' context.")
+        
+        url = f"{self.base_url}{endpoint}"
+        
+        try:
+            response = await self._client.request(method, url, json=json)
+            
+            if response.status_code == 503:
+                if retry_count < self.max_retries:
+                    import asyncio
+                    await asyncio.sleep(5)
+                    return await self._request(method, endpoint, json, retry_count + 1)
+                raise KronosClientError("Model not loaded after retries")
+            
+            response.raise_for_status()
+            return response.json()
+        
+        except httpx.ConnectError as e:
+            raise KronosClientError(f"Connection failed: {e}")
+        except httpx.TimeoutException as e:
+            raise KronosClientError(f"Request timeout: {e}")
+        except httpx.HTTPStatusError as e:
+            raise KronosClientError(f"HTTP error {e.response.status_code}: {e.response.text}")
+    
+    async def health(self) -> HealthResult:
+        """健康检查"""
+        data = await self._request("GET", "/health")
+        return HealthResult.from_dict(data)
+    
+    async def predict(
+        self,
+        data: Union[List[Dict], List[OHLCVData], pd.DataFrame],
+        pred_len: int = 24,
+        n_paths: int = 30,
+        temperature: float = 1.0,
+        top_p: float = 0.9
+    ) -> PredictResult:
+        """预测 BTC 价格走势"""
+        ohlcv_list = self._convert_data(data)
+        
+        if len(ohlcv_list) < 100:
+            raise KronosClientError(f"At least 100 data points required")
+        
+        request_data = {
+            "data": ohlcv_list,
+            "pred_len": pred_len,
+            "n_paths": n_paths,
+            "temperature": temperature,
+            "top_p": top_p
+        }
+        
+        response = await self._request("POST", "/predict", json=request_data)
+        return PredictResult.from_dict(response)
+    
+    async def get_signal(
+        self,
+        data: Union[List[Dict], List[OHLCVData], pd.DataFrame],
+        buy_threshold: float = 0.58,
+        sell_threshold: float = 0.42,
+        stop_loss: float = 0.03,
+        take_profit: float = 0.08,
+        n_paths: int = 30
+    ) -> SignalResult:
+        """获取交易信号"""
+        ohlcv_list = self._convert_data(data)
+        
+        if len(ohlcv_list) < 100:
+            raise KronosClientError(f"At least 100 data points required")
+        
+        request_data = {
+            "data": ohlcv_list,
+            "buy_threshold": buy_threshold,
+            "sell_threshold": sell_threshold,
+            "stop_loss": stop_loss,
+            "take_profit": take_profit,
+            "n_paths": n_paths
+        }
+        
+        response = await self._request("POST", "/signal", json=request_data)
+        return SignalResult.from_dict(response)
+    
+    def _convert_data(
+        self,
+        data: Union[List[Dict], List[OHLCVData], pd.DataFrame]
+    ) -> List[Dict[str, Any]]:
+        """转换数据格式"""
+        if isinstance(data, pd.DataFrame):
+            return self._dataframe_to_list(data)
+        elif isinstance(data, list):
+            if len(data) == 0:
+                return []
+            if isinstance(data[0], OHLCVData):
+                return [d.to_dict() for d in data]
+            elif isinstance(data[0], dict):
+                return data
+        
+        raise KronosClientError(f"Unsupported data type: {type(data)}")
+    
+    def _dataframe_to_list(self, df: pd.DataFrame) -> List[Dict[str, Any]]:
+        """将 DataFrame 转换为列表"""
+        required_cols = ["open", "high", "low", "close", "volume"]
+        for col in required_cols:
+            if col not in df.columns:
+                raise KronosClientError(f"Missing required column: {col}")
+        
+        result = []
+        for _, row in df.iterrows():
+            if "timestamp" in df.columns:
+                ts = row["timestamp"]
+                if isinstance(ts, pd.Timestamp):
+                    ts = ts.isoformat()
+                elif isinstance(ts, datetime):
+                    ts = ts.isoformat()
+                else:
+                    ts = str(ts)
+            else:
+                ts = datetime.utcnow().isoformat()
+            
+            result.append({
+                "timestamp": ts,
+                "open": float(row["open"]),
+                "high": float(row["high"]),
+                "low": float(row["low"]),
+                "close": float(row["close"]),
+                "volume": float(row["volume"]),
+                "amount": float(row["amount"]) if "amount" in df.columns else None
+            })
+        
+        return result
+
+
+# ==================== 使用示例 ====================
+
+if __name__ == "__main__":
+    import asyncio
+    
+    # 示例数据 (实际使用时替换为真实数据)
+    def create_sample_data() -> pd.DataFrame:
+        """创建示例数据"""
+        import numpy as np
+        
+        np.random.seed(42)
+        n = 120
+        
+        dates = pd.date_range(start="2024-01-01", periods=n, freq="1h")
+        base_price = 42000
+        
+        prices = [base_price]
+        for i in range(1, n):
+            change = np.random.randn() * 100
+            prices.append(prices[-1] + change)
+        
+        df = pd.DataFrame({
+            "timestamp": dates,
+            "open": prices,
+            "high": [p + np.random.rand() * 50 for p in prices],
+            "low": [p - np.random.rand() * 50 for p in prices],
+            "close": [p + np.random.randn() * 20 for p in prices],
+            "volume": np.random.rand(n) * 1000 + 100
+        })
+        
+        return df
+    
+    # 同步示例
+    def sync_example():
+        print("=== 同步客户端示例 ===")
+        
+        # 创建客户端
+        client = KronosClient("http://localhost:7860")
+        
+        try:
+            # 健康检查
+            health = client.health()
+            print(f"Status: {health.status}")
+            print(f"Model loaded: {health.model_loaded}")
+            
+            # 创建示例数据
+            df = create_sample_data()
+            print(f"\n数据点数: {len(df)}")
+            
+            # 预测
+            prediction = client.predict(df, pred_len=24)
+            print(f"\n预测结果:")
+            print(prediction)
+            
+            # 交易信号
+            signal = client.get_signal(df)
+            print(f"\n交易信号:")
+            print(signal)
+            
+        except KronosClientError as e:
+            print(f"Error: {e}")
+        finally:
+            client.close()
+    
+    # 异步示例
+    async def async_example():
+        print("\n=== 异步客户端示例 ===")
+        
+        async with AsyncKronosClient("http://localhost:7860") as client:
+            try:
+                # 健康检查
+                health = await client.health()
+                print(f"Status: {health.status}")
+                
+                # 创建示例数据
+                df = create_sample_data()
+                
+                # 并发预测和信号
+                prediction, signal = await asyncio.gather(
+                    client.predict(df),
+                    client.get_signal(df)
+                )
+                
+                print(f"\n预测结果:")
+                print(prediction)
+                print(f"\n交易信号:")
+                print(signal)
+                
+            except KronosClientError as e:
+                print(f"Error: {e}")
+    
+    # 运行示例
+    print("Kronos Client SDK 示例\n")
+    print("注意: 请确保 API 服务已启动 (python app.py)\n")
+    
+    # 仅打印帮助信息，不实际运行
+    print("同步使用:")
+    print("  client = KronosClient('http://localhost:7860')")
+    print("  prediction = client.predict(df)")
+    print("  signal = client.get_signal(df)")
+    print()
+    print("异步使用:")
+    print("  async with AsyncKronosClient('http://localhost:7860') as client:")
+    print("      prediction = await client.predict(df)")
+    print("      signal = await client.get_signal(df)")

model/__init__.py

@@ -0,0 +1,3 @@
+from .kronos import KronosTokenizer, Kronos, calc_time_stamps
+
+__all__ = ['KronosTokenizer', 'Kronos', 'calc_time_stamps']

model/kronos.py

@@ -0,0 +1,662 @@
+import numpy as np
+import pandas as pd
+import torch
+from huggingface_hub import PyTorchModelHubMixin
+import sys
+
+from tqdm import trange
+
+sys.path.append("../")
+from model.module import *
+
+
+class KronosTokenizer(nn.Module, PyTorchModelHubMixin):
+    """
+    KronosTokenizer module for tokenizing input data using a hybrid quantization approach.
+
+    This tokenizer utilizes a combination of encoder and decoder Transformer blocks
+    along with the Binary Spherical Quantization (BSQuantizer) to compress and decompress input data.
+
+    Args:
+           d_in (int): Input dimension.
+           d_model (int): Model dimension.
+           n_heads (int): Number of attention heads.
+           ff_dim (int): Feed-forward dimension.
+           n_enc_layers (int): Number of encoder layers.
+           n_dec_layers (int): Number of decoder layers.
+           ffn_dropout_p (float): Dropout probability for feed-forward networks.
+           attn_dropout_p (float): Dropout probability for attention mechanisms.
+           resid_dropout_p (float): Dropout probability for residual connections.
+           s1_bits (int): Number of bits for the pre token in BSQuantizer.
+           s2_bits (int): Number of bits for the post token in BSQuantizer.
+           beta (float): Beta parameter for BSQuantizer.
+           gamma0 (float): Gamma0 parameter for BSQuantizer.
+           gamma (float): Gamma parameter for BSQuantizer.
+           zeta (float): Zeta parameter for BSQuantizer.
+           group_size (int): Group size parameter for BSQuantizer.
+
+    """
+
+    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):
+
+        super().__init__()
+        self.d_in = d_in
+        self.d_model = d_model
+        self.n_heads = n_heads
+        self.ff_dim = ff_dim
+        self.enc_layers = n_enc_layers
+        self.dec_layers = n_dec_layers
+        self.ffn_dropout_p = ffn_dropout_p
+        self.attn_dropout_p = attn_dropout_p
+        self.resid_dropout_p = resid_dropout_p
+
+        self.s1_bits = s1_bits
+        self.s2_bits = s2_bits
+        self.codebook_dim = s1_bits + s2_bits # Total dimension of the codebook after quantization
+        self.embed = nn.Linear(self.d_in, self.d_model)
+        self.head = nn.Linear(self.d_model, self.d_in)
+
+        # Encoder Transformer Blocks
+        self.encoder = nn.ModuleList([
+            TransformerBlock(self.d_model, self.n_heads, self.ff_dim, self.ffn_dropout_p, self.attn_dropout_p, self.resid_dropout_p)
+            for _ in range(self.enc_layers - 1)
+        ])
+        # Decoder Transformer Blocks
+        self.decoder = nn.ModuleList([
+            TransformerBlock(self.d_model, self.n_heads, self.ff_dim, self.ffn_dropout_p, self.attn_dropout_p, self.resid_dropout_p)
+            for _ in range(self.dec_layers - 1)
+        ])
+        self.quant_embed = nn.Linear(in_features=self.d_model, out_features=self.codebook_dim) # Linear layer before quantization
+        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)
+        self.post_quant_embed = nn.Linear(in_features=self.codebook_dim, out_features=self.d_model) # Linear layer after quantization (full codebook)
+        self.tokenizer = BSQuantizer(self.s1_bits, self.s2_bits, beta, gamma0, gamma, zeta, group_size) # BSQuantizer module
+
+    def forward(self, x):
+        """
+        Forward pass of the KronosTokenizer.
+
+        Args:
+            x (torch.Tensor): Input tensor of shape (batch_size, seq_len, d_in).
+
+        Returns:
+            tuple: A tuple containing:
+                - tuple: (z_pre, z) - Reconstructed outputs from decoder with s1_bits and full codebook respectively,
+                         both of shape (batch_size, seq_len, d_in).
+                - torch.Tensor: bsq_loss - Loss from the BSQuantizer.
+                - torch.Tensor: quantized - Quantized representation from BSQuantizer.
+                - torch.Tensor: z_indices - Indices from the BSQuantizer.
+        """
+        z = self.embed(x)
+
+        for layer in self.encoder:
+            z = layer(z)
+
+        z = self.quant_embed(z) # (B, T, codebook)
+
+        bsq_loss, quantized, z_indices = self.tokenizer(z)
+
+        quantized_pre = quantized[:, :, :self.s1_bits] # Extract the first part of quantized representation (s1_bits)
+        z_pre = self.post_quant_embed_pre(quantized_pre)
+
+        z = self.post_quant_embed(quantized)
+
+        # Decoder layers (for pre part - s1 bits)
+        for layer in self.decoder:
+            z_pre = layer(z_pre)
+        z_pre = self.head(z_pre)
+
+        # Decoder layers (for full codebook)
+        for layer in self.decoder:
+            z = layer(z)
+        z = self.head(z)
+
+        return (z_pre, z), bsq_loss, quantized, z_indices
+
+    def indices_to_bits(self, x, half=False):
+        """
+        Converts indices to bit representations and scales them.
+
+        Args:
+            x (torch.Tensor): Indices tensor.
+            half (bool, optional): Whether to process only half of the codebook dimension. Defaults to False.
+
+        Returns:
+            torch.Tensor: Bit representation tensor.
+        """
+        if half:
+            x1 = x[0] # Assuming x is a tuple of indices if half is True
+            x2 = x[1]
+            mask = 2 ** torch.arange(self.codebook_dim//2, device=x1.device, dtype=torch.long) # Create a mask for bit extraction
+            x1 = (x1.unsqueeze(-1) & mask) != 0 # Extract bits for the first half
+            x2 = (x2.unsqueeze(-1) & mask) != 0 # Extract bits for the second half
+            x = torch.cat([x1, x2], dim=-1) # Concatenate the bit representations
+        else:
+            mask = 2 ** torch.arange(self.codebook_dim, device=x.device, dtype=torch.long) # Create a mask for bit extraction
+            x = (x.unsqueeze(-1) & mask) != 0 # Extract bits
+
+        x = x.float() * 2 - 1 # Convert boolean to bipolar (-1, 1)
+        q_scale = 1. / (self.codebook_dim ** 0.5) # Scaling factor
+        x = x * q_scale
+        return x
+
+    def encode(self, x, half=False):
+        """
+        Encodes the input data into quantized indices.
+
+        Args:
+            x (torch.Tensor): Input tensor of shape (batch_size, seq_len, d_in).
+            half (bool, optional): Whether to use half quantization in BSQuantizer. Defaults to False.
+
+        Returns:
+            torch.Tensor: Quantized indices from BSQuantizer.
+        """
+        z = self.embed(x)
+        for layer in self.encoder:
+            z = layer(z)
+        z = self.quant_embed(z)
+
+        bsq_loss, quantized, z_indices = self.tokenizer(z, half=half, collect_metrics=False)
+        return z_indices
+
+    def decode(self, x, half=False):
+        """
+        Decodes quantized indices back to the input data space.
+
+        Args:
+            x (torch.Tensor): Quantized indices tensor.
+            half (bool, optional): Whether the indices were generated with half quantization. Defaults to False.
+
+        Returns:
+            torch.Tensor: Reconstructed output tensor of shape (batch_size, seq_len, d_in).
+        """
+        quantized = self.indices_to_bits(x, half)
+        z = self.post_quant_embed(quantized)
+        for layer in self.decoder:
+            z = layer(z)
+        z = self.head(z)
+        return z
+
+
+class Kronos(nn.Module, PyTorchModelHubMixin):
+    """
+    Kronos Model.
+
+    Args:
+        s1_bits (int): Number of bits for pre tokens.
+        s2_bits (int): Number of bits for post tokens.
+        n_layers (int): Number of Transformer blocks.
+        d_model (int): Dimension of the model's embeddings and hidden states.
+        n_heads (int): Number of attention heads in the MultiheadAttention layers.
+        ff_dim (int): Dimension of the feedforward network in the Transformer blocks.
+        ffn_dropout_p (float): Dropout probability for the feedforward network.
+        attn_dropout_p (float): Dropout probability for the attention layers.
+        resid_dropout_p (float): Dropout probability for residual connections.
+        token_dropout_p (float): Dropout probability for token embeddings.
+        learn_te (bool): Whether to use learnable temporal embeddings.
+    """
+
+    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):
+        super().__init__()
+        self.s1_bits = s1_bits
+        self.s2_bits = s2_bits
+        self.n_layers = n_layers
+        self.d_model = d_model
+        self.n_heads = n_heads
+        self.learn_te = learn_te
+        self.ff_dim = ff_dim
+        self.ffn_dropout_p = ffn_dropout_p
+        self.attn_dropout_p = attn_dropout_p
+        self.resid_dropout_p = resid_dropout_p
+        self.token_dropout_p = token_dropout_p
+
+        self.s1_vocab_size = 2 ** self.s1_bits
+        self.token_drop = nn.Dropout(self.token_dropout_p)
+        self.embedding = HierarchicalEmbedding(self.s1_bits, self.s2_bits, self.d_model)
+        self.time_emb = TemporalEmbedding(self.d_model, self.learn_te)
+        self.transformer = nn.ModuleList([
+            TransformerBlock(self.d_model, self.n_heads, self.ff_dim, self.ffn_dropout_p, self.attn_dropout_p, self.resid_dropout_p)
+            for _ in range(self.n_layers)
+        ])
+        self.norm = RMSNorm(self.d_model)
+        self.dep_layer = DependencyAwareLayer(self.d_model)
+        self.head = DualHead(self.s1_bits, self.s2_bits, self.d_model)
+        self.apply(self._init_weights)
+
+    def _init_weights(self, module):
+
+        if isinstance(module, nn.Linear):
+            nn.init.xavier_normal_(module.weight)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            nn.init.normal_(module.weight, mean=0, std=self.embedding.d_model ** -0.5)
+        elif isinstance(module, nn.LayerNorm):
+            nn.init.ones_(module.weight)
+            nn.init.zeros_(module.bias)
+        elif isinstance(module, RMSNorm):
+            nn.init.ones_(module.weight)
+
+    def forward(self, s1_ids, s2_ids, stamp=None, padding_mask=None, use_teacher_forcing=False, s1_targets=None):
+        """
+        Args:
+            s1_ids (torch.Tensor): Input tensor of s1 token IDs. Shape: [batch_size, seq_len]
+            s2_ids (torch.Tensor): Input tensor of s2 token IDs. Shape: [batch_size, seq_len]
+            stamp (torch.Tensor, optional): Temporal stamp tensor. Shape: [batch_size, seq_len]. Defaults to None.
+            padding_mask (torch.Tensor, optional): Mask for padding tokens. Shape: [batch_size, seq_len]. Defaults to None.
+            use_teacher_forcing (bool, optional): Whether to use teacher forcing for s1 decoding. Defaults to False.
+            s1_targets (torch.Tensor, optional): Target s1 token IDs for teacher forcing. Shape: [batch_size, seq_len]. Defaults to None.
+
+        Returns:
+            Tuple[torch.Tensor, torch.Tensor]:
+                - s1 logits: Logits for s1 token predictions. Shape: [batch_size, seq_len, s1_vocab_size]
+                - s2_logits: Logits for s2 token predictions, conditioned on s1. Shape: [batch_size, seq_len, s2_vocab_size]
+        """
+        x = self.embedding([s1_ids, s2_ids])
+        if stamp is not None:
+            time_embedding = self.time_emb(stamp)
+            x = x + time_embedding
+        x = self.token_drop(x)
+
+        for layer in self.transformer:
+            x = layer(x, key_padding_mask=padding_mask)
+
+        x = self.norm(x)
+
+        s1_logits = self.head(x)
+
+        if use_teacher_forcing:
+            sibling_embed = self.embedding.emb_s1(s1_targets)
+        else:
+            s1_probs = F.softmax(s1_logits.detach(), dim=-1)
+            sample_s1_ids = torch.multinomial(s1_probs.view(-1, self.s1_vocab_size), 1).view(s1_ids.shape)
+            sibling_embed = self.embedding.emb_s1(sample_s1_ids)
+
+        x2 = self.dep_layer(x, sibling_embed, key_padding_mask=padding_mask) # Dependency Aware Layer: Condition on s1 embeddings
+        s2_logits = self.head.cond_forward(x2)
+        return s1_logits, s2_logits
+
+    def decode_s1(self, s1_ids, s2_ids, stamp=None, padding_mask=None):
+        """
+        Decodes only the s1 tokens.
+
+        This method performs a forward pass to predict only s1 tokens. It returns the s1 logits
+        and the context representation from the Transformer, which can be used for subsequent s2 decoding.
+
+        Args:
+            s1_ids (torch.Tensor): Input tensor of s1 token IDs. Shape: [batch_size, seq_len]
+            s2_ids (torch.Tensor): Input tensor of s2 token IDs. Shape: [batch_size, seq_len]
+            stamp (torch.Tensor, optional): Temporal stamp tensor. Shape: [batch_size, seq_len]. Defaults to None.
+            padding_mask (torch.Tensor, optional): Mask for padding tokens. Shape: [batch_size, seq_len]. Defaults to None.
+
+        Returns:
+            Tuple[torch.Tensor, torch.Tensor]:
+                - s1 logits: Logits for s1 token predictions. Shape: [batch_size, seq_len, s1_vocab_size]
+                - context: Context representation from the Transformer. Shape: [batch_size, seq_len, d_model]
+        """
+        x = self.embedding([s1_ids, s2_ids])
+        if stamp is not None:
+            time_embedding = self.time_emb(stamp)
+            x = x + time_embedding
+        x = self.token_drop(x)
+
+        for layer in self.transformer:
+            x = layer(x, key_padding_mask=padding_mask)
+
+        x = self.norm(x)
+
+        s1_logits = self.head(x)
+        return s1_logits, x
+
+    def decode_s2(self, context, s1_ids, padding_mask=None):
+        """
+        Decodes the s2 tokens, conditioned on the context and s1 tokens.
+
+        This method decodes s2 tokens based on a pre-computed context representation (typically from `decode_s1`)
+        and the s1 token IDs. It uses the dependency-aware layer and the conditional s2 head to predict s2 tokens.
+
+        Args:
+            context (torch.Tensor): Context representation from the transformer (output of decode_s1).
+                                     Shape: [batch_size, seq_len, d_model]
+            s1_ids (torch.torch.Tensor): Input tensor of s1 token IDs. Shape: [batch_size, seq_len]
+            padding_mask (torch.Tensor, optional): Mask for padding tokens. Shape: [batch_size, seq_len]. Defaults to None.
+
+        Returns:
+            torch.Tensor: s2 logits. Shape: [batch_size, seq_len, s2_vocab_size]
+        """
+        sibling_embed = self.embedding.emb_s1(s1_ids)
+        x2 = self.dep_layer(context, sibling_embed, key_padding_mask=padding_mask)
+        return self.head.cond_forward(x2)
+
+
+def top_k_top_p_filtering(
+        logits,
+        top_k: int = 0,
+        top_p: float = 1.0,
+        filter_value: float = -float("Inf"),
+        min_tokens_to_keep: int = 1,
+):
+    """Filter a distribution of logits using top-k and/or nucleus (top-p) filtering
+    Args:
+        logits: logits distribution shape (batch size, vocabulary size)
+        if top_k > 0: keep only top k tokens with highest probability (top-k filtering).
+        if top_p < 1.0: keep the top tokens with cumulative probability >= top_p (nucleus filtering).
+            Nucleus filtering is described in Holtzman et al. (http://arxiv.org/abs/1904.09751)
+        Make sure we keep at least min_tokens_to_keep per batch example in the output
+    From: https://gist.github.com/thomwolf/1a5a29f6962089e871b94cbd09daf317
+    """
+    if top_k > 0:
+        top_k = min(max(top_k, min_tokens_to_keep), logits.size(-1))  # Safety check
+        # Remove all tokens with a probability less than the last token of the top-k
+        indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
+        logits[indices_to_remove] = filter_value
+        return logits
+
+    if top_p < 1.0:
+        sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+        cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
+
+        # Remove tokens with cumulative probability above the threshold (token with 0 are kept)
+        sorted_indices_to_remove = cumulative_probs > top_p
+        if min_tokens_to_keep > 1:
+            # Keep at least min_tokens_to_keep (set to min_tokens_to_keep-1 because we add the first one below)
+            sorted_indices_to_remove[..., :min_tokens_to_keep] = 0
+        # Shift the indices to the right to keep also the first token above the threshold
+        sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
+        sorted_indices_to_remove[..., 0] = 0
+
+        # scatter sorted tensors to original indexing
+        indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
+        logits[indices_to_remove] = filter_value
+        return logits
+
+
+def sample_from_logits(logits, temperature=1.0, top_k=None, top_p=None, sample_logits=True):
+    logits = logits / temperature
+    if top_k is not None or top_p is not None:
+        if top_k > 0 or top_p < 1.0:
+            logits = top_k_top_p_filtering(logits, top_k=top_k, top_p=top_p)
+
+    probs = F.softmax(logits, dim=-1)
+
+    if not sample_logits:
+        _, x = top_k(probs, k=1, dim=-1)
+    else:
+        x = torch.multinomial(probs, num_samples=1)
+
+    return x
+
+
+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):
+    with torch.no_grad():
+        x = torch.clip(x, -clip, clip)
+
+        device = x.device
+        x = x.unsqueeze(1).repeat(1, sample_count, 1, 1).reshape(-1, x.size(1), x.size(2)).to(device)
+        x_stamp = x_stamp.unsqueeze(1).repeat(1, sample_count, 1, 1).reshape(-1, x_stamp.size(1), x_stamp.size(2)).to(device)
+        y_stamp = y_stamp.unsqueeze(1).repeat(1, sample_count, 1, 1).reshape(-1, y_stamp.size(1), y_stamp.size(2)).to(device)
+
+        x_token = tokenizer.encode(x, half=True)
+        
+        initial_seq_len = x.size(1)
+        batch_size = x_token[0].size(0)
+        total_seq_len = initial_seq_len + pred_len
+        full_stamp = torch.cat([x_stamp, y_stamp], dim=1)
+
+        generated_pre = x_token[0].new_empty(batch_size, pred_len)
+        generated_post = x_token[1].new_empty(batch_size, pred_len)
+
+        pre_buffer = x_token[0].new_zeros(batch_size, max_context)
+        post_buffer = x_token[1].new_zeros(batch_size, max_context)
+        buffer_len = min(initial_seq_len, max_context)
+        if buffer_len > 0:
+            start_idx = max(0, initial_seq_len - max_context)
+            pre_buffer[:, :buffer_len] = x_token[0][:, start_idx:start_idx + buffer_len]
+            post_buffer[:, :buffer_len] = x_token[1][:, start_idx:start_idx + buffer_len]
+
+        if verbose:
+            ran = trange
+        else:
+            ran = range
+        for i in ran(pred_len):
+            current_seq_len = initial_seq_len + i
+            window_len = min(current_seq_len, max_context)
+
+            if current_seq_len <= max_context:
+                input_tokens = [
+                    pre_buffer[:, :window_len],
+                    post_buffer[:, :window_len]
+                ]
+            else:
+                input_tokens = [pre_buffer, post_buffer]
+
+            context_end = current_seq_len
+            context_start = max(0, context_end - max_context)
+            current_stamp = full_stamp[:, context_start:context_end, :].contiguous()
+
+            s1_logits, context = model.decode_s1(input_tokens[0], input_tokens[1], current_stamp)
+            s1_logits = s1_logits[:, -1, :]
+            sample_pre = sample_from_logits(s1_logits, temperature=T, top_k=top_k, top_p=top_p, sample_logits=True)
+
+            s2_logits = model.decode_s2(context, sample_pre)
+            s2_logits = s2_logits[:, -1, :]
+            sample_post = sample_from_logits(s2_logits, temperature=T, top_k=top_k, top_p=top_p, sample_logits=True)
+
+            generated_pre[:, i] = sample_pre.squeeze(-1)
+            generated_post[:, i] = sample_post.squeeze(-1)
+
+            if current_seq_len < max_context:
+                pre_buffer[:, current_seq_len] = sample_pre.squeeze(-1)
+                post_buffer[:, current_seq_len] = sample_post.squeeze(-1)
+            else:
+                pre_buffer.copy_(torch.roll(pre_buffer, shifts=-1, dims=1))
+                post_buffer.copy_(torch.roll(post_buffer, shifts=-1, dims=1))
+                pre_buffer[:, -1] = sample_pre.squeeze(-1)
+                post_buffer[:, -1] = sample_post.squeeze(-1)
+
+        full_pre = torch.cat([x_token[0], generated_pre], dim=1)
+        full_post = torch.cat([x_token[1], generated_post], dim=1)
+
+        context_start = max(0, total_seq_len - max_context)
+        input_tokens = [
+            full_pre[:, context_start:total_seq_len].contiguous(),
+            full_post[:, context_start:total_seq_len].contiguous()
+        ]
+        z = tokenizer.decode(input_tokens, half=True)
+        z = z.reshape(-1, sample_count, z.size(1), z.size(2))
+        preds = z.cpu().numpy()
+        preds = np.mean(preds, axis=1)
+
+        return preds
+
+
+def calc_time_stamps(x_timestamp):
+    time_df = pd.DataFrame()
+    time_df['minute'] = x_timestamp.dt.minute
+    time_df['hour'] = x_timestamp.dt.hour
+    time_df['weekday'] = x_timestamp.dt.weekday
+    time_df['day'] = x_timestamp.dt.day
+    time_df['month'] = x_timestamp.dt.month
+    return time_df
+
+
+class KronosPredictor:
+
+    def __init__(self, model, tokenizer, device=None, max_context=512, clip=5):
+        self.tokenizer = tokenizer
+        self.model = model
+        self.max_context = max_context
+        self.clip = clip
+        self.price_cols = ['open', 'high', 'low', 'close']
+        self.vol_col = 'volume'
+        self.amt_vol = 'amount'
+        self.time_cols = ['minute', 'hour', 'weekday', 'day', 'month']
+        
+        # Auto-detect device if not specified
+        if device is None:
+            if torch.cuda.is_available():
+                device = "cuda:0"
+            elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
+                device = "mps"
+            else:
+                device = "cpu"
+        
+        self.device = device
+
+        self.tokenizer = self.tokenizer.to(self.device)
+        self.model = self.model.to(self.device)
+
+    def generate(self, x, x_stamp, y_stamp, pred_len, T, top_k, top_p, sample_count, verbose):
+
+        x_tensor = torch.from_numpy(np.array(x).astype(np.float32)).to(self.device)
+        x_stamp_tensor = torch.from_numpy(np.array(x_stamp).astype(np.float32)).to(self.device)
+        y_stamp_tensor = torch.from_numpy(np.array(y_stamp).astype(np.float32)).to(self.device)
+
+        preds = auto_regressive_inference(self.tokenizer, self.model, x_tensor, x_stamp_tensor, y_stamp_tensor, self.max_context, pred_len,
+                                          self.clip, T, top_k, top_p, sample_count, verbose)
+        preds = preds[:, -pred_len:, :]
+        return preds
+
+    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):
+
+        if not isinstance(df, pd.DataFrame):
+            raise ValueError("Input must be a pandas DataFrame.")
+
+        if not all(col in df.columns for col in self.price_cols):
+            raise ValueError(f"Price columns {self.price_cols} not found in DataFrame.")
+
+        df = df.copy()
+        if self.vol_col not in df.columns:
+            df[self.vol_col] = 0.0  # Fill missing volume with zeros
+            df[self.amt_vol] = 0.0  # Fill missing amount with zeros
+        if self.amt_vol not in df.columns and self.vol_col in df.columns:
+            df[self.amt_vol] = df[self.vol_col] * df[self.price_cols].mean(axis=1)
+
+        if df[self.price_cols + [self.vol_col, self.amt_vol]].isnull().values.any():
+            raise ValueError("Input DataFrame contains NaN values in price or volume columns.")
+
+        x_time_df = calc_time_stamps(x_timestamp)
+        y_time_df = calc_time_stamps(y_timestamp)
+
+        x = df[self.price_cols + [self.vol_col, self.amt_vol]].values.astype(np.float32)
+        x_stamp = x_time_df.values.astype(np.float32)
+        y_stamp = y_time_df.values.astype(np.float32)
+
+        x_mean, x_std = np.mean(x, axis=0), np.std(x, axis=0)
+
+        x = (x - x_mean) / (x_std + 1e-5)
+        x = np.clip(x, -self.clip, self.clip)
+
+        x = x[np.newaxis, :]
+        x_stamp = x_stamp[np.newaxis, :]
+        y_stamp = y_stamp[np.newaxis, :]
+
+        preds = self.generate(x, x_stamp, y_stamp, pred_len, T, top_k, top_p, sample_count, verbose)
+
+        preds = preds.squeeze(0)
+        preds = preds * (x_std + 1e-5) + x_mean
+
+        pred_df = pd.DataFrame(preds, columns=self.price_cols + [self.vol_col, self.amt_vol], index=y_timestamp)
+        return pred_df
+
+
+    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):
+        """
+        Perform parallel (batch) prediction on multiple time series. All series must have the same historical length and prediction length (pred_len).
+
+        Args:
+            df_list (List[pd.DataFrame]): List of input DataFrames, each containing price columns and optional volume/amount columns.
+            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.
+            y_timestamp_list (List[pd.DatetimeIndex or Series]): List of future prediction timestamps, length should equal pred_len.
+            pred_len (int): Number of prediction steps.
+            T (float): Sampling temperature.
+            top_k (int): Top-k filtering threshold.
+            top_p (float): Top-p (nucleus sampling) threshold.
+            sample_count (int): Number of parallel samples per series, automatically averaged internally.
+            verbose (bool): Whether to display autoregressive progress.
+
+        Returns:
+            List[pd.DataFrame]: List of prediction results in the same order as input, each DataFrame contains
+                                `open, high, low, close, volume, amount` columns, indexed by corresponding `y_timestamp`.
+        """
+        # Basic validation
+        if not isinstance(df_list, (list, tuple)) or not isinstance(x_timestamp_list, (list, tuple)) or not isinstance(y_timestamp_list, (list, tuple)):
+            raise ValueError("df_list, x_timestamp_list, y_timestamp_list must be list or tuple types.")
+        if not (len(df_list) == len(x_timestamp_list) == len(y_timestamp_list)):
+            raise ValueError("df_list, x_timestamp_list, y_timestamp_list must have consistent lengths.")
+
+        num_series = len(df_list)
+
+        x_list = []
+        x_stamp_list = []
+        y_stamp_list = []
+        means = []
+        stds = []
+        seq_lens = []
+        y_lens = []
+
+        for i in range(num_series):
+            df = df_list[i]
+            if not isinstance(df, pd.DataFrame):
+                raise ValueError(f"Input at index {i} is not a pandas DataFrame.")
+            if not all(col in df.columns for col in self.price_cols):
+                raise ValueError(f"DataFrame at index {i} is missing price columns {self.price_cols}.")
+
+            df = df.copy()
+            if self.vol_col not in df.columns:
+                df[self.vol_col] = 0.0
+                df[self.amt_vol] = 0.0
+            if self.amt_vol not in df.columns and self.vol_col in df.columns:
+                df[self.amt_vol] = df[self.vol_col] * df[self.price_cols].mean(axis=1)
+
+            if df[self.price_cols + [self.vol_col, self.amt_vol]].isnull().values.any():
+                raise ValueError(f"DataFrame at index {i} contains NaN values in price or volume columns.")
+
+            x_timestamp = x_timestamp_list[i]
+            y_timestamp = y_timestamp_list[i]
+
+            x_time_df = calc_time_stamps(x_timestamp)
+            y_time_df = calc_time_stamps(y_timestamp)
+
+            x = df[self.price_cols + [self.vol_col, self.amt_vol]].values.astype(np.float32)
+            x_stamp = x_time_df.values.astype(np.float32)
+            y_stamp = y_time_df.values.astype(np.float32)
+
+            if x.shape[0] != x_stamp.shape[0]:
+                raise ValueError(f"Inconsistent lengths at index {i}: x has {x.shape[0]} vs x_stamp has {x_stamp.shape[0]}.")
+            if y_stamp.shape[0] != pred_len:
+                raise ValueError(f"y_timestamp length at index {i} should equal pred_len={pred_len}, got {y_stamp.shape[0]}.")
+
+            x_mean, x_std = np.mean(x, axis=0), np.std(x, axis=0)
+            x_norm = (x - x_mean) / (x_std + 1e-5)
+            x_norm = np.clip(x_norm, -self.clip, self.clip)
+
+            x_list.append(x_norm)
+            x_stamp_list.append(x_stamp)
+            y_stamp_list.append(y_stamp)
+            means.append(x_mean)
+            stds.append(x_std)
+
+            seq_lens.append(x_norm.shape[0])
+            y_lens.append(y_stamp.shape[0])
+
+        # Require all series to have consistent historical and prediction lengths for batch processing
+        if len(set(seq_lens)) != 1:
+            raise ValueError(f"Parallel prediction requires all series to have consistent historical lengths, got: {seq_lens}")
+        if len(set(y_lens)) != 1:
+            raise ValueError(f"Parallel prediction requires all series to have consistent prediction lengths, got: {y_lens}")
+
+        x_batch = np.stack(x_list, axis=0).astype(np.float32)           # (B, seq_len, feat)
+        x_stamp_batch = np.stack(x_stamp_list, axis=0).astype(np.float32) # (B, seq_len, time_feat)
+        y_stamp_batch = np.stack(y_stamp_list, axis=0).astype(np.float32) # (B, pred_len, time_feat)
+
+        preds = self.generate(x_batch, x_stamp_batch, y_stamp_batch, pred_len, T, top_k, top_p, sample_count, verbose)
+        # preds: (B, pred_len, feat)
+
+        pred_dfs = []
+        for i in range(num_series):
+            preds_i = preds[i] * (stds[i] + 1e-5) + means[i]
+            pred_df = pd.DataFrame(preds_i, columns=self.price_cols + [self.vol_col, self.amt_vol], index=y_timestamp_list[i])
+            pred_dfs.append(pred_df)
+
+        return pred_dfs
+

model/module.py

@@ -0,0 +1,570 @@
+import math
+
+from einops import rearrange, reduce
+import torch
+import torch.nn as nn
+from torch.autograd import Function
+import torch.nn.functional as F
+
+
+class DifferentiableEntropyFunction(Function):
+    @staticmethod
+    def forward(ctx, zq, basis, K, eps):
+        zb = (zq + 1) / 2
+        zi = ((zb * basis).sum(-1)).to(torch.int64)
+        cnt = torch.scatter_reduce(torch.zeros(2 ** K, device=zq.device, dtype=zq.dtype),
+                                   0,
+                                   zi.flatten(),
+                                   torch.ones_like(zi.flatten()).to(zq.dtype),
+                                   'sum')
+        prob = (cnt + eps) / (cnt + eps).sum()
+        H = -(prob * torch.log(prob)).sum()
+        ctx.save_for_backward(zq, zi, prob)
+        ctx.K = K
+        return H
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        zq, zi, prob = ctx.saved_tensors
+        grad_array = -grad_output * (torch.log(prob) + 1) / zi.numel() / ctx.K
+        reord_grad = grad_array[zi.flatten()].reshape(zi.shape)
+        grad_input = reord_grad.unsqueeze(-1) * zq
+        return grad_input, None, None, None, None
+
+
+def codebook_entropy(zq, basis, K, eps=1e-4):
+    return DifferentiableEntropyFunction.apply(zq, basis, K, eps)
+
+
+class BinarySphericalQuantizer(nn.Module):
+    def __init__(self, embed_dim, beta, gamma0, gamma, zeta,
+                 input_format='bchw',
+                 soft_entropy=True, group_size=9,
+                 persample_entropy_compute='analytical',
+                 cb_entropy_compute='group',
+                 l2_norm=True,
+                 inv_temperature=1):
+        """
+        Paper link: https://arxiv.org/pdf/2406.07548.pdf
+        Here we use the official implementation of the BinarySphericalQuantizer.
+        """
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.beta = beta  # loss weight for commit loss
+        self.gamma0 = gamma0  # loss weight for entropy penalty
+        self.gamma = gamma  # loss weight for entropy penalty
+        self.zeta = zeta  # loss weight for entire entropy penalty
+        self.input_format = input_format
+        assert self.embed_dim % group_size == 0, "embed_dim must be divisible by group_size"
+        self.num_groups = self.embed_dim // group_size
+        self.group_size = group_size
+        assert persample_entropy_compute in ['group', 'analytical'], "persample_entropy_compute must be either 'group' or 'analytical'"
+        assert cb_entropy_compute in ['group', 'nce'], "cb_entropy_compute must be either 'group' or 'nce'"
+        self.persample_entropy_compute = persample_entropy_compute
+        self.cb_entropy_compute = cb_entropy_compute
+        self.l2_norm = l2_norm
+        self.inv_temperature = inv_temperature
+
+        self.register_buffer('basis', 2 ** torch.arange(embed_dim - 1, -1, -1))
+        self.register_buffer('group_basis', 2 ** torch.arange(group_size - 1, -1, -1))
+
+        self.num_dimensions = 2 ** embed_dim
+        self.bits_per_index = embed_dim
+
+        # we only need to keep the codebook portion up to the group size
+        # because we approximate the H loss with this subcode
+        group_codes = torch.arange(2 ** self.group_size)
+        group_codebook = self.indexes_to_codes(group_codes).float()[:, -group_size:]
+        self.register_buffer('group_codebook', group_codebook, persistent=False)
+
+        self.soft_entropy = soft_entropy  # soft_entropy: Sec 3.2 of https://arxiv.org/pdf/1911.05894.pdf
+
+    def quantize(self, z):
+        assert z.shape[-1] == self.embed_dim, f"Expected {self.embed_dim} dimensions, got {z.shape[-1]}"
+
+        zhat = torch.where(z > 0,
+                           torch.tensor(1, dtype=z.dtype, device=z.device),
+                           torch.tensor(-1, dtype=z.dtype, device=z.device))
+        return z + (zhat - z).detach()
+
+    def forward(self, z, collect_metrics=True):
+        # if self.input_format == 'bchw':
+        #     z = rearrange(z, 'b c h w -> b h w c')
+        zq = self.quantize(z)
+
+        q_scale = 1. / (self.embed_dim ** 0.5) if self.l2_norm else 1.
+
+        zq = zq * q_scale
+
+        if not collect_metrics:
+            return zq, zq.new_zeros(()), {}
+
+        indices = self.codes_to_indexes(zq.detach())
+        group_indices = self.codes_to_group_indexes(zq.detach())
+        if not self.training:
+            used_codes = torch.unique(indices, return_counts=False)
+        else:
+            used_codes = None
+
+        if self.soft_entropy:
+            persample_entropy, cb_entropy, avg_prob = self.soft_entropy_loss(z)
+            entropy_penalty = self.gamma0 * persample_entropy - self.gamma * cb_entropy
+        else:
+            zb_by_sample = ((zq + 1) / 2).reshape(z.shape[0], -1, z.shape[-1]).to(torch.float32)
+            persample_entropy = self.get_hard_per_sample_entropy(zb_by_sample)
+            cb_entropy = codebook_entropy(zq, self.basis, self.embed_dim)
+            entropy_penalty = self.gamma0 * persample_entropy - self.gamma * cb_entropy
+
+        # commit loss
+        commit_loss = self.beta * torch.mean(((zq.detach() - z) ** 2).sum(dim=-1))
+
+        # if self.input_format == 'bchw':
+        #     zq = rearrange(zq, 'b h w c -> b c h w')
+
+        return (
+            zq,
+            commit_loss + self.zeta * entropy_penalty / self.inv_temperature,
+            {"H": cb_entropy, "used_codes": used_codes, "indices": indices, "group_indices": group_indices,
+             "avg_prob": avg_prob}
+        )
+
+    def soft_entropy_loss(self, z):
+        # if we divide the code in subgroups of size group_size, the codebook will be of size 2 ** group_size
+        # the sub-code is the last group_size bits of the full code
+        group_code_book = self.group_codebook / (self.embed_dim ** 0.5 if self.l2_norm else 1)
+        divided_z = rearrange(z, '... (g c) -> ... g c', c=self.group_size)
+
+        # we calculate the distance between the divided_z and the codebook for each subgroup
+        distance = - 2 * torch.einsum('... g c, d c ->... g d', divided_z, group_code_book)
+        prob = (-distance * self.inv_temperature).softmax(dim=-1)
+        if self.persample_entropy_compute == 'analytical':
+            if self.l2_norm:
+                p = torch.sigmoid(-4 * z / (self.embed_dim ** 0.5) * self.inv_temperature)
+            else:
+                p = torch.sigmoid(-4 * z * self.inv_temperature)
+            prob = torch.stack([p, 1 - p], dim=-1)
+            per_sample_entropy = self.get_entropy(prob, dim=-1, normalize=False).sum(dim=-1).mean()
+        else:
+            per_sample_entropy = self.get_entropy(prob, dim=-1, normalize=False).sum(dim=-1).mean()
+
+        # macro average of the probability of each subgroup
+        avg_prob = reduce(prob, '... g d ->g d', 'mean')
+        codebook_entropy = self.get_entropy(avg_prob, dim=-1, normalize=False)
+
+        # the approximation of the entropy is the sum of the entropy of each subgroup
+        return per_sample_entropy, codebook_entropy.sum(), avg_prob
+
+    def get_hard_per_sample_entropy(self, zb_by_sample):
+        probs_per_dim = zb_by_sample.sum(1) / zb_by_sample.shape[1]
+        persample_entropy = - probs_per_dim * torch.log(probs_per_dim + 1e-8) - (1 - probs_per_dim) * torch.log(1 - probs_per_dim + 1e-8)
+        persample_entropy = persample_entropy.sum(-1)
+        return persample_entropy.mean()
+
+    def codes_to_indexes(self, zhat):
+        """Converts a `code` to an index in the codebook.
+        Args:
+            zhat: A tensor of shape (B, ..., C) containing the codes. must be in {-1, 1}
+        """
+        assert zhat.shape[-1] == self.embed_dim, f"Expected {self.embed_dim} dimensions, got {zhat.shape[-1]}"
+        return ((zhat + 1) / 2 * self.basis).sum(axis=-1).to(torch.int64)
+
+    def codes_to_group_indexes(self, zhat):
+        """Converts a `code` to a list of indexes (in groups) in the codebook.
+        Args:
+            zhat: A tensor of shape (B, ..., C) containing the codes. must be in {-1, 1}
+        """
+        zhat_in_group = rearrange(zhat, 'b ... (g c) -> b ... g c', c=self.group_size)
+        return ((zhat_in_group + 1) / 2 * self.group_basis).sum(axis=-1).to(torch.int64)
+
+    def indexes_to_codes(self, indices):
+        """Inverse of `indexes_to_codes`."""
+        indices = indices.unsqueeze(-1)
+        codes_non_centered = torch.remainder(
+            torch.floor_divide(indices, self.basis), 2
+        )
+        return codes_non_centered * 2 - 1
+
+    def group_indexes_to_codes(self, group_indices):
+        """Inverse of `group_indexes_to_codes`."""
+        group_indices = group_indices.unsqueeze(-1)
+        codes_non_centered = torch.remainder(
+            torch.floor_divide(group_indices, self.group_basis), 2
+        )
+        codes_non_centered = rearrange(codes_non_centered, 'b ... g c -> b ... (g c)')
+        return codes_non_centered * 2 - 1
+
+    def get_entropy(self, count, dim=-1, eps=1e-4, normalize=True):
+        if normalize:
+            probs = (count + eps) / (count + eps).sum(dim=dim, keepdim=True)
+        else:
+            probs = count
+        H = -(probs * torch.log(probs + 1e-8)).sum(dim=dim)
+        return H
+
+    def get_group_codebook_entry(self, group_indices):
+        z_q = self.group_indexes_to_codes(group_indices)
+        q_scale = 1. / (self.embed_dim ** 0.5) if self.l2_norm else 1.
+        z_q = z_q * q_scale
+        if self.input_format == 'bchw':
+            h, w = int(z_q.shape[1] ** 0.5)
+            assert h * w == z_q.shape[1], 'Invalid sequence length'
+            z_q = rearrange(z_q, 'b (h w) c -> b c h w', h=h)
+        return z_q
+
+    def get_codebook_entry(self, indices):
+        z_q = self.indexes_to_codes(indices)
+        q_scale = 1. / (self.embed_dim ** 0.5) if self.l2_norm else 1.
+        z_q = z_q * q_scale
+        if self.input_format == 'bchw':
+            h, w = int(z_q.shape[1] ** 0.5)
+            assert h * w == z_q.shape[1], 'Invalid sequence length'
+            z_q = rearrange(z_q, 'b (h w) c -> b c h w', h=h)
+        return z_q
+
+
+class BSQuantizer(nn.Module):
+
+    def __init__(self, s1_bits, s2_bits, beta, gamma0, gamma, zeta, group_size):
+        super().__init__()
+        self.codebook_dim = s1_bits + s2_bits
+        self.s1_bits = s1_bits
+        self.s2_bits = s2_bits
+        self.bsq = BinarySphericalQuantizer(self.codebook_dim, beta, gamma0, gamma, zeta, group_size=group_size)
+
+    def bits_to_indices(self, bits):
+        bits = (bits >= 0).to(torch.long)
+        indices = 2 ** torch.arange(
+            0,
+            bits.shape[-1],
+            1,
+            dtype=torch.long,
+            device=bits.device,
+        )
+        return (bits * indices).sum(-1)
+
+    def forward(self, z, half=False, collect_metrics=True):
+        z = F.normalize(z, dim=-1)
+        quantized, bsq_loss, metrics = self.bsq(z, collect_metrics=collect_metrics)
+        if half:
+            q_pre = quantized[:, :, :self.s1_bits]
+            q_post = quantized[:, :, self.s1_bits:]
+            z_indices = [self.bits_to_indices(q_pre), self.bits_to_indices(q_post)]
+        else:
+            z_indices = self.bits_to_indices(quantized)
+        return bsq_loss, quantized, z_indices
+
+
+class RMSNorm(torch.nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-5):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def _norm(self, x):
+        return x * torch.rsqrt(torch.mean(x * x, dim=-1, keepdim=True) + self.eps)
+
+    def forward(self, x):
+        output = self._norm(x.float()).type_as(x)
+        return output * self.weight
+
+
+class FeedForward(nn.Module):
+    def __init__(self, d_model, ff_dim, ffn_dropout_p=0.0):
+        super().__init__()
+
+        self.w1 = nn.Linear(d_model, ff_dim, bias=False)
+        self.w3 = nn.Linear(d_model, ff_dim, bias=False)
+        self.w2 = nn.Linear(ff_dim, d_model, bias=False)
+        self.ffn_dropout = nn.Dropout(ffn_dropout_p)
+
+    def forward(self, x):
+        return self.ffn_dropout(self.w2(F.silu(self.w1(x)) * self.w3(x)))
+
+
+class RotaryPositionalEmbedding(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2).float() / dim))
+        self.register_buffer("inv_freq", inv_freq)
+        self.seq_len_cached = None
+        self.cos_cached = None
+        self.sin_cached = None
+
+    def _update_cos_sin_cache(self, x, seq_len):
+        if seq_len != self.seq_len_cached:
+            self.seq_len_cached = seq_len
+            t = torch.arange(seq_len, device=x.device).type_as(self.inv_freq)
+            freqs = torch.einsum('i,j->ij', t, self.inv_freq)
+            emb = torch.cat((freqs, freqs), dim=-1).to(x.device)
+            self.cos_cached = emb.cos()[None, None, :, :]
+            self.sin_cached = emb.sin()[None, None, :, :]
+        return self.cos_cached, self.sin_cached
+
+    def forward(self, q, k):
+        cos, sin = self._update_cos_sin_cache(q, q.shape[-2])
+        return (
+            (q * cos) + (self._rotate_half(q) * sin),
+            (k * cos) + (self._rotate_half(k) * sin),
+        )
+
+    def _rotate_half(self, x):
+        x1, x2 = x.chunk(2, dim=-1)
+        return torch.cat((-x2, x1), dim=-1)
+
+
+class MultiHeadAttentionWithRoPE(nn.Module):
+    def __init__(self, d_model, n_heads, attn_dropout_p=0.0, resid_dropout_p=0.0):
+        super().__init__()
+        self.d_model = d_model
+        self.n_heads = n_heads
+        self.head_dim = d_model // n_heads
+
+        self.q_proj = nn.Linear(d_model, d_model)
+        self.k_proj = nn.Linear(d_model, d_model)
+        self.v_proj = nn.Linear(d_model, d_model)
+        self.out_proj = nn.Linear(d_model, d_model)
+        self.rotary = RotaryPositionalEmbedding(self.head_dim)
+        self.attn_dropout_p = attn_dropout_p
+        self.resid_dropout = nn.Dropout(resid_dropout_p)
+
+    def forward(self, x, key_padding_mask=None):
+        batch_size, seq_len, _ = x.shape
+
+        q = self.q_proj(x).view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2)
+        k = self.k_proj(x).view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2)
+        v = self.v_proj(x).view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2)
+
+        q, k = self.rotary(q, k)
+
+        if key_padding_mask is not None:
+            attn_mask = key_padding_mask.unsqueeze(1).unsqueeze(2)  # [batch, 1, 1, seq_len]
+            attn_mask = attn_mask.expand(-1, self.n_heads, seq_len, -1)  # [batch, n_heads, q_len, k_len]
+        else:
+            attn_mask = None
+
+        attn_output = F.scaled_dot_product_attention(
+            q, k, v,
+            attn_mask=attn_mask,
+            dropout_p=self.attn_dropout_p if self.training else 0.0,
+            is_causal=True
+        )
+
+        attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model)
+        return self.resid_dropout(self.out_proj(attn_output))
+
+
+class MultiHeadCrossAttentionWithRoPE(nn.Module):
+    def __init__(self, d_model, n_heads, attn_dropout_p=0.0, resid_dropout=0.0):
+        super().__init__()
+        self.d_model = d_model
+        self.n_heads = n_heads
+        self.head_dim = d_model // n_heads
+
+        self.q_proj = nn.Linear(d_model, d_model)
+        self.k_proj = nn.Linear(d_model, d_model)
+        self.v_proj = nn.Linear(d_model, d_model)
+        self.out_proj = nn.Linear(d_model, d_model)
+        self.rotary = RotaryPositionalEmbedding(self.head_dim)
+        self.attn_dropout_p = attn_dropout_p
+        self.resid_dropout = nn.Dropout(resid_dropout)
+
+    def forward(self, query, key, value, key_padding_mask=None):
+        batch_size, q_len, _ = query.shape
+        _, seq_len, _ = key.shape
+
+        q = self.q_proj(query).view(batch_size, q_len, self.n_heads, self.head_dim).transpose(1, 2)
+        k = self.k_proj(key).view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2)
+        v = self.v_proj(value).view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2)
+
+        q, k = self.rotary(q, k)
+
+        if key_padding_mask is not None:
+            attn_mask = key_padding_mask.unsqueeze(1).unsqueeze(2)
+            attn_mask = attn_mask.expand(-1, self.n_heads, q_len, -1)
+        else:
+            attn_mask = None
+
+        is_causal_flag = self.training
+
+        attn_output = F.scaled_dot_product_attention(
+            q, k, v,
+            attn_mask=attn_mask,
+            dropout_p=self.attn_dropout_p if self.training else 0.0,
+            is_causal=is_causal_flag
+        )
+
+        attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, q_len, self.d_model)
+        return self.resid_dropout(self.out_proj(attn_output))
+
+
+class HierarchicalEmbedding(nn.Module):
+    def __init__(self, s1_bits, s2_bits, d_model=256):
+        super().__init__()
+        self.s1_bits = s1_bits
+        self.s2_bits = s2_bits
+
+        vocab_s1 = 2 ** s1_bits
+        vocab_s2 = 2 ** s2_bits
+
+        self.emb_s1 = nn.Embedding(vocab_s1, d_model)
+        self.emb_s2 = nn.Embedding(vocab_s2, d_model)
+        self.d_model = d_model
+        self.fusion_proj = nn.Linear(d_model * 2, d_model)
+
+        nn.init.normal_(self.emb_s1.weight, mean=0, std=d_model ** -0.5)
+        nn.init.normal_(self.emb_s2.weight, mean=0, std=d_model ** -0.5)
+
+    def split_token(self, token_ids: torch.Tensor, s2_bits: int):
+        """Inputs:
+            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].
+            s2_bits (int): Number of low bits used for the fine token (s2).
+        """
+        assert isinstance(s2_bits, int) and s2_bits > 0, "s2_bits must be a positive integer"
+
+        t = token_ids.long()
+        mask = (1 << s2_bits) - 1
+        s2_ids = t & mask           # extract low bits
+        s1_ids = t >> s2_bits       # extract high bits
+        return s1_ids, s2_ids
+
+    def forward(self, token_ids):
+        """Inputs:
+        token_ids:
+            - tuple or list: (s1_ids, s2_ids), each of shape [batch_size, seq_len], or
+            - torch.Tensor: composite token IDs of shape [batch_size, seq_len], which will be split into (s1_ids, s2_ids) internally.
+        Output: [batch_size, seq_len, d_model]
+        """
+        if isinstance(token_ids, tuple) or isinstance(token_ids, list):
+            s1_ids, s2_ids = token_ids
+        else:
+            s1_ids, s2_ids = self.split_token(token_ids, self.s2_bits)
+        s1_emb = self.emb_s1(s1_ids) * math.sqrt(self.d_model)
+        s2_emb = self.emb_s2(s2_ids) * math.sqrt(self.d_model)
+        return self.fusion_proj(torch.cat([s1_emb, s2_emb], dim=-1))
+
+
+class DependencyAwareLayer(nn.Module):
+    def __init__(self, d_model, n_heads=4, attn_dropout_p=0.0, resid_dropout=0.0):
+        super().__init__()
+        self.cross_attn = MultiHeadCrossAttentionWithRoPE(d_model, n_heads, attn_dropout_p, resid_dropout)
+        self.norm = RMSNorm(d_model)
+
+    def forward(self, hidden_states, sibling_embed, key_padding_mask=None):
+        """hidden_states: [batch, seq_len, d_model]
+        sibling_embed: Embedding from another subtoken
+        """
+        attn_out = self.cross_attn(
+            query=sibling_embed,
+            key=hidden_states,
+            value=hidden_states,
+            key_padding_mask=key_padding_mask
+        )
+        return self.norm(hidden_states + attn_out)
+
+
+class TransformerBlock(nn.Module):
+    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):
+        super().__init__()
+        self.norm1 = RMSNorm(d_model)
+        self.self_attn = MultiHeadAttentionWithRoPE(d_model, n_heads, attn_dropout_p, resid_dropout_p)
+        self.norm2 = RMSNorm(d_model)
+        self.ffn = FeedForward(d_model, ff_dim, ffn_dropout_p)
+
+    def forward(self, x, key_padding_mask=None):
+        residual = x
+        x = self.norm1(x)
+        attn_out = self.self_attn(x, key_padding_mask=key_padding_mask)
+        x = residual + attn_out
+
+        residual = x
+        x = self.norm2(x)
+        ffn_out = self.ffn(x)
+        x = residual + ffn_out
+        return x
+
+
+class DualHead(nn.Module):
+    def __init__(self, s1_bits, s2_bits, d_model):
+        super().__init__()
+        self.vocab_s1 = 2 ** s1_bits
+        self.vocab_s2 = 2 ** s2_bits
+        self.proj_s1 = nn.Linear(d_model, self.vocab_s1)
+        self.proj_s2 = nn.Linear(d_model, self.vocab_s2)
+
+    def compute_loss(self, s1_logits, s2_logits, s1_targets, s2_targets, padding_mask=None):
+        if padding_mask is not None:
+            valid_mask = (padding_mask == 0)
+            s1_logits = s1_logits[valid_mask]
+            s2_logits = s2_logits[valid_mask]
+            s1_targets = s1_targets[valid_mask]
+            s2_targets = s2_targets[valid_mask]
+            ce_s1 = F.cross_entropy(s1_logits, s1_targets)
+            ce_s2 = F.cross_entropy(s2_logits, s2_targets)
+        else:
+            ce_s1 = F.cross_entropy(s1_logits.reshape(-1, self.vocab_s1), s1_targets.reshape(-1))
+            ce_s2 = F.cross_entropy(s2_logits.reshape(-1, self.vocab_s2), s2_targets.reshape(-1))
+        ce_loss = (ce_s1 + ce_s2) / 2
+        return ce_loss, ce_s1, ce_s2
+
+    def forward(self, x):
+        return self.proj_s1(x)
+
+    def cond_forward(self, x2):
+        return self.proj_s2(x2)
+
+
+class FixedEmbedding(nn.Module):
+    def __init__(self, c_in, d_model):
+        super(FixedEmbedding, self).__init__()
+
+        w = torch.zeros(c_in, d_model).float()
+        w.require_grad = False
+
+        position = torch.arange(0, c_in).float().unsqueeze(1)
+        div_term = (torch.arange(0, d_model, 2).float() * -(math.log(10000.0) / d_model)).exp()
+
+        w[:, 0::2] = torch.sin(position * div_term)
+        w[:, 1::2] = torch.cos(position * div_term)
+
+        self.emb = nn.Embedding(c_in, d_model)
+        self.emb.weight = nn.Parameter(w, requires_grad=False)
+
+    def forward(self, x):
+        return self.emb(x).detach()
+
+
+class TemporalEmbedding(nn.Module):
+    def __init__(self, d_model, learn_pe):
+        super(TemporalEmbedding, self).__init__()
+
+        minute_size = 60
+        hour_size = 24
+        weekday_size = 7
+        day_size = 32
+        month_size = 13
+
+        Embed = FixedEmbedding if not learn_pe else nn.Embedding
+        self.minute_embed = Embed(minute_size, d_model)
+        self.hour_embed = Embed(hour_size, d_model)
+        self.weekday_embed = Embed(weekday_size, d_model)
+        self.day_embed = Embed(day_size, d_model)
+        self.month_embed = Embed(month_size, d_model)
+
+    def forward(self, x):
+        x = x.long()
+
+        minute_x = self.minute_embed(x[:, :, 0])
+        hour_x = self.hour_embed(x[:, :, 1])
+        weekday_x = self.weekday_embed(x[:, :, 2])
+        day_x = self.day_embed(x[:, :, 3])
+        month_x = self.month_embed(x[:, :, 4])
+
+        return hour_x + weekday_x + day_x + month_x + minute_x
+
+
+
+
+
+
+
+

models/predictor/README.md

@@ -0,0 +1,10 @@
+---
+tags:
+- model_hub_mixin
+- pytorch_model_hub_mixin
+---
+
+This model has been pushed to the Hub using the [PytorchModelHubMixin](https://huggingface.co/docs/huggingface_hub/package_reference/mixins#huggingface_hub.PyTorchModelHubMixin) integration:
+- Code: [More Information Needed]
+- Paper: [More Information Needed]
+- Docs: [More Information Needed]

models/predictor/config.json

@@ -0,0 +1,13 @@
+{
+  "attn_dropout_p": 0.0,
+  "d_model": 256,
+  "ff_dim": 512,
+  "ffn_dropout_p": 0.2,
+  "learn_te": true,
+  "n_heads": 4,
+  "n_layers": 4,
+  "resid_dropout_p": 0.2,
+  "s1_bits": 10,
+  "s2_bits": 10,
+  "token_dropout_p": 0.0
+}

models/predictor/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0fb3d6ca23cb183f692e232ba427c801f1b706d862ae3226241c9a6ba86796fc
+size 16440776

models/tokenizer/README.md

@@ -0,0 +1,10 @@
+---
+tags:
+- model_hub_mixin
+- pytorch_model_hub_mixin
+---
+
+This model has been pushed to the Hub using the [PytorchModelHubMixin](https://huggingface.co/docs/huggingface_hub/package_reference/mixins#huggingface_hub.PyTorchModelHubMixin) integration:
+- Code: [More Information Needed]
+- Paper: [More Information Needed]
+- Docs: [More Information Needed]

models/tokenizer/config.json

@@ -0,0 +1,18 @@
+{
+  "attn_dropout_p": 0.0,
+  "beta": 0.05,
+  "d_in": 6,
+  "d_model": 256,
+  "ff_dim": 512,
+  "ffn_dropout_p": 0.0,
+  "gamma": 1.1,
+  "gamma0": 1.0,
+  "group_size": 5,
+  "n_dec_layers": 4,
+  "n_enc_layers": 4,
+  "n_heads": 4,
+  "resid_dropout_p": 0.0,
+  "s1_bits": 10,
+  "s2_bits": 10,
+  "zeta": 0.05
+}

models/tokenizer/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f559eb686c32b87b12b63bdadcdf069e53690f05c31b4c3bd810ba3176229f6a
+size 15842376

requirements.txt

@@ -0,0 +1,23 @@
+# Kronos BTC Prediction API - Dependencies
+# Optimized for HuggingFace Spaces (CPU)
+
+# Core
+fastapi==0.104.1
+uvicorn[standard]==0.24.0
+pydantic==2.5.2
+
+# ML/DL
+torch==2.1.0
+numpy>=1.24.0,<2.0.0
+pandas>=2.0.0
+
+# Model loading
+safetensors>=0.4.0
+huggingface-hub>=0.19.0
+
+# HTTP client (for client SDK)
+httpx>=0.25.0
+aiohttp>=3.9.0
+
+# Utilities
+python-dateutil>=2.8.0