Ash
48 lines
1.6 KiB
Python
48 lines
1.6 KiB
Python
|
|
import pandas as pd
|
|
from sklearn.linear_model import LogisticRegression
|
|
import skl2onnx
|
|
from skl2onnx.common.data_types import FloatTensorType
|
|
|
|
def main():
|
|
"""
|
|
This script performs the following steps:
|
|
1. Loads the historical price data from a CSV file.
|
|
2. Performs feature engineering to create a 'feature_price_change' feature.
|
|
3. Performs target engineering to create a 'y_target_direction' target variable.
|
|
4. Creates and trains a logistic regression model.
|
|
5. Exports the trained model to an ONNX file named 'benchmark_logistic_model.onnx'.
|
|
"""
|
|
# Load the data
|
|
raw_price_df = pd.read_csv("raw_price_data.csv")
|
|
|
|
# Feature Engineering
|
|
raw_price_df["feature_price_change"] = raw_price_df["Close"].diff()
|
|
|
|
# Target Engineering
|
|
raw_price_df["y_target_direction"] = (raw_price_df["Close"].shift(-1) > raw_price_df["Close"]).astype(int)
|
|
|
|
# Drop rows with NaN values
|
|
raw_price_df.dropna(inplace=True)
|
|
|
|
# Separate features and target
|
|
X = raw_price_df[['feature_price_change']]
|
|
y = raw_price_df['y_target_direction']
|
|
|
|
# Create and train the logistic regression model
|
|
log_reg_model = LogisticRegression()
|
|
log_reg_model.fit(X, y)
|
|
|
|
# Convert the model to ONNX format
|
|
initial_type = [('float_input', FloatTensorType([None, 1]))]
|
|
onnx_model = skl2onnx.convert_sklearn(log_reg_model, initial_types=initial_type)
|
|
|
|
# Save the ONNX model
|
|
with open("benchmark_logistic_model.onnx", "wb") as f:
|
|
f.write(onnx_model.SerializeToString())
|
|
|
|
print("Model successfully exported to benchmark_logistic_model.onnx")
|
|
|
|
if __name__ == "__main__":
|
|
main()
|