API

You gotta have an API

API for Machine Learning #

Fast api is pretty awesome for fast ML api. Other options could be mlflow serve.

from fastapi import FastAPI
from pydantic import BaseModel
import joblib

# Initialize the app
app = FastAPI()

# Load the model
model = joblib.load("model.pkl")

# Define a request body using Pydantic BaseModel
class Input(BaseModel):
    feature_1: float
    feature_2: float
    feature_3: float

# Define an endpoint
@app.post("/predict")
async def predict(input: Input):
    # Convert input to a dictionary and reshape
    input_dict = input.dict()
    features = [[input_dict["feature_1"], input_dict["feature_2"], input_dict["feature_3"]]]

    # Make a prediction using the model
    prediction = model.predict(features)[0]

    # Return the prediction as a JSON response
    return {"prediction": prediction}

For your model loader - you can easily load a model from a model registry like mlflow or s3 alike. Class Input could be replaced with a datamodel from pydantic for fast feature checking before you use bad data to predict.