btc_predictor / model_preparation.py
alexray's picture
Update from Jenkins
79cf85f
Raw
History Blame Contribute Delete
1.11 kB
from sklearn.svm import LinearSVR
from sklearn.metrics import mean_squared_error, r2_score
import pandas as pd
import pickle
def prepare_model(data_dir="data", model_name="linear_svr_model.pkl"):
# Load training data
X_train = pd.read_csv(f"{data_dir}/train_features.csv", index_col=0)
y_train = pd.read_csv(f"{data_dir}/train_target.csv", index_col=0)
y_train = y_train.values.ravel()
print(X_train.shape)
print(y_train.shape)
model = LinearSVR(random_state=42, max_iter=10000)
model.fit(X_train, y_train)
predictions = model.predict(X_train)
mse = mean_squared_error(y_train, predictions)
r2 = r2_score(y_train, predictions)
print("Training:")
print("Mean Squared Error:", mse)
print("R-squared:", r2)
with open(model_name, "wb") as model_file:
pickle.dump(model, model_file)
predictions_df = pd.DataFrame(predictions, index=X_train.index,
columns=["Prediction"])
predictions_df.to_csv(f"{data_dir}/train_prediction.csv", index=True)
if __name__ == '__main__':
prepare_model()