| """ |
| 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) |
|
|
| |
| 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() |
|
|
