inference.py

from typing import Dict, List
import joblib
import numpy as np

def load_model():
    model = joblib.load("model.pkl")
    return model

def predict(inputs: Dict):
    """
    Prediction function for the API
    """
    model = load_model()
    features = np.array(inputs['inputs']).reshape(1, -1)
    prediction = model.predict(features)
    probability = model.predict_proba(features).max()
    
    return {
        "prediction": int(prediction[0]),
        "confidence": float(probability)
    }