feat: refactor inference (#338)

* feat: refactor inference

* test coverage

* test coverage

* add 'tests' extra so skops and scikit-learn get installed in CI

Author: stephantul

Makefile

@@ -8,7 +8,7 @@ install:
 	uv run pre-commit install
 
 install-no-pre-commit:
-	uv pip install ".[dev,distill,inference,train,onnx,quantization,integration]"
+	uv pip install ".[dev,distill,train,onnx,quantization,integration,tests]"
 
 install-base:
 	uv sync --extra dev

model2vec/inference/README.md

@@ -1,6 +1,6 @@
 # Inference
 
-This subpackage mainly contains helper functions for inference with trained models that have been exported to `scikit-learn` compatible pipelines.
+This subpackage mainly contains helper functions for inference with trained classifier/projector heads, persisted as a `safetensors` file and `config.json` metadata.
 
 If you're looking for information on how to train a model, see [here](../train/README.md).
 
@@ -16,3 +16,16 @@ label = classifier.predict("Attitudes towards cattle in the Alps: a study in let
 ```
 
 This should just work.
+
+# Migrating a legacy pipeline
+
+Pipelines saved by older versions of model2vec store the head as a `scikit-learn`/`skops` `pipeline.skops` file instead of `head.safetensors`. `from_pretrained` still loads these automatically, falling back to the legacy format and emitting a warning. This requires `scikit-learn` and `skops` to be installed.
+
+To upgrade a pipeline to the current format (and silence the warning), convert it with `convert_legacy_pipeline` and save the result:
+
+```python
+from model2vec.inference import convert_legacy_pipeline
+
+pipeline = convert_legacy_pipeline("path/or/repo-id/of/legacy/pipeline")
+pipeline.save_pretrained("path/to/save")
+```

model2vec/inference/__init__.py

@@ -1,10 +1,4 @@
-from model2vec.utils import get_package_extras, importable
+from model2vec.inference.evaluation import evaluate_single_or_multi_label
+from model2vec.inference.model import StaticModelPipeline, convert_legacy_pipeline
 
-_REQUIRED_EXTRA = "inference"
-
-for extra_dependency in get_package_extras("model2vec", _REQUIRED_EXTRA):
-    importable(extra_dependency, _REQUIRED_EXTRA)
-
-from model2vec.inference.model import StaticModelPipeline, evaluate_single_or_multi_label
-
-__all__ = ["StaticModelPipeline", "evaluate_single_or_multi_label"]
+__all__ = ["StaticModelPipeline", "convert_legacy_pipeline", "evaluate_single_or_multi_label"]

model2vec/inference/evaluation.py

@@ -0,0 +1,113 @@
+from __future__ import annotations
+
+from collections.abc import Iterable, Sequence
+from typing import Any, cast
+
+import numpy as np
+
+
+def _is_multi_label_shaped(y: list[int] | list[str] | list[list[int]] | list[list[str]]) -> bool:
+    """Check if the labels are in a multi-label shape."""
+    return isinstance(y, (list, tuple)) and len(y) > 0 and isinstance(y[0], (list, tuple, set))
+
+
+def _one_hot(labels: Sequence[Any], classes: Sequence[Any]) -> np.ndarray:
+    """One-hot encode a flat sequence of labels against a fixed set of classes."""
+    index = {label: position for position, label in enumerate(classes)}
+    encoded = np.zeros((len(labels), len(classes)), dtype=int)
+    for row, label in enumerate(labels):
+        encoded[row, index[label]] = 1
+    return encoded
+
+
+def _multi_hot(label_lists: Iterable[Iterable[Any]], classes: Sequence[Any]) -> np.ndarray:
+    """Multi-hot encode a sequence of label lists against a fixed set of classes."""
+    index = {label: position for position, label in enumerate(classes)}
+    label_lists = list(label_lists)
+    encoded = np.zeros((len(label_lists), len(classes)), dtype=int)
+    for row, labels in enumerate(label_lists):
+        for label in labels:
+            encoded[row, index[label]] = 1
+    return encoded
+
+
+def _precision_recall_f1_support(
+    y_true: np.ndarray, y_pred: np.ndarray
+) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
+    """Compute per-class precision, recall, f1, and support from one-hot / multi-hot encoded labels."""
+    true_positive = ((y_true == 1) & (y_pred == 1)).sum(axis=0).astype(float)
+    false_positive = ((y_true == 0) & (y_pred == 1)).sum(axis=0).astype(float)
+    false_negative = ((y_true == 1) & (y_pred == 0)).sum(axis=0).astype(float)
+    support = y_true.sum(axis=0)
+
+    predicted_positive = true_positive + false_positive
+    actual_positive = true_positive + false_negative
+    precision = np.divide(
+        true_positive, predicted_positive, out=np.zeros_like(true_positive), where=predicted_positive > 0
+    )
+    recall = np.divide(true_positive, actual_positive, out=np.zeros_like(true_positive), where=actual_positive > 0)
+    precision_plus_recall = precision + recall
+    f1 = np.divide(
+        2 * precision * recall, precision_plus_recall, out=np.zeros_like(precision), where=precision_plus_recall > 0
+    )
+
+    return precision, recall, f1, support
+
+
+def evaluate_single_or_multi_label(
+    predictions: np.ndarray,
+    y: list[int] | list[str] | list[list[int]] | list[list[str]],
+) -> dict[str, dict[str, float]]:
+    """Evaluate the classifier on a given dataset using a classification report.
+
+    This function computes per-class precision, recall and f1-score (via one-vs-rest / multi-hot encoding), plus
+    overall accuracy, macro average, and weighted average.
+
+    :param predictions: The predictions.
+    :param y: The ground truth labels.
+    :return: A classification report, as a dictionary.
+    """
+    if _is_multi_label_shaped(y):
+        y = cast(list[list[str]] | list[list[int]], y)
+        predictions = cast(np.ndarray, predictions)
+        y_labels = {label for labels in y for label in labels}
+        predicted_labels = {label for labels in predictions for label in labels}
+        classes = sorted(y_labels | predicted_labels)
+        y_transformed = _multi_hot(y, classes)
+        predictions_transformed = _multi_hot(predictions, classes)
+    else:
+        y = cast(list[str] | list[int], y)
+        classes = sorted(set(y) | set(predictions.tolist()))
+        y_transformed = _one_hot(y, classes)
+        predictions_transformed = _one_hot(predictions.tolist(), classes)
+
+    target_names = [str(c) for c in classes]
+    precision, recall, f1, support = _precision_recall_f1_support(y_transformed, predictions_transformed)
+    total_support = float(support.sum())
+    accuracy = float(np.all(y_transformed == predictions_transformed, axis=1).mean())
+
+    report: dict[str, Any] = {
+        name: {
+            "precision": float(precision[idx]),
+            "recall": float(recall[idx]),
+            "f1-score": float(f1[idx]),
+            "support": float(support[idx]),
+        }
+        for idx, name in enumerate(target_names)
+    }
+    report["accuracy"] = accuracy
+    report["macro avg"] = {
+        "precision": float(precision.mean()),
+        "recall": float(recall.mean()),
+        "f1-score": float(f1.mean()),
+        "support": total_support,
+    }
+    weights = support / total_support if total_support > 0 else np.zeros_like(support, dtype=float)
+    report["weighted avg"] = {
+        "precision": float((precision * weights).sum()),
+        "recall": float((recall * weights).sum()),
+        "f1-score": float((f1 * weights).sum()),
+        "support": total_support,
+    }
+
+    return report

model2vec/inference/mlp.py

@@ -0,0 +1,79 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from enum import Enum
+
+import numpy as np
+
+
+class Activation(str, Enum):
+    SOFTMAX = "softmax"
+    SIGMOID = "sigmoid"
+    IDENTITY = "identity"
+
+
+def _softmax(x: np.ndarray) -> np.ndarray:
+    """Numerically stable softmax over the last axis."""
+    shifted = x - x.max(axis=-1, keepdims=True)
+    exponentiated = np.exp(shifted)
+    return exponentiated / exponentiated.sum(axis=-1, keepdims=True)
+
+
+def _sigmoid(x: np.ndarray) -> np.ndarray:
+    """Numerically stable sigmoid."""
+    return np.where(x >= 0, 1 / (1 + np.exp(-x)), np.exp(x) / (1 + np.exp(x)))
+
+
+@dataclass
+class Layer:
+    weight: np.ndarray
+    bias: np.ndarray
+
+    def __call__(self, x: np.ndarray) -> np.ndarray:
+        """Apply the linear transformation."""
+        return x @ self.weight.T + self.bias
+
+
+class MLPHead:
+    def __init__(
+        self,
+        layers: list[Layer],
+        activation: Activation,
+        classes: np.ndarray | None = None,
+    ) -> None:
+        """An MLP with ReLU activation.
+
+        :param layers: The linear layers, in order.
+        :param activation: The output activation.
+        :param classes: The classes, if the task is a classification task.
+        """
+        self.layers = layers
+        self.activation = activation
+        self.classes_ = classes
+
+    def _logits(self, X: np.ndarray) -> np.ndarray:
+        """Run the forward through the layers."""
+        out = X
+        *hidden_layers, last_layer = self.layers
+        for layer in hidden_layers:
+            out = np.maximum(layer(out), 0.0)
+        return last_layer(out)
+
+    def predict_proba(self, X: np.ndarray) -> np.ndarray:
+        """Predict probabilities, applying the output activation to the raw logits."""
+        logits = self._logits(X)
+        match self.activation:
+            case Activation.SOFTMAX:
+                return _softmax(logits)
+            case Activation.SIGMOID:
+                return _sigmoid(logits)
+            case Activation.IDENTITY:
+                return logits
+
+    def predict_index(self, X: np.ndarray) -> np.ndarray:
+        """Predict the index of the most likely class."""
+        return self._logits(X).argmax(axis=1)
+
+    def predict_regression(self, X: np.ndarray) -> np.ndarray:
+        """Predict the raw (identity-activation) output, e.g. for a projector head."""
+        return self._logits(X)