mltk - ML Test Kit

pytest for ML -- unified testing across the entire ML lifecycle.

Installation note (temporary): pip install mltk is not yet available — the PyPI name is pending a transfer claim. Until resolved, install directly from GitHub:

pip install git+https://github.com/Liorrr/mltk

All imports (import mltk), CLI commands (mltk scan), and functionality are identical.

Install

Homebrew (macOS/Linux):

brew tap Liorrr/mltk
brew install mltk

pip:

pip install mlspec  # PyPI name — 'mltk' pending transfer

Why mltk?

ML systems fail silently. A model can train on corrupt data, produce confident predictions from stale features, and pass every unit test while being completely wrong in production. Traditional testing does not catch these failures.

mltk gives you 244 assertion definitions (238 unique names; 229 listed by mltk list / discovery) covering the entire ML lifecycle -- data quality, model validation, drift detection, fairness testing, inference benchmarking, training bug detection, LLM evaluation, behavioral consistency, NER-based PII detection, red team security, multimodal evaluation, observability, and production monitoring. One toolkit, one pip install, native pytest integration. No more gluing together 5 different tools.

Quick Start

import pandas as pd
import pytest
from mltk.data import (
    assert_schema, assert_no_nulls, assert_no_pii,
)
from mltk.model import (
    assert_metric, assert_no_regression, assert_no_bias,
)
from mltk.inference import assert_latency, assert_throughput

@pytest.mark.ml_data
def test_training_data():
    df = pd.read_csv("data/training.csv")
    assert_schema(
        df,
        {"id": "int64", "text": "object", "label": "int64"},
    )
    assert_no_nulls(df, columns=["label"])
    assert_no_pii(df, columns=["text"])

@pytest.mark.ml_model
def test_model_quality(y_true, y_pred):
    assert_metric(
        y_true, y_pred, metric="f1", threshold=0.85,
    )
    assert_no_regression(
        y_true, y_pred, baseline=0.90, tolerance=0.02,
    )
    assert_no_bias(
        y_true, y_pred,
        sensitive_feature=gender,
        method="demographic_parity",
    )

@pytest.mark.ml_inference
def test_inference_performance():
    assert_latency(model.predict, X_test, p95=50.0)
    assert_throughput(model.predict, X_single, min_rps=100)

Run with HTML report:

pytest --mltk-report

What's Included (v0.12.4)

Behavioral consistency testing — 6 assertions (paraphrase invariance, format invariance, output stability, semantic equivalence, directional expectation, retrieval consistency) plus built-in ParaphraseGenerator that catch models memorizing phrasing instead of learning concepts. Multi-method evaluation: lexical (token F1), embedding, NLI, LLM-as-Judge. No other ML testing tool ships these as pytest assertions.

NER PII detection — assert_no_pii(method="ner"|"gliner"|"hybrid") adds Named Entity Recognition to PII scanning. Presidio + spaCy catches names, organizations, and locations that regex cannot find. GLiNER provides zero-shot detection for domain-specific entities (healthcare MRN, legal case numbers). Hybrid mode runs regex + NER with intelligent deduplication. Install with pip install mltk[ner].

YAML test definitions — write ML tests in YAML, no Python required. Run with mltk test tests.yaml. Supports 11 data assertions with env:VAR_NAME data source for CI/CD.

EU AI Act compliance reports — mltk compliance results.json auto-maps test results to EU AI Act articles, classifies risk level (unacceptable/high/limited/minimal), and generates an HTML evidence table with gap analysis.

mltk doctor — 9 diagnostic checks that inspect your Python version, dependencies, config, Rust extension, and pytest plugin, with actionable fix hints for each failure.

Data contracts — define column types, ranges, and uniqueness rules in YAML. Use mltk contract generate-tests to auto-generate a pytest file from the contract — zero boilerplate.

LLM evaluation — assert_semantic_similarity, assert_no_toxicity, assert_no_hallucination, assert_ttft, assert_itl, RAG evaluation (faithfulness, context precision/recall, RAGAS composite), agentic testing (task completion, tool selection), multi-turn conversation, coherence, BERTScore, and text quality assertions for GenAI systems.

Training bug detection — catch data leakage (P0), gradient failures (P1), numerical instability, plus P2 bugs: augmentation corruption, checkpoint integrity, distributed training sync, memory leaks, and training-serving skew detection.

Server platform — mltk server starts a self-hosted result-tracking server with REST API, live dashboard, API key auth, webhooks, run comparison, and GitHub CI integration (PR comments + check runs).

Cloud monitoring — AWS SageMaker/CloudWatch, GCP Vertex AI, Azure ML, and Prometheus/Triton adapters for production endpoint health, latency, and error rate assertions.

FDA audit trail — mltk fda-audit generates a device-grade audit trail from test results for regulatory submissions.

Compliance PDF — mltk compliance-pdf exports any HTML compliance report (EU AI Act, OWASP) to PDF for auditors.

Model Card Generator — mltk model-card results.json generates a Google-format Model Card in Markdown from test results.

Chat interface — mltk chat provides interactive Q&A about test results with no LLM or external API required.

Environment variable config — all config keys available as MLTK_* env vars, highest priority in the cascade. Works out of the box in CI/CD pipelines.

JSON export — --mltk-export-json flag exports full test results to JSON for downstream tooling.

Feature Matrix (244 assertion definitions)

Installation

# Core (data + model assertions)
pip install mltk

# With CLI
pip install mltk[cli]

# With HTML reports
pip install mltk[report]

# With server platform
pip install mltk[server]

# NER PII detection
pip install mltk[ner]         # Presidio + spaCy NER

# Domain kits
pip install mltk[cv]          # Computer Vision
pip install mltk[nlp]         # NLP
pip install mltk[speech]      # Speech Recognition
pip install mltk[contracts]   # Data contracts (YAML)
pip install mltk[torch]       # Gradient inspection
pip install mltk[llm]         # LLM evaluation

# Integrations
pip install mltk[integrations]  # Jira adapter
pip install mltk[mlflow]        # MLflow logging
pip install mltk[aws]           # AWS SageMaker/CloudWatch
pip install mltk[gcp]           # GCP Vertex AI
pip install mltk[azure]         # Azure ML monitoring
pip install mltk[pdf]           # Compliance PDF export
pip install mltk[server]        # Self-hosted server platform

# Everything
pip install mltk[all]

CLI

# Core
mltk version                              # Show version
mltk init                                 # Scaffold config + tests
mltk scan data.csv                        # Quick data quality scan
mltk drift ref.csv cur.csv                # Drift analysis
mltk score                                # ML Test Score
mltk doctor                               # Diagnose environment

# Test execution & reporting
mltk test tests.yaml                      # Run YAML-defined tests
mltk model-card results.json              # Generate Google Model Card
mltk compliance results.json              # EU AI Act compliance report
mltk fda-audit results.json              # FDA device audit trail export
mltk compliance-pdf report.html          # Export compliance report to PDF

# Data contracts
mltk contract init                        # Scaffold contract YAML
mltk contract validate data.csv          # Validate against contract
mltk contract generate-tests             # Generate pytest from contract

# Documentation
mltk docs serve                           # Serve docs locally (hot reload)
mltk docs build                           # Build static HTML docs
mltk docs open                            # Build, serve, and open in browser

# Test registry
mltk registry push my_fixtures            # Save test files to registry
mltk registry pull my_fixtures            # Restore from registry
mltk registry list                        # List saved collections

# Notifications
mltk notify slack --results-json r.json  # Send results to Slack

# Server platform
mltk server                               # Start result-tracking server
mltk server-create-key --project myproj  # Generate API key for server

# Interactive
mltk chat --results-json results.json    # Q&A about test results

pytest Plugin

Auto-registered on install. Provides ML-specific markers and --mltk-report:

pytest -m ml_data                 # Run only data quality tests
pytest -m ml_model                # Run only model quality tests
pytest -m "not ml_slow"           # Skip long-running tests
pytest --mltk-report              # Generate HTML report
pytest --mltk-export-json         # Export results to JSON

Markers: ml_data, ml_model, ml_drift, ml_inference, ml_smoke, ml_slow, ml_gpu, ml_nondeterministic

Rust Acceleration

Optional Rust backend for 10-100x speedup on drift detection (KS test, PSI). Falls back to scipy/numpy automatically.

Comparison

Documentation

• Getting Started
• Configuration
• API Reference
• Examples

License

Elastic License 2.0 — free for internal use; commercial licensing available at lior1cc@gmail.com. See LICENSE-COMMERCIAL.