feat(signal_v2-phase3b): chronos_predictor + 4 mock tests

ChronosPredictor wraps HuggingFace ChronosPipeline. Batch predict
returns ChronosPrediction(median, q10, q90, conf, as_of) per ticker.
Confidence = 1 - clamp(spread/2, 0, 1) where spread = (q90-q10) / |median|.
Lazy import of chronos lib (heavy). GPU auto-detect with FP16.

44 tests pass.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
2026-05-16 18:00:46 +09:00
parent c5a88fab66
commit 28f9c8c3a6
2 changed files with 161 additions and 0 deletions

View File

@@ -0,0 +1,78 @@
"""Chronos-2 zero-shot forecaster wrapper."""
from __future__ import annotations
import logging
from dataclasses import dataclass
from datetime import datetime
from zoneinfo import ZoneInfo
import numpy as np
logger = logging.getLogger(__name__)
KST = ZoneInfo("Asia/Seoul")
@dataclass
class ChronosPrediction:
median: float
q10: float
q90: float
conf: float
as_of: str
class ChronosPredictor:
"""HuggingFace Chronos-2 zero-shot forecaster."""
def __init__(self, model_name: str = "amazon/chronos-2", device: str | None = None):
from chronos import ChronosPipeline
import torch
self._device = device or ("cuda" if torch.cuda.is_available() else "cpu")
logger.info("Loading Chronos pipeline: %s on %s", model_name, self._device)
self._pipeline = ChronosPipeline.from_pretrained(
model_name,
device_map=self._device,
torch_dtype=torch.float16 if self._device == "cuda" else torch.float32,
)
logger.info("Chronos pipeline loaded.")
def predict_batch(
self,
daily_ohlcv_dict: dict[str, list[dict]],
prediction_length: int = 1,
num_samples: int = 100,
) -> dict[str, ChronosPrediction]:
"""종목별 1-day return 분포 예측."""
import torch
tickers = list(daily_ohlcv_dict.keys())
if not tickers:
return {}
contexts = [
torch.tensor([bar["close"] for bar in daily_ohlcv_dict[t]], dtype=torch.float32)
for t in tickers
]
forecasts = self._pipeline.predict(
context=contexts,
prediction_length=prediction_length,
num_samples=num_samples,
)
# Convert to numpy if tensor
forecasts_np = forecasts.numpy() if hasattr(forecasts, "numpy") else np.asarray(forecasts)
now_iso = datetime.now(KST).isoformat()
results: dict[str, ChronosPrediction] = {}
for i, ticker in enumerate(tickers):
samples = forecasts_np[i, :, 0]
last_close = daily_ohlcv_dict[ticker][-1]["close"]
returns = (samples - last_close) / last_close
median = float(np.quantile(returns, 0.5))
q10 = float(np.quantile(returns, 0.1))
q90 = float(np.quantile(returns, 0.9))
spread = (q90 - q10) / max(abs(median), 0.001)
conf = float(max(0.0, min(1.0, 1.0 - spread / 2.0)))
results[ticker] = ChronosPrediction(
median=median, q10=q10, q90=q90, conf=conf, as_of=now_iso,
)
return results