inference.py

import xgboost as xgb
import json
import numpy as np

def model_fn(model_dir):
    # Load the model from Hugging Face Hub
    model = xgb.Booster()
    model.load_model(f"{model_dir}/xgboost_model.json")
    return model

def predict_fn(data, model):
    # Convert input data into DMatrix
    dmatrix = xgb.DMatrix(np.array(data['inputs']))
    prediction = model.predict(dmatrix)
    return prediction.tolist()

if __name__ == "__main__":
    # Example of testing locally
    model = model_fn(".")
    sample_data = {"inputs": [[1, 2, 3], [4, 5, 6]]}  # Replace with your input features
    predictions = predict_fn(sample_data, model)
    print(predictions)