""" grid_seeker v6.6 三级模型推理管道 Rank → Top → Stacking → stacking_probability (最终排序) """ import pickle import numpy as np import pandas as pd from pathlib import Path from datetime import date from core.scoring.config import ( get_model_path, RANK_MODEL, TOP_MODEL, STACKING_MODEL, STACKING_THRESHOLD, ) from core.scoring.features.pipeline import FeaturePipeline from core.scoring.models import ScoringResult from core.database import db from core.logger import LogLevel, PrintLog # ============================================================ # Rank 模型输入特征 (52维, v3.4, 直接从模型文件的 selected_features 读取) # ============================================================ def _get_rank_features() -> list: import pickle from core.scoring.config import get_model_path path = get_model_path(RANK_MODEL) with open(path, 'rb') as f: obj = pickle.load(f) if isinstance(obj, dict): sf = obj.get('selected_features', []) if sf: return sf raise RuntimeError("无法从 rank.pkl 读取 selected_features") RANK_FEATURE_COLS = _get_rank_features() class GridSeekerPipeline: """ grid_seeker v6.6 三级模型评分管道。 Usage: engine = GridSeekerPipeline() rankings = engine.run(trade_date=date.today()) # 返回 DataFrame: stock_code, stacking_probability, rank 等 """ def __init__(self, model_dir: Path = None): self._rank_model = None self._top_model = None self._stacking_model = None # ---- 模型加载 ---- def _load_model(self, name: str): """加载单个 .pkl 模型""" path = get_model_path(name) if not path.exists(): raise FileNotFoundError(f'模型文件不存在: {path}') with open(path, 'rb') as f: obj = pickle.load(f) # 支持 dict 格式 {"model": lgbm_model, ...} 或直接返回模型对象 if isinstance(obj, dict): return obj.get('model', obj) return obj @property def rank_model(self): if self._rank_model is None: self._rank_model = self._load_model(RANK_MODEL) return self._rank_model @property def top_model(self): if self._top_model is None: self._top_model = self._load_model(TOP_MODEL) return self._top_model @property def stacking_model(self): if self._stacking_model is None: self._stacking_model = self._load_model(STACKING_MODEL) return self._stacking_model # ---- 预测 ---- def _predict_with_model(self, model, X: pd.DataFrame, feature_cols: list) -> np.ndarray: """ 使用模型预测。自动选择特征子集，兼容 sklearn API (predict/predict_proba)。 """ available = [c for c in feature_cols if c in X.columns] missing = set(feature_cols) - set(available) if missing: PrintLog(LogLevel.WARNING, f'[scorer] 缺少特征列 ({len(missing)}): {list(missing)[:5]}...') X_sub = X[available].fillna(0).values try: if hasattr(model, 'predict_proba'): proba = model.predict_proba(X_sub) if proba.shape[1] >= 2: return proba[:, 1] return proba[:, 0] elif hasattr(model, 'predict'): return model.predict(X_sub) else: return model.predict(X_sub) except Exception as e: PrintLog(LogLevel.ERROR, f'[scorer] 模型预测失败: {e}') raise # ---- 主流程 ---- def run(self, trade_date: date) -> pd.DataFrame: """ 执行完整的 3 级评分管道。 Returns: DataFrame indexed by stock_code, 含 stacking_probability / rank 等列, 按 stacking_probability 降序排列 """ PrintLog(LogLevel.INFO, f'[scorer] ===== grid_seeker v6.6 评分开始 ({trade_date}) =====') # 1. 特征工程 pipeline = FeaturePipeline(trade_date) feature_df = pipeline.run() if feature_df.empty: PrintLog(LogLevel.WARNING, '[scorer] 无股票通过特征工程, 终止') return pd.DataFrame() PrintLog(LogLevel.INFO, f'[scorer] 特征工程完成: {len(feature_df)} stocks, ' f'{len(feature_df.columns)} dims') # 2. Stage 1: Rank 模型 → rank_predicted_rounds (52维) PrintLog(LogLevel.INFO, '[scorer] Stage 1/3: Rank 模型...') feature_df['rank_predicted_rounds'] = self._predict_with_model( self.rank_model, feature_df, RANK_FEATURE_COLS ) # 3. Stage 2: Top 模型 → top_elite_prob (53维 = 52 + rank_predicted_rounds) PrintLog(LogLevel.INFO, '[scorer] Stage 2/3: Top 模型...') top_cols = RANK_FEATURE_COLS + ['rank_predicted_rounds'] feature_df['top_elite_prob'] = self._predict_with_model( self.top_model, feature_df, top_cols ) # 4. Stage 3: Stacking 模型 → stacking_probability (54维 = 52 + rank + top) PrintLog(LogLevel.INFO, '[scorer] Stage 3/3: Stacking 模型...') stk_cols = RANK_FEATURE_COLS + ['rank_predicted_rounds', 'top_elite_prob'] feature_df['stacking_probability'] = self._predict_with_model( self.stacking_model, feature_df, stk_cols ) # 5. 排序（直接用 stacking_probability） feature_df['score_rank'] = feature_df['stacking_probability'].rank( ascending=False, method='min' ).astype(int) feature_df['candidate_count'] = len(feature_df) feature_df = feature_df.sort_values('score_rank') n_above = (feature_df['stacking_probability'] >= STACKING_THRESHOLD).sum() PrintLog(LogLevel.INFO, f'[scorer] 评分完成: {len(feature_df)} 只候选, ' f'{n_above} 只高于阈值 {STACKING_THRESHOLD}') PrintLog(LogLevel.INFO, f'[scorer] Top-5: ' f'{feature_df.head(5)[["stacking_probability", "rank_predicted_rounds"]].to_dict("index")}') return feature_df def persist(self, rankings: pd.DataFrame, trade_date: date): """将评分结果持久化到 ScoringResult 表""" if rankings.empty: return records = [] for code, row in rankings.iterrows(): records.append({ 'stock_code': str(code), 'trade_date': trade_date, 'predicted_profit': float(row.get('stacking_probability', 0)), 'rank_predicted_rounds': float(row.get('rank_predicted_rounds', 0)) if 'rank_predicted_rounds' in row else None, 'top_elite_prob': float(row.get('top_elite_prob', 0)) if 'top_elite_prob' in row else None, 'stacking_probability': float(row.get('stacking_probability', 0)) if 'stacking_probability' in row else None, 'score_rank': int(row.get('score_rank', 0)), 'candidate_count': int(row.get('candidate_count', 0)), }) with db.atomic(): for batch in _chunked(records, 500): ScoringResult.insert_many(batch).on_conflict_replace().execute() PrintLog(LogLevel.INFO, f'[scorer] 评分结果已持久化: {len(records)} 条') def get_top_n(self, trade_date: date, n: int = 50) -> list[dict]: """查询历史评分 Top-N""" rows = (ScoringResult .select() .where( (ScoringResult.trade_date == trade_date) & (ScoringResult.score_rank <= n) ) .order_by(ScoringResult.score_rank) .dicts()) return list(rows) def _chunked(lst: list, n: int): for i in range(0, len(lst), n): yield lst[i:i + n]