Update app.py

app.py

@@ -3,8 +3,8 @@
 import pandas as pd
 from pycaret.regression import load_model, predict_model
 from fastapi import FastAPI
+from pydantic import BaseModel
 import uvicorn
-from pydantic import create_model
 
 # Create the app
 app = FastAPI()
@@ -13,13 +13,27 @@ app = FastAPI()
 model = load_model("lr_api")
 
 # Create input/output pydantic models
-input_model = create_model("lr_api_input", **{'rownames': 1030, 'year': 1994, 'violent': 304.5, 'murder': 2.9000000953674316, 'prisoners': 152, 'afam': 1.769081950187683, 'cauc': 70.66014862060547, 'male': 18.20832061767578, 'population': 1.9304360151290894, 'income': 12036.8603515625, 'density': 0.023493800312280655, 'state': 'Utah', 'law': 'yes'})
-output_model = create_model("lr_api_output", prediction=63.6)
-
+class InputModel(BaseModel):
+    rownames: int
+    year: int
+    violent: float
+    murder: float
+    prisoners: int
+    afam: float
+    cauc: float
+    male: float
+    population: float
+    income: float
+    density: float
+    state: str
+    law: str
+
+class OutputModel(BaseModel):
+    prediction: float
 
 # Define predict function
-@app.post("/predict", response_model=output_model)
-def predict(data: input_model):
+@app.post("/predict", response_model=OutputModel)
+def predict(data: InputModel):
     data = pd.DataFrame([data.dict()])
     predictions = predict_model(model, data=data)
     return {"prediction": predictions["prediction_label"].iloc[0]}