dagshub_mlops_project/fastapi_with_pycaret/diamond_predict_api.py

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

import pandas as pd
from pycaret.regression 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("diamond_predict_api")

# Create input/output pydantic models
input_model = create_model("diamond_predict_api_input", **{'Carat Weight': 1.149999976158142, 'Cut': 'Ideal', 'Color': 'H', 'Clarity': 'SI1', 'Polish': 'VG', 'Symmetry': 'VG', 'Report': 'GIA'})
output_model = create_model("diamond_predict_api_output", prediction=5169)


# 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=8000)