liaozhaorun
990a77ec6c
- 添加 FactorConfig 类，支持链式 API 和动态因子列表 - 重构 prepare_data 使用 FactorEngine 计算因子，确保防泄露机制生效 - 新增 prepare_data_and_train / train_and_predict 分步执行接口 - 修复 factors/__init__.py docstring 语法错误
668 lines
20 KiB
Python
668 lines
20 KiB
Python
"""训练管道 - 包含数据处理、模型训练和预测功能
|
||
|
||
本模块提供：
|
||
1. 数据准备：使用 FactorEngine 从因子计算结果中准备训练/测试数据
|
||
2. 数据处理：Fillna(0) -> Dropna
|
||
3. 模型训练：使用LightGBM训练分类模型
|
||
4. 预测和选股：输出每日top5股票池
|
||
|
||
注意：本模块使用 src.factors 框架进行因子计算，确保防泄露机制生效。
|
||
|
||
因子配置示例：
|
||
from src.factors import MovingAverageFactor, ReturnRankFactor
|
||
from src.training.pipeline import prepare_data
|
||
|
||
# 直接传入因子实例列表 - 简单直观
|
||
factors = [
|
||
MovingAverageFactor(period=5),
|
||
MovingAverageFactor(period=10),
|
||
ReturnRankFactor(period=5),
|
||
]
|
||
|
||
train_data, val_data, test_data, factor_config = prepare_data(
|
||
factors=factors,
|
||
...
|
||
)
|
||
|
||
# 或者使用 FactorConfig 包装（支持链式添加）
|
||
from src.training.pipeline import FactorConfig
|
||
|
||
config = FactorConfig()
|
||
.add(MovingAverageFactor(period=5))
|
||
.add(MovingAverageFactor(period=10))
|
||
.add(ReturnRankFactor(period=5))
|
||
"""
|
||
|
||
from datetime import datetime
|
||
from pathlib import Path
|
||
from typing import List, Optional
|
||
|
||
import numpy as np
|
||
import polars as pl
|
||
|
||
from src.factors import DataLoader, FactorEngine, BaseFactor
|
||
from src.factors.data_spec import DataSpec
|
||
from src.factors.momentum import MovingAverageFactor, ReturnRankFactor
|
||
from src.pipeline import (
|
||
DropNAProcessor,
|
||
FillNAProcessor,
|
||
LightGBMModel,
|
||
PipelineStage,
|
||
ProcessingPipeline,
|
||
TaskType,
|
||
)
|
||
|
||
|
||
# ========== 因子配置类 ==========
|
||
|
||
|
||
class FactorConfig:
|
||
"""因子配置类 - 管理因子列表
|
||
|
||
用于包装因子实例列表，支持链式添加。
|
||
|
||
示例：
|
||
# 方式1：初始化时传入列表
|
||
config = FactorConfig([
|
||
MovingAverageFactor(period=5),
|
||
ReturnRankFactor(period=5),
|
||
])
|
||
|
||
# 方式2：链式添加
|
||
config = FactorConfig()
|
||
.add(MovingAverageFactor(period=5))
|
||
.add(ReturnRankFactor(period=5))
|
||
|
||
# 获取因子实例列表
|
||
factors = config.get_factors()
|
||
|
||
# 获取特征列名
|
||
feature_cols = config.get_feature_names()
|
||
"""
|
||
|
||
def __init__(self, factors: Optional[List[BaseFactor]] = None):
|
||
"""初始化因子配置
|
||
|
||
Args:
|
||
factors: 因子实例列表
|
||
"""
|
||
self._factors: List[BaseFactor] = factors or []
|
||
|
||
def add(self, factor: BaseFactor) -> "FactorConfig":
|
||
"""添加因子到配置
|
||
|
||
支持链式调用：
|
||
config = FactorConfig()
|
||
.add(MovingAverageFactor(period=5))
|
||
.add(ReturnRankFactor(period=5))
|
||
|
||
Args:
|
||
factor: 因子实例
|
||
|
||
Returns:
|
||
self，支持链式调用
|
||
"""
|
||
if not isinstance(factor, BaseFactor):
|
||
raise ValueError(f"必须是 BaseFactor 实例， got {type(factor)}")
|
||
self._factors.append(factor)
|
||
return self
|
||
|
||
def get_factors(self) -> List[BaseFactor]:
|
||
"""获取因子实例列表
|
||
|
||
Returns:
|
||
因子实例列表
|
||
"""
|
||
return self._factors
|
||
|
||
def get_feature_names(self) -> List[str]:
|
||
"""获取所有因子的特征列名
|
||
|
||
Returns:
|
||
特征列名列表
|
||
"""
|
||
return [f.name for f in self._factors]
|
||
|
||
def get_max_lookback(self) -> int:
|
||
"""获取所有因子中最大的 lookback 天数
|
||
|
||
Returns:
|
||
最大 lookback 天数
|
||
"""
|
||
max_lookback = 0
|
||
for factor in self._factors:
|
||
for spec in factor.data_specs:
|
||
max_lookback = max(max_lookback, spec.lookback_days)
|
||
return max_lookback
|
||
|
||
def __len__(self) -> int:
|
||
return len(self._factors)
|
||
|
||
def __repr__(self) -> str:
|
||
names = [f.name for f in self._factors]
|
||
return f"FactorConfig({names})"
|
||
|
||
|
||
def prepare_data(
|
||
factors: Optional[List[BaseFactor]] = None,
|
||
data_dir: str = "data",
|
||
train_start: str = "20180101",
|
||
train_end: str = "20230101",
|
||
val_start: str = "20230101",
|
||
val_end: str = "20230601",
|
||
test_start: str = "20230601",
|
||
test_end: str = "20240101",
|
||
) -> tuple[pl.DataFrame, pl.DataFrame, pl.DataFrame, FactorConfig]:
|
||
"""准备训练、验证和测试数据
|
||
|
||
使用 FactorEngine 计算因子，确保防泄露机制生效。
|
||
|
||
Args:
|
||
factors: 因子实例列表，默认为 None（使用 MA5, MA10, ReturnRank5）
|
||
data_dir: 数据目录
|
||
train_start: 训练集开始日期
|
||
train_end: 训练集结束日期
|
||
val_start: 验证集开始日期
|
||
val_end: 验证集结束日期
|
||
test_start: 测试集开始日期
|
||
test_end: 测试集结束日期
|
||
|
||
Returns:
|
||
(train_data, val_data, test_data, factor_config):
|
||
训练集、验证集、测试集的DataFrame，以及使用的因子配置
|
||
"""
|
||
from src.data.storage import Storage
|
||
|
||
storage = Storage()
|
||
|
||
# 1. 处理因子配置
|
||
if factors is None:
|
||
# 默认因子配置
|
||
factor_config = FactorConfig(
|
||
[
|
||
MovingAverageFactor(period=5),
|
||
MovingAverageFactor(period=10),
|
||
ReturnRankFactor(period=5),
|
||
]
|
||
)
|
||
elif isinstance(factors, FactorConfig):
|
||
factor_config = factors
|
||
elif isinstance(factors, list):
|
||
# 转换为 FactorConfig
|
||
factor_config = FactorConfig(factors)
|
||
else:
|
||
raise ValueError(
|
||
f"factors 必须是 List[BaseFactor] 或 FactorConfig， got {type(factors)}"
|
||
)
|
||
|
||
factors_list = factor_config.get_factors()
|
||
feature_cols = factor_config.get_feature_names()
|
||
|
||
if not factors_list:
|
||
raise ValueError("至少需要提供一个因子")
|
||
|
||
print(f"[PrepareData] 使用因子: {feature_cols}")
|
||
|
||
# 2. 初始化 FactorEngine
|
||
loader = DataLoader(data_dir=data_dir)
|
||
engine = FactorEngine(loader)
|
||
|
||
# 3. 计算需要的回溯天数
|
||
max_lookback = factor_config.get_max_lookback()
|
||
start_with_lookback = str(int(train_start) - 10000) # 往前取一年
|
||
|
||
# 获取所有股票列表
|
||
start_dt = f"{start_with_lookback[:4]}-{start_with_lookback[4:6]}-{start_with_lookback[6:8]}"
|
||
end_dt = f"{test_end[:4]}-{test_end[4:6]}-{test_end[6:8]}"
|
||
all_stocks_query = f"""
|
||
SELECT DISTINCT ts_code FROM daily
|
||
WHERE trade_date >= '{start_dt}' AND trade_date <= '{end_dt}'
|
||
"""
|
||
all_stocks_df = storage._connection.sql(all_stocks_query).pl()
|
||
all_stocks = all_stocks_df["ts_code"].to_list()
|
||
|
||
print(f"[PrepareData] 股票数量: {len(all_stocks)}")
|
||
|
||
# 4. 计算所有因子并合并
|
||
all_features = None
|
||
|
||
for factor in factors_list:
|
||
print(f"[PrepareData] 计算因子: {factor.name} ({factor.factor_type})")
|
||
|
||
if factor.factor_type == "time_series":
|
||
# 时序因子需要传入股票列表
|
||
result = engine.compute(
|
||
factor,
|
||
stock_codes=all_stocks,
|
||
start_date=start_with_lookback,
|
||
end_date=test_end,
|
||
)
|
||
else:
|
||
# 截面因子不需要股票列表
|
||
result = engine.compute(
|
||
factor,
|
||
start_date=start_with_lookback,
|
||
end_date=test_end,
|
||
)
|
||
|
||
# 合并结果
|
||
if all_features is None:
|
||
all_features = result
|
||
else:
|
||
# 确保没有重复的 _right 列
|
||
result = result.select([
|
||
c for c in result.columns if not c.endswith("_right")
|
||
])
|
||
all_features = all_features.select([
|
||
c for c in all_features.columns if not c.endswith("_right")
|
||
])
|
||
all_features = all_features.join(
|
||
result, on=["trade_date", "ts_code"], how="outer"
|
||
)
|
||
|
||
if all_features is None:
|
||
raise ValueError("没有计算任何因子")
|
||
|
||
# 5. 计算标签（未来5日收益率）
|
||
all_data = _compute_label(all_features, start_date=train_start, end_date=test_end)
|
||
|
||
# 6. 过滤不符合条件的股票
|
||
all_data = _filter_invalid_stocks(all_data)
|
||
print(f"[PrepareData] After filtering: total={len(all_data)}")
|
||
|
||
# 转换日期格式用于比较
|
||
train_start_fmt = f"{train_start[:4]}-{train_start[4:6]}-{train_start[6:8]}"
|
||
train_end_fmt = f"{train_end[:4]}-{train_end[4:6]}-{train_end[6:8]}"
|
||
val_start_fmt = f"{val_start[:4]}-{val_start[4:6]}-{val_start[6:8]}"
|
||
val_end_fmt = f"{val_end[:4]}-{val_end[4:6]}-{val_end[6:8]}"
|
||
test_start_fmt = f"{test_start[:4]}-{test_start[4:6]}-{test_start[6:8]}"
|
||
test_end_fmt = f"{test_end[:4]}-{test_end[4:6]}-{test_end[6:8]}"
|
||
|
||
# 拆分数据
|
||
train_data = all_data.filter(
|
||
(pl.col("trade_date") >= train_start_fmt)
|
||
& (pl.col("trade_date") <= train_end_fmt)
|
||
)
|
||
val_data = all_data.filter(
|
||
(pl.col("trade_date") >= val_start_fmt) & (pl.col("trade_date") <= val_end_fmt)
|
||
)
|
||
test_data = all_data.filter(
|
||
(pl.col("trade_date") >= test_start_fmt)
|
||
& (pl.col("trade_date") <= test_end_fmt)
|
||
)
|
||
|
||
print(
|
||
f"[PrepareData] Split result: train={len(train_data)}, val={len(val_data)}, test={len(test_data)}"
|
||
)
|
||
|
||
return train_data, val_data, test_data, factor_config
|
||
|
||
|
||
def _filter_invalid_stocks(df: pl.DataFrame) -> pl.DataFrame:
|
||
"""过滤不符合条件的股票
|
||
|
||
过滤规则：
|
||
1. 过滤北交所股票（ts_code 以 BJ 结尾）
|
||
2. 过滤创业板股票（ts_code 以 30 开头）
|
||
3. 过滤科创板股票（ts_code 以 68 开头）
|
||
4. 过滤退市/风险股票（ts_code 以 8 开头）
|
||
|
||
Args:
|
||
df: 原始数据
|
||
|
||
Returns:
|
||
过滤后的数据
|
||
"""
|
||
ts_code_col = pl.col("ts_code")
|
||
|
||
return df.filter(
|
||
~ts_code_col.str.ends_with("BJ")
|
||
& ~ts_code_col.str.starts_with("30")
|
||
& ~ts_code_col.str.starts_with("68")
|
||
& ~ts_code_col.str.starts_with("8")
|
||
)
|
||
|
||
|
||
def _compute_label(
|
||
features_df: pl.DataFrame,
|
||
start_date: str,
|
||
end_date: str,
|
||
) -> pl.DataFrame:
|
||
"""计算标签（未来5日收益率）
|
||
|
||
标签定义：未来5日收益率大于0为1，否则为0
|
||
|
||
Args:
|
||
features_df: 包含因子的DataFrame
|
||
start_date: 开始日期
|
||
end_date: 结束日期
|
||
|
||
Returns:
|
||
包含因子和标签的DataFrame
|
||
"""
|
||
from src.data.storage import Storage
|
||
|
||
storage = Storage()
|
||
|
||
# 从数据库获取收盘价数据用于计算标签
|
||
start_dt = f"{start_date[:4]}-{start_date[4:6]}-{start_date[6:8]}"
|
||
end_dt = f"{end_date[:4]}-{end_date[4:6]}-{end_date[6:8]}"
|
||
|
||
# 需要多取5天数据来计算未来收益率
|
||
end_dt_extended = f"{end_date[:4]}-{end_date[4:6]}-{int(end_date[6:8]) + 5}"
|
||
|
||
price_query = f"""
|
||
SELECT ts_code, trade_date, close
|
||
FROM daily
|
||
WHERE trade_date >= '{start_dt}' AND trade_date <= '{end_dt_extended}'
|
||
ORDER BY ts_code, trade_date
|
||
"""
|
||
price_data = storage._connection.sql(price_query).pl()
|
||
price_data = price_data.with_columns(
|
||
pl.col("trade_date").dt.strftime("%Y-%m-%d").alias("trade_date")
|
||
)
|
||
|
||
# 按股票计算未来5日收益率
|
||
result_list = []
|
||
for ts_code in price_data["ts_code"].unique():
|
||
stock_data = price_data.filter(pl.col("ts_code") == ts_code).sort("trade_date")
|
||
|
||
if len(stock_data) < 6:
|
||
continue
|
||
|
||
# 计算未来5日收益率
|
||
future_return = stock_data["close"].shift(-5) - stock_data["close"]
|
||
future_return_pct = future_return / stock_data["close"]
|
||
stock_data = stock_data.with_columns(
|
||
[
|
||
future_return_pct.alias("future_return_5"),
|
||
]
|
||
)
|
||
|
||
# 生成标签：收益率>0为1，否则为0
|
||
stock_data = stock_data.with_columns(
|
||
[
|
||
(pl.col("future_return_5") > 0).cast(pl.Int8).alias("label"),
|
||
]
|
||
)
|
||
|
||
result_list.append(stock_data.select(["trade_date", "ts_code", "label"]))
|
||
|
||
if not result_list:
|
||
return pl.DataFrame()
|
||
|
||
label_df = pl.concat(result_list)
|
||
|
||
# 将标签合并到因子数据
|
||
result = features_df.join(label_df, on=["trade_date", "ts_code"], how="inner")
|
||
|
||
# 过滤有效日期范围
|
||
result = result.filter(
|
||
(pl.col("trade_date") >= start_dt) & (pl.col("trade_date") <= end_dt)
|
||
)
|
||
|
||
return result
|
||
|
||
|
||
def create_pipeline() -> ProcessingPipeline:
|
||
"""创建数据处理流水线
|
||
|
||
处理流程：
|
||
1. FillNA(0): 将缺失值填充为0
|
||
|
||
注意：不使用 Dropna，因为会导致训练和预测时的行数不匹配
|
||
|
||
Returns:
|
||
配置好的ProcessingPipeline
|
||
"""
|
||
processors = [
|
||
FillNAProcessor(method="zero"), # 缺失值填充为0
|
||
]
|
||
return ProcessingPipeline(processors)
|
||
|
||
|
||
def train_model(
|
||
train_data: pl.DataFrame,
|
||
val_data: Optional[pl.DataFrame],
|
||
feature_cols: List[str],
|
||
label_col: str = "label",
|
||
model_params: Optional[dict] = None,
|
||
) -> tuple[LightGBMModel, ProcessingPipeline]:
|
||
"""训练LightGBM分类模型
|
||
|
||
Args:
|
||
train_data: 训练数据
|
||
val_data: 验证数据（用于早停）
|
||
feature_cols: 特征列名列表
|
||
label_col: 标签列名
|
||
model_params: 模型参数字典
|
||
|
||
Returns:
|
||
(训练好的模型, 处理流水线)
|
||
"""
|
||
# 创建处理流水线
|
||
pipeline = create_pipeline()
|
||
print("[TrainModel] Pipeline created: FillNA(0)")
|
||
|
||
# 准备训练特征和标签
|
||
X_train = train_data.select(feature_cols)
|
||
y_train = train_data[label_col]
|
||
print(f"[TrainModel] Train samples: {len(X_train)}, features: {feature_cols}")
|
||
|
||
# 处理训练数据
|
||
X_train_processed = pipeline.fit_transform(X_train, stage=PipelineStage.TRAIN)
|
||
print(f"[TrainModel] After processing: {len(X_train_processed)} samples")
|
||
|
||
# 过滤训练集有效标签（排除-1等无效值）
|
||
valid_mask = y_train.is_in([0, 1])
|
||
X_train_processed = X_train_processed.filter(valid_mask)
|
||
y_train = y_train.filter(valid_mask)
|
||
print(
|
||
f"[TrainModel] After filtering valid labels: {len(X_train_processed)} samples"
|
||
)
|
||
print(
|
||
f"[TrainModel] Train label distribution: {dict(y_train.value_counts().sort('label').iter_rows())}"
|
||
)
|
||
|
||
# 准备验证集
|
||
X_val_processed = None
|
||
y_val = None
|
||
if val_data is not None and len(val_data) > 0:
|
||
X_val = val_data.select(feature_cols)
|
||
y_val = val_data[label_col]
|
||
print(f"[TrainModel] Val samples: {len(X_val)}")
|
||
|
||
# 处理验证集数据（使用训练集的参数）
|
||
X_val_processed = pipeline.transform(X_val, stage=PipelineStage.TEST)
|
||
|
||
# 过滤验证集有效标签
|
||
val_valid_mask = y_val.is_in([0, 1])
|
||
X_val_processed = X_val_processed.filter(val_valid_mask)
|
||
y_val = y_val.filter(val_valid_mask)
|
||
print(f"[TrainModel] Val after filtering: {len(X_val_processed)} samples")
|
||
print(
|
||
f"[TrainModel] Val label distribution: {dict(y_val.value_counts().sort('label').iter_rows())}"
|
||
)
|
||
|
||
# 创建模型
|
||
params = model_params or {
|
||
"n_estimators": 100,
|
||
"learning_rate": 0.05,
|
||
"max_depth": 5,
|
||
"num_leaves": 31,
|
||
}
|
||
print(f"[TrainModel] Model params: {params}")
|
||
model = LightGBMModel(
|
||
task_type="classification",
|
||
params=params,
|
||
)
|
||
|
||
# 训练模型（使用验证集早停）
|
||
print("[TrainModel] Training LightGBM...")
|
||
if X_val_processed is not None and y_val is not None:
|
||
print("[TrainModel] Using validation set for early stopping")
|
||
model.fit(X_train_processed, y_train, X_val_processed, y_val)
|
||
else:
|
||
model.fit(X_train_processed, y_train)
|
||
print("[TrainModel] Training completed!")
|
||
|
||
return model, pipeline
|
||
|
||
|
||
def predict_top_stocks(
|
||
model: LightGBMModel,
|
||
pipeline: ProcessingPipeline,
|
||
test_data: pl.DataFrame,
|
||
feature_cols: List[str],
|
||
top_n: int = 5,
|
||
) -> pl.DataFrame:
|
||
"""预测并选出每日top N股票
|
||
|
||
Args:
|
||
model: 训练好的模型
|
||
pipeline: 数据处理流水线
|
||
test_data: 测试数据
|
||
feature_cols: 特征列名
|
||
top_n: 每日选出的股票数量
|
||
|
||
Returns:
|
||
包含日期和股票代码的DataFrame
|
||
"""
|
||
# 准备特征和必要列
|
||
X_test = test_data.select(feature_cols)
|
||
key_cols = ["trade_date", "ts_code"]
|
||
key_data = test_data.select(key_cols)
|
||
|
||
print(f"[Predict] Test samples: {len(X_test)}, top_n: {top_n}")
|
||
|
||
# 处理数据（使用训练阶段的参数）
|
||
X_test_processed = pipeline.transform(X_test, stage=PipelineStage.TEST)
|
||
print(f"[Predict] Data processed, shape: {X_test_processed.shape}")
|
||
|
||
# 预测概率
|
||
probs = model.predict_proba(X_test_processed)
|
||
print(f"[Predict] Predictions generated, probability shape: {probs.shape}")
|
||
|
||
# 使用 key_data 添加预测结果，保持行数一致
|
||
result = key_data.with_columns(
|
||
pl.Series(
|
||
name="pred_prob",
|
||
values=probs[:, 1]
|
||
if len(probs.shape) > 1 and probs.shape[1] > 1
|
||
else probs.flatten(),
|
||
),
|
||
)
|
||
|
||
# 每日选出top N
|
||
top_stocks = []
|
||
for date in result["trade_date"].unique().sort():
|
||
day_data = result.filter(pl.col("trade_date") == date)
|
||
|
||
# 按概率降序排序，选出top N
|
||
day_top = day_data.sort("pred_prob", descending=True).head(top_n)
|
||
|
||
top_stocks.append(
|
||
day_top.select(["trade_date", "pred_prob", "ts_code"]).rename(
|
||
{"pred_prob": "score"}
|
||
)
|
||
)
|
||
|
||
return pl.concat(top_stocks)
|
||
|
||
|
||
def save_top_stocks(top_stocks: pl.DataFrame, output_path: str) -> None:
|
||
"""保存选股结果到TSV文件
|
||
|
||
Args:
|
||
top_stocks: 选股结果
|
||
output_path: 输出文件路径
|
||
"""
|
||
# 转换为pandas并保存为TSV
|
||
df = top_stocks.to_pandas()
|
||
df.to_csv(output_path, sep="\t", index=False)
|
||
print(f"[Training] Top stocks saved to: {output_path}")
|
||
|
||
|
||
def run_training(
|
||
factors: Optional[List[BaseFactor]] = None,
|
||
data_dir: str = "data",
|
||
output_path: str = "output/top_stocks.tsv",
|
||
train_start: str = "20180101",
|
||
train_end: str = "20230101",
|
||
val_start: str = "20230101",
|
||
val_end: str = "20230601",
|
||
test_start: str = "20230601",
|
||
test_end: str = "20240101",
|
||
top_n: int = 5,
|
||
) -> pl.DataFrame:
|
||
"""运行完整训练流程
|
||
|
||
Args:
|
||
factors: 因子实例列表，默认为 None（使用 MA5, MA10, ReturnRank5）
|
||
data_dir: 数据目录
|
||
output_path: 输出文件路径
|
||
train_start: 训练集开始日期
|
||
train_end: 训练集结束日期
|
||
val_start: 验证集开始日期
|
||
val_end: 验证集结束日期
|
||
test_start: 测试集开始日期
|
||
test_end: 测试集结束日期
|
||
top_n: 每日选股数量
|
||
|
||
Returns:
|
||
选股结果DataFrame
|
||
"""
|
||
print(f"[Training] Starting training pipeline...")
|
||
print(f"[Training] Train period: {train_start} -> {train_end}")
|
||
print(f"[Training] Val period: {val_start} -> {val_end}")
|
||
print(f"[Training] Test period: {test_start} -> {test_end}")
|
||
|
||
# 1. 准备数据
|
||
print("[Training] Preparing data...")
|
||
train_data, val_data, test_data, factor_config = prepare_data(
|
||
factors=factors,
|
||
data_dir=data_dir,
|
||
train_start=train_start,
|
||
train_end=train_end,
|
||
val_start=val_start,
|
||
val_end=val_end,
|
||
test_start=test_start,
|
||
test_end=test_end,
|
||
)
|
||
print(f"[Training] Train samples: {len(train_data)}")
|
||
print(f"[Training] Val samples: {len(val_data)}")
|
||
print(f"[Training] Test samples: {len(test_data)}")
|
||
|
||
# 2. 获取特征列名
|
||
feature_cols = factor_config.get_feature_names()
|
||
label_col = "label"
|
||
print(f"[Training] Feature columns: {feature_cols}")
|
||
|
||
# 3. 训练模型
|
||
print("[Training] Training model...")
|
||
model, pipeline = train_model(
|
||
train_data=train_data,
|
||
val_data=val_data,
|
||
feature_cols=feature_cols,
|
||
label_col=label_col,
|
||
)
|
||
|
||
# 4. 测试集预测
|
||
print("[Training] Predicting on test set...")
|
||
top_stocks = predict_top_stocks(
|
||
model=model,
|
||
pipeline=pipeline,
|
||
test_data=test_data,
|
||
feature_cols=feature_cols,
|
||
top_n=top_n,
|
||
)
|
||
|
||
# 5. 保存结果
|
||
print(f"[Training] Saving results to {output_path}...")
|
||
Path(output_path).parent.mkdir(parents=True, exist_ok=True)
|
||
save_top_stocks(top_stocks, output_path)
|
||
|
||
print("[Training] Training completed!")
|
||
|
||
return top_stocks
|