batman

car_price_guesser/train_model.py

@@ -0,0 +1,197 @@
+import pandas as pd
+import numpy as np
+from sklearn.model_selection import train_test_split
+from sklearn.linear_model import LinearRegression
+from sklearn.ensemble import RandomForestRegressor
+from sklearn.metrics import mean_squared_error, r2_score
+from sklearn.preprocessing import OneHotEncoder
+from sklearn.compose import ColumnTransformer
+from sklearn.pipeline import Pipeline
+from sklearn.impute import SimpleImputer
+import joblib
+
+
+def load_and_prepare_data():
+    """
+    Load the dataset and prepare it for modeling.
+    
+    Returns:
+        tuple: X (features), y (target)
+    """
+    print("Loading data...")
+    df = pd.read_csv("New_York_cars.csv")
+    
+    # Remove rows where target 'money' is missing
+    df = df.dropna(subset=["money"])
+    
+    # Define features and target
+    features = [
+        "Year",
+        "Mileage",
+        "brand",
+        "new&used",
+        "Exterior color",
+        "Interior color",
+        "Drivetrain",
+        "MPG",
+        "Fuel type",
+        "Transmission",
+        "Engine",
+        "Convenience",
+        "Entertainment",
+        "Exterior",
+        "Safety",
+        "Seating",
+        "Accidents or damage",
+        "Clean title",
+        "1-owner vehicle",
+        "Personal use only",
+        "Model",
+    ]
+    target = "money"
+    
+    X = df[features]
+    y = df[target]
+    
+    return X, y
+
+
+def create_preprocessor():
+    """
+    Create a preprocessing pipeline for handling both numeric and categorical features.
+    
+    Returns:
+        ColumnTransformer: Preprocessing pipeline
+    """
+    # Define numeric and categorical features
+    numeric_features = ["Year", "Mileage"]
+    categorical_features = [
+        "brand",
+        "MPG",
+        "Fuel type",
+        "Transmission",
+        "Engine",
+        "Safety",
+        "Accidents or damage",
+        "Clean title",
+        "1-owner vehicle",
+        "Personal use only",
+        "Model",
+    ]
+    
+    # Numeric preprocessing pipeline
+    numeric_transformer = Pipeline(
+        steps=[("imputer", SimpleImputer(strategy="median"))]
+    )
+    
+    # Categorical preprocessing pipeline
+    categorical_transformer = Pipeline(
+        steps=[
+            ("imputer", SimpleImputer(strategy="constant", fill_value="Unknown")),
+            ("onehot", OneHotEncoder(handle_unknown="ignore", sparse_output=False)),
+        ]
+    )
+    
+    # Combine preprocessors
+    preprocessor = ColumnTransformer(
+        transformers=[
+            ("num", numeric_transformer, numeric_features),
+            ("cat", categorical_transformer, categorical_features),
+        ]
+    )
+    
+    return preprocessor
+
+
+def train_model(X_train, y_train, preprocessor):
+    """
+    Train a linear regression model with preprocessing.
+    
+    Args:
+        X_train: Training features
+        y_train: Training target
+        preprocessor: Preprocessing pipeline
+        
+    Returns:
+        Pipeline: Trained model pipeline
+    """
+    print("Training Linear Regression...")
+    
+    # Create model pipeline
+    lr_pipeline = Pipeline(
+        steps=[("preprocessor", preprocessor), ("regressor", LinearRegression())]
+    )
+    
+    # Fit the model
+    lr_pipeline.fit(X_train, y_train)
+    
+    return lr_pipeline
+
+
+def evaluate_model(model, X_test, y_test):
+    """
+    Evaluate the trained model using MSE and R2 metrics.
+    
+    Args:
+        model: Trained model pipeline
+        X_test: Test features
+        y_test: Test target
+        
+    Returns:
+        tuple: (mse, r2)
+    """
+    # Make predictions
+    y_pred = model.predict(X_test)
+    
+    # Calculate metrics
+    mse = mean_squared_error(y_test, y_pred)
+    rmse = np.sqrt(mse)
+    r2 = r2_score(y_test, y_pred)
+    
+    print(f"Linear Regression MSE: {mse}")
+    print(f"Linear Regression RMSE: {rmse:.4f}")
+    print(f"Linear Regression R2: {r2}")
+    
+    return mse, r2
+
+
+def save_model(model, filename="car_price_model.joblib"):
+    """
+    Save the trained model to disk.
+    
+    Args:
+        model: Trained model pipeline
+        filename: Output filename
+    """
+    joblib.dump(model, filename)
+    print(f"Model saved to {filename}")
+
+
+def train():
+    """
+    Main training function that orchestrates the entire process.
+    """
+    # Load and prepare data
+    X, y = load_and_prepare_data()
+    
+    # Split data into train and test sets
+    print("Splitting data...")
+    X_train, X_test, y_train, y_test = train_test_split(
+        X, y, test_size=0.2, random_state=42
+    )
+    
+    # Create preprocessing pipeline
+    preprocessor = create_preprocessor()
+    
+    # Train the model
+    model = train_model(X_train, y_train, preprocessor)
+    
+    # Evaluate the model
+    mse, r2 = evaluate_model(model, X_test, y_test)
+    
+    # Save the model
+    save_model(model)
+
+
+if __name__ == "__main__":
+    train()