forked from bitosslab/linuxrustcommit
94 lines
4.1 KiB
Python
94 lines
4.1 KiB
Python
import os
|
|
import argparse
|
|
import logging
|
|
from datetime import datetime
|
|
|
|
from src.data_preprocessing import load_data, preprocess_data, create_time_series
|
|
from src.feature_engineering import engineer_features
|
|
from src.model import create_model
|
|
from src.train import train_model, evaluate_model
|
|
from src.predict import predict_churn, interpret_predictions
|
|
|
|
logging.basicConfig(filename='churn_prediction.log', level=logging.INFO)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
def setup_argparse():
|
|
parser = argparse.ArgumentParser(description='开源项目贡献者流失预测系统')
|
|
parser.add_argument('--data-dir', type=str, default='data', help='数据目录路径')
|
|
parser.add_argument('--output-dir', type=str, default='results', help='输出目录路径')
|
|
parser.add_argument('--epochs', type=int, default=100, help='训练轮数')
|
|
parser.add_argument('--batch-size', type=int, default=32, help='批次大小')
|
|
parser.add_argument('--learning-rate', type=float, default=0.001, help='学习率')
|
|
return parser.parse_args()
|
|
|
|
def main():
|
|
args = setup_argparse()
|
|
|
|
logger.info(f"开始运行流失预测系统 - {datetime.now()}")
|
|
|
|
try:
|
|
# 数据预处理
|
|
logger.info("开始数据预处理")
|
|
linux_df = load_data(os.path.join(args.data_dir, 'linux_commits.csv'))
|
|
rust_df = load_data(os.path.join(args.data_dir, 'rust_commits.csv'))
|
|
|
|
linux_df = preprocess_data(linux_df)
|
|
rust_df = preprocess_data(rust_df)
|
|
|
|
linux_X, linux_y = create_time_series(linux_df)
|
|
rust_X, rust_y = create_time_series(rust_df)
|
|
|
|
# 特征工程
|
|
logger.info("开始特征工程")
|
|
linux_X_engineered = engineer_features(linux_X)
|
|
rust_X_engineered = engineer_features(rust_X)
|
|
|
|
# 创建模型
|
|
logger.info("创建模型")
|
|
input_shape = (linux_X_engineered.shape[1], linux_X_engineered.shape[2])
|
|
model = create_model(input_shape)
|
|
|
|
# 训练模型
|
|
logger.info("开始训练模型")
|
|
linux_model, _ = train_model(model, linux_X_engineered, linux_y,
|
|
epochs=args.epochs, batch_size=args.batch_size)
|
|
rust_model, _ = train_model(model, rust_X_engineered, rust_y,
|
|
epochs=args.epochs, batch_size=args.batch_size)
|
|
|
|
# 评估模型
|
|
logger.info("评估模型")
|
|
linux_loss, linux_accuracy = evaluate_model(linux_model, linux_X_engineered, linux_y)
|
|
rust_loss, rust_accuracy = evaluate_model(rust_model, rust_X_engineered, rust_y)
|
|
|
|
logger.info(f"Linux模型 - 损失: {linux_loss:.4f}, 准确率: {linux_accuracy:.4f}")
|
|
logger.info(f"Rust模型 - 损失: {rust_loss:.4f}, 准确率: {rust_accuracy:.4f}")
|
|
|
|
# 预测
|
|
logger.info("开始预测")
|
|
linux_predictions = predict_churn(linux_model, linux_X_engineered)
|
|
rust_predictions = predict_churn(rust_model, rust_X_engineered)
|
|
|
|
linux_churn_prob, linux_churn_status = interpret_predictions(linux_predictions)
|
|
rust_churn_prob, rust_churn_status = interpret_predictions(rust_predictions)
|
|
|
|
# 保存结果
|
|
os.makedirs(args.output_dir, exist_ok=True)
|
|
with open(os.path.join(args.output_dir, 'linux_results.txt'), 'w') as f:
|
|
f.write(f"Linux项目贡献者流失预测:\n")
|
|
f.write(f"流失概率: {linux_churn_prob[:5]}\n")
|
|
f.write(f"流失状态: {linux_churn_status[:5]}\n")
|
|
|
|
with open(os.path.join(args.output_dir, 'rust_results.txt'), 'w') as f:
|
|
f.write(f"Rust项目贡献者流失预测:\n")
|
|
f.write(f"流失概率: {rust_churn_prob[:5]}\n")
|
|
f.write(f"流失状态: {rust_churn_status[:5]}\n")
|
|
|
|
logger.info(f"预测结果已保存到 {args.output_dir}")
|
|
logger.info(f"流失预测系统运行完成 - {datetime.now()}")
|
|
|
|
except Exception as e:
|
|
logger.error(f"运行过程中发生错误: {str(e)}")
|
|
raise
|
|
|
|
if __name__ == '__main__':
|
|
main() |