trustbench 0.2.0__py3-none-any.whl

trustbench/__init__.py

@@ -0,0 +1,121 @@
+"""TrustBench: Evaluation framework for machine learning model reliability."""
+
+__version__ = "0.2.0"
+
+from trustbench.benchmark import TrustBench
+from trustbench.core.evaluator import evaluate
+from trustbench.core.exceptions import (
+    SecurityError,
+    ShapeError,
+    TrustBenchError,
+    UnsupportedCapabilityError,
+)
+from trustbench.core.protocols import DatasetAdapter, ModelAdapter
+from trustbench.core.results import EvaluationResult
+from trustbench.evaluation.distribution_shift import (
+    ShiftType,
+    compute_label_distribution_shift,
+    compute_ood_detection_metrics,
+    evaluate_distribution_shift,
+)
+from trustbench.evaluation.fairness import compute_fairness_metrics
+from trustbench.evaluation.regression import RegressionResult, evaluate_regression
+from trustbench.evaluation.retrieval import RetrievalResult, evaluate_retrieval
+from trustbench.evaluation.segmentation import (
+    SegmentationResult,
+    evaluate_segmentation,
+)
+from trustbench.metrics.regression import compute_regression_metrics
+from trustbench.metrics.retrieval import compute_retrieval_metrics
+from trustbench.metrics.segmentation import compute_segmentation_metrics
+from trustbench.registry import (
+    GenericDatasetAdapter,
+    GenericModelAdapter,
+    discover_plugins,
+    list_registered_corruptions,
+    list_registered_metrics,
+    load_plugin,
+    register_corruption,
+    register_dataset_adapter,
+    register_metric,
+    register_model_adapter,
+    validate_dataset_adapter,
+    validate_model_adapter,
+)
+
+__all__ = [
+    "DatasetAdapter",
+    "EvaluationResult",
+    "GenericDatasetAdapter",
+    "GenericModelAdapter",
+    "JAXModelAdapter",
+    "KerasModelAdapter",
+    "ModelAdapter",
+    "ONNXModelAdapter",
+    "PyTorchModelAdapter",
+    "RegressionResult",
+    "RetrievalResult",
+    "SecurityError",
+    "SegmentationResult",
+    "ShapeError",
+    "ShiftType",
+    "SklearnModelAdapter",
+    "TrustBench",
+    "TrustBenchError",
+    "UnsupportedCapabilityError",
+    "__version__",
+    "compute_fairness_metrics",
+    "compute_label_distribution_shift",
+    "compute_ood_detection_metrics",
+    "compute_regression_metrics",
+    "compute_retrieval_metrics",
+    "compute_segmentation_metrics",
+    "discover_plugins",
+    "evaluate",
+    "evaluate_distribution_shift",
+    "evaluate_regression",
+    "evaluate_retrieval",
+    "evaluate_segmentation",
+    "list_registered_corruptions",
+    "list_registered_metrics",
+    "load_plugin",
+    "register_corruption",
+    "register_dataset_adapter",
+    "register_metric",
+    "register_model_adapter",
+    "validate_dataset_adapter",
+    "validate_model_adapter",
+]
+
+
+def __getattr__(name: str):
+    if name == "ExperimentManager":
+        from trustbench.experiments.manager import ExperimentManager
+
+        return ExperimentManager
+    if name in ("FilesystemExperimentStore", "SQLiteExperimentStore"):
+        from trustbench import storage
+
+        return getattr(storage, name)
+    if name == "SklearnModelAdapter":
+        from trustbench.models.sklearn import SklearnModelAdapter
+
+        return SklearnModelAdapter
+    if name == "PyTorchModelAdapter":
+        from trustbench.models.pytorch import PyTorchModelAdapter
+
+        return PyTorchModelAdapter
+    if name == "ONNXModelAdapter":
+        from trustbench.models.onnx import ONNXModelAdapter
+
+        return ONNXModelAdapter
+    if name == "KerasModelAdapter":
+        from trustbench.models.keras import KerasModelAdapter
+
+        return KerasModelAdapter
+    if name == "JAXModelAdapter":
+        from trustbench.models.jax import JAXModelAdapter
+
+        return JAXModelAdapter
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+

trustbench/api/__init__.py

@@ -0,0 +1,5 @@
+"""API module for TrustBench."""
+
+from trustbench.api.app import app, create_app
+
+__all__ = ["app", "create_app"]

trustbench/api/app.py

@@ -0,0 +1,45 @@
+"""FastAPI application factory for TrustBench API."""
+
+from contextlib import asynccontextmanager
+
+from fastapi import FastAPI
+from fastapi.middleware.cors import CORSMiddleware
+
+from trustbench.api.routes.experiments import router as experiments_router
+from trustbench.storage.database import init_db
+
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    # Initialize database tables on startup
+    init_db()
+    yield
+
+
+def create_app() -> FastAPI:
+    """Create and configure TrustBench FastAPI application."""
+    app = FastAPI(
+        title="TrustBench API",
+        description="ML Reliability, Robustness & Uncertainty Evaluation Framework API",
+        version="0.1.0",
+        lifespan=lifespan,
+    )
+
+    app.add_middleware(
+        CORSMiddleware,
+        allow_origins=["*"],
+        allow_credentials=True,
+        allow_methods=["*"],
+        allow_headers=["*"],
+    )
+
+    app.include_router(experiments_router)
+
+    @app.get("/health", tags=["health"])
+    def health_check():
+        return {"status": "ok", "service": "trustbench-api"}
+
+    return app
+
+
+app = create_app()

trustbench/api/routes/experiments.py

@@ -0,0 +1,82 @@
+"""FastAPI routes for experiment management, metrics querying, artifacts, and comparisons."""
+
+from fastapi import APIRouter, Depends, HTTPException, Query
+from sqlalchemy.orm import Session
+
+from trustbench.experiments.manager import ExperimentManager
+from trustbench.experiments.schemas import (
+    ExperimentComparisonRequest,
+    ExperimentComparisonResponse,
+    ExperimentDetail,
+    ExperimentResponse,
+)
+from trustbench.storage.database import get_db_session
+
+router = APIRouter(prefix="/experiments", tags=["experiments"])
+
+
+def get_manager(session: Session = Depends(get_db_session)) -> ExperimentManager:
+    """Dependency for ExperimentManager."""
+    return ExperimentManager(session=session)
+
+
+@router.get("", response_model=list[ExperimentResponse])
+def list_experiments(
+    experiment_type: str | None = Query(None, description="Filter by experiment type (e.g. baseline, robustness)"),
+    status: str | None = Query(None, description="Filter by status (pending, running, completed, failed)"),
+    checkpoint_sha256: str | None = Query(None, description="Filter by checkpoint SHA-256 hash"),
+    manager: ExperimentManager = Depends(get_manager),
+):
+    """List all registered experiments."""
+    return manager.list_experiments(
+        experiment_type=experiment_type,
+        status=status,
+        checkpoint_sha256=checkpoint_sha256,
+    )
+
+
+@router.get("/{experiment_id}", response_model=ExperimentDetail)
+def get_experiment_detail(
+    experiment_id: str,
+    manager: ExperimentManager = Depends(get_manager),
+):
+    """Get full details of a specific experiment."""
+    exp = manager.get_experiment_detail(experiment_id)
+    if not exp:
+        raise HTTPException(status_code=404, detail=f"Experiment '{experiment_id}' not found.")
+    return exp
+
+
+@router.post("/compare", response_model=ExperimentComparisonResponse)
+def compare_experiments(
+    req: ExperimentComparisonRequest,
+    manager: ExperimentManager = Depends(get_manager),
+):
+    """Compare multiple experiments across shared conditions and metrics."""
+    result = manager.compare_experiments(
+        experiment_ids=req.experiment_ids,
+        conditions=req.conditions,
+        metrics=req.metrics,
+    )
+    return result
+
+
+@router.post("/index-directory")
+def index_artifact_directory(
+    directory_path: str = Query(..., description="Path to on-disk artifact directory"),
+    experiment_name: str | None = Query(None, description="Optional custom name"),
+    manager: ExperimentManager = Depends(get_manager),
+):
+    """Index an on-disk artifact folder into the database registry."""
+    try:
+        model = manager.registry.index_artifact_directory(directory_path, experiment_name=experiment_name)
+        return {
+            "message": "Experiment indexed successfully.",
+            "experiment_id": model.id,
+            "name": model.name,
+            "status": model.status,
+        }
+    except FileNotFoundError as e:
+        raise HTTPException(status_code=404, detail=str(e))
+    except (ValueError, RuntimeError, OSError) as e:
+        raise HTTPException(status_code=500, detail=str(e))

trustbench/benchmark.py

@@ -0,0 +1,264 @@
+"""TrustBench — public high-level API."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any
+
+from trustbench.core.evaluator import evaluate as _core_evaluate
+from trustbench.core.protocols import DatasetAdapter, ModelAdapter
+from trustbench.core.results import EvaluationResult
+
+
+def _wrap_model(model: Any, **kwargs: Any) -> ModelAdapter:
+    """Coerce *model* to a :class:`ModelAdapter`."""
+    # Already a ModelAdapter — use directly
+    if isinstance(model, ModelAdapter):
+        return model
+
+    # torch.nn.Module — wrap in PyTorchModelAdapter
+    try:
+        from torch import nn
+
+        from trustbench.models.pytorch import PyTorchModelAdapter
+    except ImportError as exc:
+        raise ImportError(
+            "PyTorch is required to auto-wrap an nn.Module. "
+            "Install it with: pip install torch"
+        ) from exc
+
+    if isinstance(model, nn.Module):
+        num_classes = kwargs.get("num_classes")
+        if num_classes is None:
+            raise ValueError(
+                "Pass num_classes=<int> when providing a raw nn.Module to TrustBench."
+            )
+        device = kwargs.get("device")
+        return PyTorchModelAdapter(model, num_classes=num_classes, device=device)
+
+    # Scikit-learn estimator
+    if hasattr(model, "predict") and hasattr(model, "fit"):
+        from trustbench.models.sklearn import SklearnModelAdapter
+
+        task = kwargs.get("task", "classification")
+        num_classes = kwargs.get("num_classes")
+        return SklearnModelAdapter(model, num_classes=num_classes, task=task)
+
+    # ONNX InferenceSession
+    if hasattr(model, "run") and hasattr(model, "get_inputs"):
+        from trustbench.models.onnx import ONNXModelAdapter
+
+        num_classes = kwargs.get("num_classes", 2)
+        task = kwargs.get("task", "classification")
+        return ONNXModelAdapter(model, num_classes=num_classes, task=task)
+
+    raise TypeError(
+        f"model must be a ModelAdapter, nn.Module, scikit-learn estimator, or ONNX session, got {type(model).__name__}."
+    )
+
+
+def _wrap_dataset(dataset: Any) -> DatasetAdapter:
+    """Coerce *dataset* to a :class:`DatasetAdapter`."""
+    if isinstance(dataset, DatasetAdapter):
+        return dataset
+
+    try:
+        from torch.utils.data import DataLoader, Dataset
+
+        from trustbench.datasets.torch_dataset import TorchDatasetAdapter
+    except ImportError as exc:
+        raise ImportError(
+            "PyTorch is required to auto-wrap a Dataset/DataLoader. "
+            "Install it with: pip install torch"
+        ) from exc
+
+    if isinstance(dataset, DataLoader):
+        # Extract the underlying Dataset from the DataLoader
+        dataset = dataset.dataset  # type: ignore[assignment]
+
+    if isinstance(dataset, Dataset):
+        return TorchDatasetAdapter(dataset)
+
+    raise TypeError(
+        f"dataset must be a DatasetAdapter, torch Dataset, or DataLoader, "
+        f"got {type(dataset).__name__}."
+    )
+
+
+class TrustBench:
+    """High-level evaluation interface.
+
+    Example::
+
+        import torch.nn as nn
+        from trustbench import TrustBench
+
+        model = nn.Linear(4, 3)
+        # dataset is any torch.utils.data.Dataset or DatasetAdapter
+        tb = TrustBench(model, dataset, num_classes=3)
+        result = tb.evaluate()
+
+    Args:
+        model: A :class:`~trustbench.core.protocols.ModelAdapter`, ``nn.Module``
+               (requires *num_classes* kwarg), DataLoader, or Dataset.
+        dataset: A :class:`~trustbench.core.protocols.DatasetAdapter`,
+                 ``torch.utils.data.Dataset``, or ``torch.utils.data.DataLoader``.
+        metrics: Metric group names to compute.  Defaults to
+                 ``("classification", "calibration", "uncertainty")``.
+        batch_size: Inference batch size.
+        class_names: Optional list of human-readable class labels.
+        seed: Global random seed.
+        output_dir: If provided, artifacts are saved here after :meth:`evaluate`.
+        num_classes: Required when *model* is a raw ``nn.Module``.
+        device: Device string passed to :class:`PyTorchModelAdapter`.
+    """
+
+    def __init__(
+        self,
+        model: Any,
+        dataset: Any,
+        *,
+        metrics: list[str] | tuple[str, ...] | None = None,
+        batch_size: int = 100,
+        class_names: list[str] | tuple[str, ...] | None = None,
+        seed: int | None = None,
+        output_dir: str | Path | None = None,
+        # forwarded to _wrap_model when model is nn.Module
+        num_classes: int | None = None,
+        device: str | None = None,
+    ) -> None:
+        self._model_adapter: ModelAdapter = _wrap_model(
+            model, num_classes=num_classes, device=device
+        )
+        self._dataset_adapter: DatasetAdapter = _wrap_dataset(dataset)
+        self._metrics = metrics
+        self._batch_size = batch_size
+        self._class_names = class_names
+        self._seed = seed
+        self._output_dir = Path(output_dir) if output_dir is not None else None
+
+    
+    # Class-method constructor from config
+    
+
+    @classmethod
+    def from_config(cls, config_path: str | Path) -> TrustBench:
+        """Construct a :class:`TrustBench` instance from a YAML config file.
+
+        The *model* and *dataset* must be resolvable from the config (i.e. a
+        checkpoint path and architecture name must be provided).
+
+        Args:
+            config_path: Path to a ``trustbench.yaml`` config file.
+
+        Returns:
+            Configured :class:`TrustBench` instance ready to call :meth:`evaluate`.
+        """
+        from trustbench.config.loader import load_config
+        from trustbench.models.checkpoint import load_model_from_checkpoint
+        from trustbench.models.pytorch import PyTorchModelAdapter
+
+        cfg = load_config(config_path)
+        mcfg = cfg.model
+        dcfg = cfg.dataset
+        rcfg = cfg.reporting
+
+        #  Model 
+        if mcfg.checkpoint_path is None:
+            raise ValueError(
+                "config.model.checkpoint_path is required for TrustBench.from_config()."
+            )
+        if mcfg.arch_name is None:
+            raise ValueError(
+                "config.model.arch_name is required for TrustBench.from_config()."
+            )
+
+        ckpt_path = cfg.resolve_path(mcfg.checkpoint_path)
+        nn_module, _ = load_model_from_checkpoint(
+            ckpt_path, model_name=mcfg.arch_name, device=mcfg.device
+        )
+        model_adapter = PyTorchModelAdapter(
+            nn_module, num_classes=mcfg.num_classes, device=mcfg.device
+        )
+
+        #  Dataset 
+        dataset_adapter = cls._build_dataset_from_config(cfg)
+
+        output_dir: Path | None = Path(
+            cfg.resolve_path(rcfg.output_dir)
+        )
+
+        class_names = None
+        if dcfg.dataset_type.lower() == "cifar10":
+            from trustbench.datasets.cifar10 import CIFAR10_CLASSES
+
+            class_names = list(CIFAR10_CLASSES)
+
+        return cls(
+            model_adapter,
+            dataset_adapter,
+            metrics=cfg.metrics.groups,
+            batch_size=dcfg.batch_size,
+            seed=cfg.seed,
+            output_dir=output_dir,
+            class_names=class_names,
+        )
+
+    @staticmethod
+    def _build_dataset_from_config(cfg: Any) -> DatasetAdapter:
+        """Instantiate the correct DatasetAdapter from config."""
+        dcfg = cfg.dataset
+        dtype = dcfg.dataset_type.lower()
+
+        if dtype == "cifar10":
+            from trustbench.datasets.cifar10 import (
+                CIFAR10DatasetAdapter,
+                get_cifar10_datasets,
+            )
+
+            data_dir = (
+                str(cfg.resolve_path(dcfg.data_dir))
+                if dcfg.data_dir
+                else "./data"
+            )
+            trainset, testset = get_cifar10_datasets(data_dir=data_dir, download=True)
+            split_dataset = trainset if dcfg.split == "train" else testset
+            return CIFAR10DatasetAdapter(split_dataset)
+
+        raise ValueError(
+            f"Unknown dataset_type {dcfg.dataset_type!r}. "
+            "Supported types: 'cifar10'."
+        )
+
+    
+    # Evaluation
+    
+
+    def evaluate(self) -> EvaluationResult:
+        """Run the full evaluation pipeline and return an :class:`EvaluationResult`.
+
+        If *output_dir* was provided at construction, artifacts are automatically
+        saved via :func:`~trustbench.reporting.save_result`.
+
+        Returns:
+            :class:`EvaluationResult` containing metrics, predictions, and metadata.
+        """
+        result = _core_evaluate(
+            self._model_adapter,
+            self._dataset_adapter,
+            metrics=self._metrics,
+            batch_size=self._batch_size,
+            class_names=self._class_names,
+            seed=self._seed,
+        )
+
+        if self._output_dir is not None:
+            from trustbench.reporting.result_io import save_result
+
+            save_result(result, self._output_dir)
+
+        return result
+
+    def run(self) -> EvaluationResult:
+        """Alias for :meth:`evaluate`."""
+        return self.evaluate()

trustbench/cli/__init__.py

@@ -0,0 +1,5 @@
+"""TrustBench CLI package."""
+
+from trustbench.cli.main import app
+
+__all__ = ["app"]