# 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)

```bash
# 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

```python
%pip install dgtw-mlflow-utils
```

---

## Quick Start

### 1. Naming Conventions

> See full example: [`examples/01_naming_conventions.py`](examples/01_naming_conventions.py)

```python
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`](examples/02_upload_sklearn.py)

```python
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`](examples/04_upload_tensorflow.py)

```python
result = uploader.upload_tensorflow(tf_model=keras_model)
```

### 4. Upload a PyTorch Model

```python
result = uploader.upload_pytorch(pytorch_model=torch_model)
```

### 5. Upload Raw Artifact Files

> See full example: [`examples/05_upload_artifacts.py`](examples/05_upload_artifacts.py)

```python
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`](examples/03_registry_management.py)

```python
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`](examples/06_evaluator.py)

```python
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):

```python
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`](examples/06_evaluator.py)

```python
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:

```python
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`](examples/06_evaluator.py)

```python
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`)

```python
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/`](examples/) directory contains runnable scripts demonstrating each feature:

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

```bash
# 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

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