File size: 1,054 Bytes
878fe95 |
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 |
import json
import joblib
import numpy as np
from lightgbm import LGBMRegressor
# Load the model once when the endpoint starts
model = joblib.load("lgbm_model.joblib") # <-- Your saved model
def preprocess(inputs):
"""
Convert JSON input into the appropriate format for LightGBM model
inputs: dict or list of dicts
"""
# If input is a single dict, convert it to a list
if isinstance(inputs, dict):
inputs = [inputs]
# Convert to numpy array (feature order must match training!)
feature_order = sorted(inputs[0].keys())
X = np.array([[sample[f] for f in feature_order] for sample in inputs])
return X
def predict(inputs):
X = preprocess(inputs)
preds = model.predict(X)
# Return a list of predictions
return preds.tolist()
def handle(inputs):
"""
Function called by Hugging Face Inference Endpoint
"""
# If input is a JSON string
if isinstance(inputs, str):
inputs = json.loads(inputs)
preds = predict(inputs)
return {"predictions": preds}
|