Notebooks to Production

Most models are born in notebooks. They are wonderful for exploring, and terrible for production. Cells run in any order, paths point at someone's downloads folder, and nothing is tested. The step from notebook to production code is where many projects stall. Let's look at a practical way to take it, one function at a time.

• Notebooks are for exploring
• Hidden order and paths
• Refactor into functions

The notebook version

Here is a typical notebook cell. It reads a file from a personal downloads folder, with a name that suggests a long history. It filters, and fills missing values, in place. It trains a model, and prints the score, on the same data it trained on, which tells us nothing. Everything works, on one machine, today. But no one else can run it, nobody can test it, and next month even its author will not remember which cells to run first.

# notebook cell, run top to bottom, in whatever order
import pandas as pd
df = pd.read_csv("/Users/me/Downloads/churn_final_v3.csv")
df = df[df.charges > 0]
df["calls"] = df["calls"].fillna(0)
X = df[["tenure", "charges", "calls"]]
y = df["churn"]
from sklearn.linear_model import LogisticRegression
model = LogisticRegression().fit(X, y)
print(model.score(X, y))

The production version

And here is the same logic, as production code. A constant lists the features, in one place. A clean function takes a table, and returns a table, with no hidden state. A train function takes data, and a seed, and returns a fitted model. Both have type hints, and clear inputs, and outputs. And a small test checks the most important rule, that bad rows are removed. The behaviour is the same. But now it can be imported, tested, versioned, and run in a pipeline, by anyone.

# src/train.py, a function with inputs, outputs, and a test
import pandas as pd
from sklearn.linear_model import LogisticRegression

FEATURES = ["tenure", "charges", "calls"]


def clean(df: pd.DataFrame) -> pd.DataFrame:
    df = df[df["charges"] > 0].copy()
    df["calls"] = df["calls"].fillna(0)
    return df


def train(df: pd.DataFrame, seed: int = 0):
    df = clean(df)
    model = LogisticRegression(max_iter=500, random_state=seed)
    return model.fit(df[FEATURES], df["churn"])


def test_clean_removes_bad_rows():
    raw = pd.DataFrame({"tenure": [1, 2], "charges": [-5.0, 40.0], "calls": [None, 1]})
    assert len(clean(raw)) == 1

A refactoring routine

Follow a simple routine. First, make the notebook run top to bottom, in a fresh session. Second, remove hard coded paths, and move them to settings. Third, turn repeated cells into functions, with clear inputs and outputs. Fourth, move the functions into a Python module, in Git. Fifth, add tests, for the cleaning, and for the model. And sixth, call the module from the notebook, which now only explores and reports. Notebooks stay useful, as a front end, and the logic lives in code.

Summary