Model Training and Prediction Functions in ML Pipelines

Hello again, and welcome to lesson 2 of "Building Reusable Pipeline Functions"! In our previous lesson, we learned how to create modular, reusable functions for loading and preprocessing data — the critical first steps in our ML pipeline. Now that our data is properly prepared, we're ready to move on to the next essential component: model training and prediction functions.

Throughout this lesson, we'll build upon the foundation we established with our data processing functions and create equally robust functions for training machine learning models and generating predictions. Just as with data processing, following good design principles for these functions will help ensure our ML system is maintainable, flexible, and production-ready.

By the end of this lesson, you'll understand how to create model training functions that can adapt to different modeling approaches while maintaining a consistent interface — a key skill for building production ML systems that can evolve over time.

Before we dive into code, let's discuss why well-designed model training functions are critical in an ML pipeline:

1. Consistency: Standard interfaces for model training and prediction ensure that different models can be easily swapped without changing the surrounding code.
2. Experimentation: Properly structured model functions make it easier to experiment with different algorithms and hyperparameters.
3. Maintainability: Isolating model logic in dedicated functions makes the code easier to understand and update.
4. Reproducibility: Well-designed functions with proper parameter handling ensure training results can be reproduced reliably.

When building model functions for production, we need to consider both immediate and future requirements. While we might start with a single model type, our pipeline should be flexible enough to accommodate different approaches as our project evolves. This forward-thinking design is a hallmark of production-grade ML systems.

Let's start by designing the interface for our model training function. A good interface should be intuitive and flexible, accommodating different model types while maintaining a consistent structure:

def train_model(X_train, y_train, model_type="random_forest", **model_params):
    """
    Train a machine learning model on the preprocessed data.
    
    Args:
        X_train (array-like): Preprocessed training features
        y_train (array-like): Training target values
        model_type (str): Type of model to train ('random_forest' or 'linear')
        **model_params: Additional parameters to pass to the model constructor
        
    Returns:
        object: Trained model
    """
    print(f"Training a {model_type} model...")
    
    # Model selection and training logic will go here
    
    print("Model training completed!")
    return model

This function signature demonstrates several important design principles:

• Required parameters (X_train, y_train) for the training data;
• A default model type ("random_forest") that provides a sensible default choice;
• Flexible parameter passing using **model_params to accept any number of model-specific parameters;
• A clear docstring that describes what the function does, its parameters, and return value.

Now let's add the logic to select and train different model types:

import numpy as np
from sklearn.ensemble import RandomForestRegressor
from sklearn.linear_model import LinearRegression

def train_model(X_train, y_train, model_type="random_forest", **model_params):
    """
    Train a machine learning model on the preprocessed data.
    
    Args:
        X_train (array-like): Preprocessed training features
        y_train (array-like): Training target values
        model_type (str): Type of model to train ('random_forest' or 'linear')
        **model_params: Additional parameters to pass to the model constructor
        
    Returns:
        object: Trained model
    """
    print(f"Training a {model_type} model...")
    
    if model_type == "random_forest":
        # Default parameters if not specified
        model_params.setdefault('n_estimators', 10)  # Using a sensible default
        model_params.setdefault('random_state', 42)  # For reproducibility
            
        model = RandomForestRegressor(**model_params)
    
    elif model_type == "linear":
        model = LinearRegression(**model_params)
    
    else:
        raise ValueError(f"Unsupported model type: {model_type}")
    
    # Train the model
    model.fit(X_train, y_train)
    
    print("Model training completed!")
    return model

This code defines a function to train different types of machine learning models:

• It accepts training data and a model type as inputs.
• It selects the model based on the model_type parameter.
• For a "random_forest" model, it uses setdefault to apply default values for n_estimators and random_state only when they aren't already provided in model_params.
• It raises an error if an unsupported model type is specified.
• It trains the selected model using the provided data with model.fit(X_train, y_train) and returns the trained model.

A good ML pipeline needs functions not just for training models but also for making predictions with those models. Let's implement a prediction function that works with any of our trained models:

def predict_with_model(model, X):
    """
    Make predictions using a trained model.
    
    Args:
        model (object): Trained model
        X (array-like): Preprocessed features
        
    Returns:
        array: Predictions
    """
    return model.predict(X)

This function is intentionally simple, but that simplicity is a strength. By creating a dedicated function for prediction:

1. We create a consistent interface for generating predictions, regardless of the underlying model type.
2. We establish a single point of modification if we need to change how predictions are generated.
3. We follow the single responsibility principle by separating training logic from prediction logic.

Though simple now, this function could evolve to include additional functionality like input validation, error handling, prediction post-processing, logging and monitoring. By isolating prediction in its own function, we make it easier to add these capabilities in the future without disrupting the rest of our pipeline.

Finally, let's see how our model training and prediction functions integrate into a complete pipeline:

def main():
    """Main function to demonstrate the model training pipeline."""
    # Step 1: Load the data
    print("Loading diamonds dataset...")
    data_path = "diamonds.csv"
    diamonds_df = load_diamonds_data(data_path)
    print(f"Dataset loaded with shape: {diamonds_df.shape}")
    
    # Step 2: Preprocess the dataset
    print("\nPreprocessing the dataset...")
    X_train, X_test, y_train, _, _ = preprocess_diamonds_data(diamonds_df)
    print(f"Preprocessing complete.")
    print(f"Training set size: {X_train.shape[0]}, Test set size: {X_test.shape[0]}")
    
    # Step 3: Train a Random Forest model with custom parameters
    print("\n--- Training Random Forest Model ---")
    rf_model = train_model(
        X_train, 
        y_train, 
        model_type="random_forest",
        n_estimators=10,
        max_depth=5,
        random_state=42
    )
    
    # Step 4: Generate predictions with the trained model
    y_pred_rf = predict_with_model(rf_model, X_test)
    print(f"Generated {len(y_pred_rf)} predictions")

This main function demonstrates the complete workflow, from data loading to prediction:

1. We load the dataset using our function from Lesson 1;
2. We preprocess the data using our preprocessing function, also from Lesson 1;
3. We train a Random Forest model with specific hyperparameters using our new train_model function;
4. We generate predictions on the test set using our predict_with_model function.

Notice how each step builds on the previous one, creating a clear, linear workflow. This organized approach makes the code easy to understand and modify. It also illustrates the real benefit of our modular design: we can easily swap in different models or change hyperparameters without disrupting the overall pipeline structure.

In this second lesson, we've expanded our ML pipeline by creating robust, flexible functions for model training and prediction. These functions follow the same design principles we established in our data processing module: modularity, flexibility, and clear interfaces. We've designed a system that can accommodate multiple model types, handle diverse parameters, and provide consistent prediction capabilities — all while maintaining a clean, readable codebase that follows software engineering best practices.

As you move into the practice exercises, you'll have the opportunity to experiment with these functions firsthand. This hands-on experience will reinforce the concepts we've covered and help you develop the skills needed to build production-quality ML pipelines in your own projects. Happy coding!