lr_api.py

# -*- coding: utf-8 -*-

import pandas as pd
from pycaret.classification import load_model, predict_model
from fastapi import FastAPI
import uvicorn
from pydantic import create_model

# Create the app
app = FastAPI()

# Load trained Pipeline
model = load_model("lr_api")

# Create input/output pydantic models
input_model = create_model("lr_api_input", **{'Id': 216, 'WeekofPurchase': 265, 'StoreID': 7, 'PriceCH': 1.8600000143051147, 'PriceMM': 2.130000114440918, 'DiscCH': 0.3700000047683716, 'DiscMM': 0.0, 'SpecialCH': 1, 'SpecialMM': 0, 'LoyalCH': 0.974931001663208, 'SalePriceMM': 2.130000114440918, 'SalePriceCH': 1.4900000095367432, 'PriceDiff': 0.6399999856948853, 'Store7': 'Yes', 'PctDiscMM': 0.0, 'PctDiscCH': 0.19892500340938568, 'ListPriceDiff': 0.27000001072883606, 'STORE': 0})
output_model = create_model("lr_api_output", prediction='CH')


# Define predict function
@app.post("/predict", response_model=output_model)
def predict(data: input_model):
    data = pd.DataFrame([data.dict()])
    predictions = predict_model(model, data=data)
    return {"prediction": predictions["prediction_label"].iloc[0]}


if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)