generated from Python-2026Spring/assignment-05-final-project-template
31 lines
937 B
Python
31 lines
937 B
Python
# save_model.py
|
||
import joblib
|
||
from src.data import load_data, split_data
|
||
from src.model import build_preprocessor
|
||
from sklearn.ensemble import RandomForestClassifier
|
||
from sklearn.pipeline import Pipeline
|
||
|
||
if __name__ == "__main__":
|
||
# 加载数据
|
||
df = load_data()
|
||
X_train, X_test, y_train, y_test = split_data(df)
|
||
preprocessor = build_preprocessor(X_train)
|
||
|
||
# 训练调优后的随机森林模型
|
||
model = Pipeline(steps=[
|
||
('preprocessor', preprocessor),
|
||
('classifier', RandomForestClassifier(
|
||
n_estimators=200,
|
||
max_depth=15,
|
||
min_samples_split=8,
|
||
min_samples_leaf=4,
|
||
class_weight='balanced_subsample',
|
||
random_state=42,
|
||
n_jobs=-1
|
||
))
|
||
])
|
||
model.fit(X_train, y_train)
|
||
|
||
# 保存模型
|
||
joblib.dump(model, 'telco_churn_model.pkl')
|
||
print("✅ 模型已保存为：telco_churn_model.pkl") |