File size: 1,932 Bytes
66ca7f0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 | """
Custom inference handler for the MLForge test sklearn LinearRegression model.
This handler lets the Hugging Face Inference API serve sklearn/joblib models
that are not natively supported by the built-in pipeline loaders.
HuggingFace loads this file automatically when the repo uses:
library_name: sklearn
pipeline_tag: tabular-regression
Expected request format:
{"inputs": [[5.0, 3.0]]} → single or multi-sample list of feature vectors
"""
from pathlib import Path
from typing import Any
import joblib
import numpy as np
class EndpointHandler:
"""Custom HF Endpoint handler for sklearn regression models."""
def __init__(self, path: str = ""):
model_dir = Path(path)
# Try the canonical name first, then fall back to any .joblib file
candidate = model_dir / "sklearn_model.joblib"
if not candidate.exists():
matches = list(model_dir.glob("*.joblib"))
if not matches:
raise FileNotFoundError(
f"No .joblib model file found in {model_dir}"
)
candidate = matches[0]
self.model = joblib.load(candidate)
self._model_file = candidate.name
def __call__(self, data: dict[str, Any]) -> list[float]:
"""Run inference.
Args:
data: Dict with key ``"inputs"`` containing a list of feature vectors,
e.g. ``{"inputs": [[5.0, 3.0], [1.0, 2.0]]}``
Returns:
List of predicted float values, one per input row.
"""
raw = data.get("inputs", data.get("data"))
if raw is None:
raise ValueError(
"Request must contain 'inputs' key with a list of feature vectors."
)
X = np.array(raw, dtype=float)
if X.ndim == 1:
X = X.reshape(1, -1)
predictions = self.model.predict(X)
return predictions.tolist()
|