MLOps Project: Part 1: Machine Learning Experiment Tracking Using MLflow

This series follows a small sentiment-analysis service from training to deployment and monitoring. The model receives an e-commerce clothing review and predicts whether the reviewer recommends the product. The Women’s Clothing E-Commerce Reviews dataset is used as a teaching dataset; results in this post are examples from this workflow, not a claim about the dataset in general.

This first post focuses on experiment tracking with MLflow. Later posts cover deployment and monitoring. The MLOps Zoomcamp playlist provides the course context.

What is experiment tracking?

Experiment tracking records enough information to compare a run and reproduce it later:

• the data and split used for training;
• source code and environment;
• model family and hyperparameters;
• evaluation metrics;
• saved models, tokenizers, and other artifacts.

That record makes it easier to identify a promising run, explain why it was selected, and recreate it when the data or code changes. It does not make an experiment reproducible by itself: the dataset version, random seed, dependency versions, and preprocessing steps must also be recorded.

MLflow Tracking provides an API and UI for runs. MLflow also has model packaging, model-serving, and model-registry components. They solve different parts of the lifecycle.

Start a local tracking server

Install MLflow and the packages used by the examples:

python -m pip install mlflow==3.5.0 pandas==2.2.3 scikit-learn==1.5.2 nltk==3.9.1 tensorflow==2.18.0
python -m nltk.downloader stopwords

Start a local server backed by SQLite:

mlflow server --host 127.0.0.1 --port 5000 \
  --backend-store-uri sqlite:///mlflow.db \
  --artifacts-destination ./mlartifacts

Open http://127.0.0.1:5000 in a browser. The MLflow quickstart documents the current installation and tracking commands. For a shared setup on Google Cloud, use the authenticated MLflow on GCP walkthrough.

Split before fitting preprocessing

The vectorizer and tokenizer are part of the model pipeline. They must be fitted on training text only. If they see the validation or test text first, their vocabulary and term statistics leak information from the evaluation set.

The helper below creates disjoint training, validation, and test sets. The validation set is available for choosing hyperparameters or early stopping. The test set is read once for the final comparison.

from __future__ import annotations

import re
from pathlib import Path

import pandas as pd
from nltk.corpus import stopwords
from nltk.stem import PorterStemmer
from sklearn.model_selection import train_test_split


STOP_WORDS = set(stopwords.words("english"))
STEMMER = PorterStemmer()


def clean_review(value: str) -> str:
    words = re.sub(r"[^a-zA-Z]", " ", value).lower().split()
    return " ".join(
        STEMMER.stem(word)
        for word in words
        if word not in STOP_WORDS
    )


def load_splits(path: str | Path):
    data = pd.read_csv(path, index_col=0)
    data = data.dropna(subset=["Review Text"])
    texts = data["Review Text"].map(clean_review)
    labels = data["Recommended IND"].astype("int32")

    train_text, holdout_text, y_train, y_holdout = train_test_split(
        texts,
        labels,
        test_size=0.30,
        random_state=42,
        stratify=labels,
    )
    val_text, test_text, y_val, y_test = train_test_split(
        holdout_text,
        y_holdout,
        test_size=0.50,
        random_state=42,
        stratify=y_holdout,
    )
    return (
        train_text.reset_index(drop=True),
        val_text.reset_index(drop=True),
        test_text.reset_index(drop=True),
        y_train.reset_index(drop=True),
        y_val.reset_index(drop=True),
        y_test.reset_index(drop=True),
    )

The regular expression uses the complete alphabetic range. The split is deterministic because the seed is recorded. Replace the seed or split strategy when the data requires a time or group boundary.

Bag of Words

A Bag of Words vectorizer represents a review by the counts of its terms. Fit it on the training reviews, then transform validation and test reviews with the same vocabulary:

import mlflow
import mlflow.sklearn
from sklearn.feature_extraction.text import CountVectorizer
from sklearn.metrics import accuracy_score
from sklearn.naive_bayes import BernoulliNB
from sklearn.pipeline import Pipeline


mlflow.set_tracking_uri("http://127.0.0.1:5000")
mlflow.set_experiment("customer-sentiment-analysis")

train_text, val_text, test_text, y_train, y_val, y_test = load_splits(
    "Womens Clothing E-Commerce Reviews.csv"
)

alpha = 1.0
with mlflow.start_run(run_name="bow-bernoulli") as run:
    model = Pipeline(
        [
            ("vectorizer", CountVectorizer(max_features=3000, ngram_range=(1, 1))),
            ("classifier", BernoulliNB(alpha=alpha)),
        ]
    )
    model.fit(train_text, y_train)

    val_accuracy = accuracy_score(y_val, model.predict(val_text))
    test_accuracy = accuracy_score(y_test, model.predict(test_text))

    mlflow.log_param("algorithm", "BernoulliNB")
    mlflow.log_param("alpha", alpha)
    mlflow.log_param("max_features", 3000)
    mlflow.log_param("random_state", 42)
    mlflow.log_metric("validation_accuracy", val_accuracy)
    mlflow.log_metric("test_accuracy", test_accuracy)
    mlflow.sklearn.log_model(model, artifact_path="model")

    run_id = run.info.run_id

print(f"run ID: {run_id}")
print(f"test accuracy: {test_accuracy:.3f}")

The logged Pipeline contains both the vectorizer and classifier. Loading the one model artifact therefore applies the same vocabulary and classifier at inference time.

TF-IDF

TF-IDF downweights terms that appear in many documents. The data split and logging contract stay the same:

import mlflow
import mlflow.sklearn
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics import accuracy_score
from sklearn.naive_bayes import MultinomialNB
from sklearn.pipeline import Pipeline


alpha = 1.0
with mlflow.start_run(run_name="tfidf-multinomial") as run:
    model = Pipeline(
        [
            ("vectorizer", TfidfVectorizer(max_features=3000, ngram_range=(1, 1))),
            ("classifier", MultinomialNB(alpha=alpha)),
        ]
    )
    model.fit(train_text, y_train)

    val_accuracy = accuracy_score(y_val, model.predict(val_text))
    test_accuracy = accuracy_score(y_test, model.predict(test_text))

    mlflow.log_param("algorithm", "MultinomialNB")
    mlflow.log_param("alpha", alpha)
    mlflow.log_param("max_features", 3000)
    mlflow.log_param("random_state", 42)
    mlflow.log_metric("validation_accuracy", val_accuracy)
    mlflow.log_metric("test_accuracy", test_accuracy)
    mlflow.sklearn.log_model(model, artifact_path="model")

    run_id = run.info.run_id

print(f"run ID: {run_id}")
print(f"test accuracy: {test_accuracy:.3f}")

A small neural model

A neural model uses a tokenizer and an embedding. The tokenizer is fitted on the training reviews only, and the validation set is used for early stopping. The test set remains untouched until evaluation.

import pickle

import mlflow
import mlflow.tensorflow
import tensorflow as tf
from tensorflow.keras.preprocessing.sequence import pad_sequences
from tensorflow.keras.preprocessing.text import Tokenizer


tokenizer = Tokenizer(num_words=3000, oov_token="<unk>")
tokenizer.fit_on_texts(train_text)

X_train = pad_sequences(
    tokenizer.texts_to_sequences(train_text),
    padding="post",
)
X_val = pad_sequences(
    tokenizer.texts_to_sequences(val_text),
    maxlen=X_train.shape[1],
    padding="post",
)
X_test = pad_sequences(
    tokenizer.texts_to_sequences(test_text),
    maxlen=X_train.shape[1],
    padding="post",
)

with mlflow.start_run(run_name="embedding-average") as run:
    model = tf.keras.Sequential(
        [
            tf.keras.layers.Embedding(3000, 32, mask_zero=True),
            tf.keras.layers.GlobalAveragePooling1D(),
            tf.keras.layers.Dense(6, activation="relu"),
            tf.keras.layers.Dense(1, activation="sigmoid"),
        ]
    )
    model.compile(
        optimizer="adam",
        loss="binary_crossentropy",
        metrics=["accuracy"],
    )

    history = model.fit(
        X_train,
        y_train,
        validation_data=(X_val, y_val),
        epochs=50,
        batch_size=32,
        callbacks=[
            tf.keras.callbacks.EarlyStopping(
                monitor="val_loss",
                patience=2,
                restore_best_weights=True,
            )
        ],
        verbose=1,
    )

    test_loss, test_accuracy = model.evaluate(
        X_test,
        y_test,
        batch_size=128,
        verbose=0,
    )
    mlflow.log_param("algorithm", "embedding-average")
    mlflow.log_param("embedding_dim", 32)
    mlflow.log_param("epochs", len(history.history["loss"]))
    mlflow.log_metric("test_loss", test_loss)
    mlflow.log_metric("test_accuracy", test_accuracy)
    mlflow.tensorflow.log_model(model, artifact_path="model")

    with open("tokenizer.pickle", "wb") as handle:
        pickle.dump(tokenizer, handle, protocol=pickle.HIGHEST_PROTOCOL)
    mlflow.log_artifact("tokenizer.pickle", artifact_path="preprocessing")

    run_id = run.info.run_id

print(f"run ID: {run_id}")
print(f"test accuracy: {test_accuracy:.3f}")

Autologging can capture many framework parameters, but explicit metrics and artifact names make the later registry code easier to follow. The model artifact is always named model and the tokenizer is always under preprocessing/tokenizer.pickle.

Compare runs and manage models

Use the MLflow UI to compare the validation and test metrics, parameters, and artifacts for the three runs. Record these fields for each run:

Pick a model using the validation metric and the agreed selection rule. Report the test metric only after the selection is fixed. The selected result is a property of this dataset split and model configuration.

The deployment example in the next post uses the neural run because that run logs a tokenizer alongside the model. Choose that run only after reviewing the validation metrics. The Bag of Words and TF-IDF runs are self-contained scikit-learn pipelines and need a different serving contract.

The MLflow Model Registry workflow uses explicit model versions and aliases. The example assigns the selected version the champion alias:

import mlflow
from mlflow.tracking import MlflowClient


MODEL_NAME = "customer-sentiment-analysis"
client = MlflowClient(tracking_uri="http://127.0.0.1:5000")

# This walkthrough deploys the neural run, selected after reviewing validation.
# In a comparison, replace this with the run ID chosen by the agreed rule.
selected_run_id = run_id
model_version = mlflow.register_model(
    model_uri=f"runs:/{selected_run_id}/model",
    name=MODEL_NAME,
)
client.set_registered_model_alias(
    name=MODEL_NAME,
    alias="champion",
    version=model_version.version,
)

loaded_model = mlflow.pyfunc.load_model(
    f"models:/{MODEL_NAME}@champion"
)
print(loaded_model)

Registering a model creates a version and records the source run. Assigning an alias labels the version for consumers; it does not deploy the model or make a production decision by itself.

The tokenizer is retrieved from the same source run as the aliased model:

import pickle
from pathlib import Path

version = client.get_model_version_by_alias(MODEL_NAME, "champion")
tokenizer_dir = Path("downloaded_artifacts")
tokenizer_dir.mkdir(exist_ok=True)

tokenizer_file = client.download_artifacts(
    run_id=version.run_id,
    path="preprocessing/tokenizer.pickle",
    dst_path=str(tokenizer_dir),
)

with open(tokenizer_file, "rb") as handle:
    tokenizer = pickle.load(handle)

The registry lookup obtains the source run identifier for the aliased model. Keep the model and preprocessing artifact versioned together.

The next post uses the selected model in a service. The MLflow on GCP article shows how to move tracking metadata and artifacts to managed Google Cloud resources.