Update index.py

Author: Yusufibin

src/index.py

@@ -1,20 +1,29 @@
-from fastapi import FastAPI
+# -*- coding: utf-8 -*-
 
-from src.dtos.ISayHelloDto import ISayHelloDto
+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("rf_api")
 
-@app.get("/")
-async def root():
-    return {"message": "Hello World"}
+# Create input/output pydantic models
+input_model = create_model("rf_api_input", **{'age': 70.0, 'anaemia': 0.0, 'creatinine_phosphokinase': 582.0, 'diabetes': 1.0, 'ejection_fraction': 38.0, 'high_blood_pressure': 0.0, 'platelets': 25100.0, 'serum_creatinine': 1.100000023841858, 'serum_sodium': 140.0, 'sex': 1.0, 'smoking': 0.0, 'time': 246.0})
+output_model = create_model("rf_api_output", prediction=0)
 
 
-@app.get("/hello/{name}")
-async def say_hello(name: str):
-    return {"message": f"Hello {name}"}
+# 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]}
 
 
-@app.post("/hello")
-async def hello_message(dto: ISayHelloDto):
-    return {"message": f"Hello {dto.message}"}
+if __name__ == "__main__":
+    uvicorn.run(app, host="127.0.0.1", port=8000)