File size: 3,931 Bytes
9e03a51 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 | from __future__ import annotations
from dataclasses import dataclass
from math import sqrt
from pathlib import Path
from typing import Any
import joblib
import pandas as pd
from sklearn.compose import ColumnTransformer
from sklearn.ensemble import RandomForestRegressor
from sklearn.impute import SimpleImputer
from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score
from sklearn.model_selection import train_test_split
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import OneHotEncoder
from src.data import DEFAULT_FEATURES, TARGET_COLUMN, select_model_frame
NUMERIC_FEATURES = [
"OverallQual",
"GrLivArea",
"GarageCars",
"TotalBsmtSF",
"FullBath",
"YearBuilt",
]
CATEGORICAL_FEATURES = ["Neighborhood", "HouseStyle"]
DEFAULT_MODEL_PATH = Path("models") / "house_price_model.joblib"
@dataclass(frozen=True)
class TrainedModel:
pipeline: Pipeline
feature_names: list[str]
metrics: dict[str, float]
class HousePriceModel:
def __init__(self, trained: TrainedModel):
self._trained = trained
@property
def metrics(self) -> dict[str, float]:
return self._trained.metrics
@property
def feature_names(self) -> list[str]:
return self._trained.feature_names
@classmethod
def load(cls, path: str | Path = DEFAULT_MODEL_PATH) -> "HousePriceModel":
payload = joblib.load(path)
return cls(
TrainedModel(
pipeline=payload["pipeline"],
feature_names=list(payload["feature_names"]),
metrics=dict(payload["metrics"]),
)
)
def predict(self, features: dict[str, Any]) -> float:
frame = pd.DataFrame([{name: features.get(name) for name in self.feature_names}])
prediction = self._trained.pipeline.predict(frame)[0]
return round(float(max(prediction, 0.0)), 2)
def build_pipeline() -> Pipeline:
numeric_pipeline = Pipeline(
steps=[
("imputer", SimpleImputer(strategy="median")),
]
)
categorical_pipeline = Pipeline(
steps=[
("imputer", SimpleImputer(strategy="most_frequent")),
("encoder", OneHotEncoder(handle_unknown="ignore", sparse_output=False)),
]
)
preprocessor = ColumnTransformer(
transformers=[
("num", numeric_pipeline, NUMERIC_FEATURES),
("cat", categorical_pipeline, CATEGORICAL_FEATURES),
]
)
regressor = RandomForestRegressor(n_estimators=250, random_state=42, min_samples_leaf=2)
return Pipeline(steps=[("preprocessor", preprocessor), ("regressor", regressor)])
def train_model(frame: pd.DataFrame, artifact_path: str | Path = DEFAULT_MODEL_PATH) -> TrainedModel:
model_frame = select_model_frame(frame).dropna(subset=[TARGET_COLUMN])
features = model_frame[DEFAULT_FEATURES]
target = model_frame[TARGET_COLUMN]
if len(model_frame) >= 10:
x_train, x_test, y_train, y_test = train_test_split(
features, target, test_size=0.2, random_state=42
)
else:
x_train, x_test, y_train, y_test = features, features, target, target
pipeline = build_pipeline()
pipeline.fit(x_train, y_train)
predictions = pipeline.predict(x_test)
metrics = {
"rmse": round(sqrt(mean_squared_error(y_test, predictions)), 2),
"mae": round(mean_absolute_error(y_test, predictions), 2),
"r2": round(r2_score(y_test, predictions), 4),
}
trained = TrainedModel(pipeline=pipeline, feature_names=list(DEFAULT_FEATURES), metrics=metrics)
artifact_path = Path(artifact_path)
artifact_path.parent.mkdir(parents=True, exist_ok=True)
joblib.dump(
{
"pipeline": trained.pipeline,
"feature_names": trained.feature_names,
"metrics": trained.metrics,
},
artifact_path,
)
return trained
|