model.py 文件源码

def build_model(all_features, categorical_features, target, connection_string, filename):
    engine = create_engine(connection_string)

    query = """SELECT * FROM connection_features"""
    df = pd.read_sql_query(query, con=engine, index_col=['departurestop', 'departuredate', 'route'])
    df.index.levels[0].name = 'stationfrom'
    df.index.levels[1].name = 'date'
    df.index.levels[2].name = 'vehicle'
    df = df.reset_index()

    model = build_model_random_forest(df, all_features, categorical_features, target)
    joblib.dump(model, filename)