dgtw-mlflow-utils 1.0.4

MLflow utilities for model uploading, registry management, evaluation, and naming conventions.

dgtw-mlflow-utils

A collection of MLflow utilities for model uploading, registry management, evaluation, and standardized naming conventions — designed for the TAT Personalization project.

---

Installation

From PyPI (recommended)

# Install the core package
pip install dgtw-mlflow-utils

# Install with TensorFlow support
pip install "dgtw-mlflow-utils[tensorflow]"

# Install with all ML framework extras
pip install "dgtw-mlflow-utils[all]"

In a Jupyter notebook

%pip install dgtw-mlflow-utils

---

Quick Start

1. Naming Conventions

See full example: examples/01_naming_conventions.py

from dgtw_mlflow_utils import ModelNameBuilder, Architecture, ModelVariant, Stage, RunAction

builder = ModelNameBuilder()

# Experiment name → "TAT_Recommendation_Research"
builder.experiment_name(component="Recommendation", stage=Stage.RESEARCH)

# Registered model name → "TAT_TwoTower_Baseline"
builder.registered_model_name(
    architecture=Architecture.TWO_TOWER,
    variant=ModelVariant.BASELINE,
)

# Run name → "Train_MiniLM_20261025-1400"
builder.run_name(action=RunAction.TRAIN, use_case="MiniLM")

# Validate a name
ModelNameBuilder.validate_registered_name("TAT_TwoTower_Baseline")  # True

2. Upload a Scikit-learn Model

See full example: examples/02_upload_sklearn.py

from dgtw_mlflow_utils import ModelUploader, UploadConfig, Architecture, ModelVariant, Stage, RunAction

config = UploadConfig(
    architecture=Architecture.TWO_TOWER,
    variant=ModelVariant.BASELINE,
    component="Recommendation",
    stage=Stage.RESEARCH,
    use_case="MiniLM",
    action=RunAction.UPLOAD,
    tracking_uri="http://your-mlflow-server:5000",
    params={"embedding_dim": 256},
    tags={"developer": "ai-team"},
)

uploader = ModelUploader(config)
print(uploader.summary())

# Upload
result = uploader.upload_sklearn(sk_model=pipeline, input_example=X_test[:5])
print(result["model_uri"])

3. Upload a TensorFlow / Keras Model

See full example: examples/04_upload_tensorflow.py

result = uploader.upload_tensorflow(tf_model=keras_model)

4. Upload a PyTorch Model

result = uploader.upload_pytorch(pytorch_model=torch_model)

5. Upload Raw Artifact Files

See full example: examples/05_upload_artifacts.py

result = uploader.upload_artifacts(
    files=["vectors.npy", "config.yaml"],
    artifact_subdir="evaluation",
    extra_metrics={"recall_5": 0.82, "mrr": 0.65},
)

6. Model Registry Management

See full example: examples/03_registry_management.py

from dgtw_mlflow_utils import ModelRegistry

registry = ModelRegistry(tracking_uri="http://your-mlflow-server:5000")

# List all versions
versions = registry.list_versions("TAT_TwoTower_Baseline")

# Promote latest version to champion
registry.promote("TAT_TwoTower_Baseline", alias="champion")

# Load the champion model for inference
model = registry.load_champion("TAT_TwoTower_Baseline")
predictions = model.predict(input_data)

# Cleanup old versions (dry run first)
deleted = registry.cleanup_old_versions("TAT_TwoTower_Baseline", keep=5, dry_run=True)

# Print formatted summary
registry.print_summary("TAT_TwoTower_Baseline")

7. Training Evaluation Callback (TensorFlow / Keras)

See full example: examples/06_evaluator.py

from dgtw_mlflow_utils import TfMLflowEpochLogger

callback = TfMLflowEpochLogger(
    enable_mlflow=True,
    experiment_name="TAT_Recommendation_Research",
    run_name="Training_Run",
    # Extra kwargs are logged as MLflow params:
    log_params={
        "epochs": 50,
        "learning_rate": 0.001,
    }
)

# The callback handles the full lifecycle:
#   on_train_begin → starts MLflow run + logs params
#   on_epoch_end   → logs training metrics per epoch
#   on_test_end    → logs validation metrics
#   on_train_end   → ends MLflow run
model.fit(train_dataset, epochs=50, callbacks=[callback])

NumPy Equivalent (NumpyMLflowLogger)

For custom training loops or other frameworks (PyTorch, Scikit-learn):

from dgtw_mlflow_utils import NumpyMLflowLogger

logger = NumpyMLflowLogger(
    enable_mlflow=True,
    experiment_name="TAT_Recommendation_Research",
    run_name="Training_Run_Numpy",
)

logger.on_train_begin()
for epoch in range(50):
    # ... training ...
    logger.on_epoch_end(epoch, metrics={"loss": 0.5})
logger.on_train_end()

8. Standalone Evaluation

See full example: examples/06_evaluator.py

from dgtw_mlflow_utils import evaluate_model_tf

# Runs a batch-by-batch evaluation loop and computes Recall@K, MRR, NDCG, Loss
metrics = evaluate_model_tf(
    model=trained_model,
    dataset=test_dataset,
    k=10,
    enable_mlflow=True,
    experiment_name="TAT_Recommendation_Research",
    run_name="Final_Evaluation",
    log_params={"dataset_size": 1000},
)
# metrics → {"recall_10": 0.85, "mrr_10": 0.72, "ndcg_10": 0.78, "loss": 0.34}

NumPy Equivalent (evaluate_model_numpy)

For evaluating any model that can return predictions as a NumPy array:

from dgtw_mlflow_utils import evaluate_model_numpy

metrics = evaluate_model_numpy(
    predict_fn=lambda batch: your_model.predict(batch),
    dataset=test_dataset,
    k=10,
    enable_mlflow=True,
    experiment_name="TAT_Recommendation_Research",
)

9. Low-level Retrieval Metrics

See full example: examples/06_evaluator.py

import tensorflow as tf
from dgtw_mlflow_utils import compute_retrieval_metrics

# logits: (batch_size, batch_size) similarity matrix from in-batch negatives
logits = tf.random.normal((32, 32))

metrics = compute_retrieval_metrics(logits, k=5)
# metrics → {"recall_5": tensor, "mrr_5": tensor, "ndcg_5": tensor}

NumPy Equivalent (compute_retrieval_metrics_numpy)

import numpy as np
from dgtw_mlflow_utils import compute_retrieval_metrics_numpy

logits = np.random.normal(size=(32, 32))
metrics = compute_retrieval_metrics_numpy(logits, k=5)

---

Examples

The examples/ directory contains runnable scripts demonstrating each feature:

File	Description
01_naming_conventions.py	Generate experiment, model, and run names with validation
02_upload_sklearn.py	Train & upload a scikit-learn RandomForest on Iris
03_registry_management.py	List versions, promote champion, load model, cleanup
04_upload_tensorflow.py	Train a Keras model with epoch logging & upload
05_upload_artifacts.py	Upload raw files (configs, CSVs) without model flavour
06_evaluator.py	compute_retrieval_metrics, TfMLflowEpochLogger callback, evaluate_model_tf loop

# Run an example
python examples/01_naming_conventions.py

---

Modules

Module	Description
dgtw_mlflow_utils.naming	Enums & ModelNameBuilder for standardized MLflow naming
dgtw_mlflow_utils.uploader	ModelUploader & UploadConfig for uploading models
dgtw_mlflow_utils.registry	ModelRegistry for lifecycle management (aliases, cleanup, loading)
dgtw_mlflow_utils.evaluator	TfMLflowEpochLogger, NumpyMLflowLogger, & standalone evaluation loops
dgtw_mlflow_utils.metrics	compute_retrieval_metrics and compute_retrieval_metrics_numpy (Recall@K, MRR@K, NDCG@K)

Note: evaluator and metrics modules require TensorFlow. They are lazy-loaded so you can use the rest of the package without TensorFlow installed.

Optional Dependencies

Extra	Packages	Install command
tensorflow	tensorflow>=2.10	pip install dgtw-mlflow-utils[tensorflow]
sklearn	scikit-learn>=1.0	pip install dgtw-mlflow-utils[sklearn]
pytorch	torch>=1.12	pip install dgtw-mlflow-utils[pytorch]
all	All of the above	pip install dgtw-mlflow-utils[all]
dev	pytest, pytest-cov	pip install dgtw-mlflow-utils[dev]

Project Structure

dgtw-mlflow-utils/
├── pyproject.toml              # pip packaging config
├── README.md
├── src/                        # → installed as dgtw_mlflow_utils
│   ├── __init__.py             # package exports (lazy TF imports)
│   ├── naming.py               # ModelNameBuilder & enums
│   ├── uploader.py             # ModelUploader & UploadConfig
│   ├── registry.py             # ModelRegistry
│   ├── evaluator.py            # MLflowEpochLogger & evaluate_model_tf (and NumPy variants)
│   └── metrics.py              # compute_retrieval_metrics
└── examples/
    ├── 01_naming_conventions.py
    ├── 02_upload_sklearn.py
    ├── 03_registry_management.py
    ├── 04_upload_tensorflow.py
    ├── 05_upload_artifacts.py
    └── 06_evaluator.py

Repository

https://git.digithunworldwide.com/ai/mlflow-utils.git