78 lines
2.3 KiB
Python
78 lines
2.3 KiB
Python
import logging
|
|
|
|
import mlflow
|
|
import pandas as pd
|
|
import typer
|
|
from sklearn import linear_model
|
|
from sklearn import metrics
|
|
|
|
from .. import constants
|
|
|
|
app = typer.Typer()
|
|
|
|
logging.basicConfig(level=logging.INFO)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@app.command()
|
|
def main(
|
|
strategy: constants.ModellingStrategy,
|
|
):
|
|
training_file_path = constants.DATASET_DIR / "train.parquet"
|
|
validation_file_path = constants.DATASET_DIR / "validation.parquet"
|
|
|
|
features = constants.MODELLING_FEATURES[strategy]
|
|
|
|
train_df = pd.read_parquet(training_file_path)
|
|
validation_df = pd.read_parquet(validation_file_path)
|
|
|
|
train_df = train_df.dropna(
|
|
subset=features + [constants.TARGET]
|
|
)
|
|
validation_df = validation_df.dropna(
|
|
subset=features + [constants.TARGET]
|
|
)
|
|
|
|
X_train = train_df[features]
|
|
y_train = train_df[constants.TARGET]
|
|
|
|
X_validation = validation_df[features]
|
|
y_validation = validation_df[constants.TARGET]
|
|
|
|
logger.info(f"Training model with strategy '{strategy}' and features {features}")
|
|
|
|
with mlflow.start_run(run_name=f"modelling_{strategy.value}"):
|
|
mlflow.log_param("model_type", "linear")
|
|
mlflow.log_param("strategy", strategy.value)
|
|
mlflow.log_param("split_strategy", constants.CHOSEN_SPLIT_STRATEGY.value)
|
|
|
|
mlflow.log_param("features", ",".join(features))
|
|
model = linear_model.LinearRegression()
|
|
model.fit(X_train, y_train)
|
|
|
|
train_predictions = model.predict(X_train)
|
|
validation_predictions = model.predict(X_validation)
|
|
|
|
train_rmse = metrics.root_mean_squared_error(y_train, train_predictions)
|
|
validation_rmse = metrics.root_mean_squared_error(y_validation, validation_predictions)
|
|
|
|
train_mae = metrics.mean_absolute_error(y_train, train_predictions)
|
|
validation_mae = metrics.mean_absolute_error(y_validation, validation_predictions)
|
|
|
|
mlflow.log_metric("train_rmse", train_rmse)
|
|
mlflow.log_metric("validation_rmse", validation_rmse)
|
|
|
|
mlflow.log_metric("train_mae", train_mae)
|
|
mlflow.log_metric("validation_mae", validation_mae)
|
|
|
|
for feature_name, coefficient in zip(
|
|
features,
|
|
model.coef_,
|
|
strict=True,
|
|
):
|
|
mlflow.log_metric(f"coef_{feature_name}", float(coefficient))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
app()
|