File size: 1,762 Bytes
6f7f49c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 | from fastapi import APIRouter, Depends, HTTPException, Request, Response, UploadFile
from pydantic import BaseModel, Field
from typing import Literal
from routes.task import predict_drug
import skops.io as sio
# Create an instance of the FastAPI class
router = APIRouter(
prefix="/api",
tags=["api"],
responses={404: {"description": "Not found"}}
)
class PredictDrugInput(BaseModel):
age: int = Field(..., ge=15, le=74, description="Age of the patient (15 to 74)")
sex: Literal["M", "F"] = Field(..., description="Sex of the patient (M or F)")
blood_pressure: Literal["HIGH", "LOW", "NORMAL"] = Field(..., description="Blood pressure level")
cholesterol: Literal["HIGH", "NORMAL"] = Field(..., description="Cholesterol level")
na_to_k_ratio: float = Field(..., ge=6.2, le=38.2, description="Sodium-to-potassium ratio in blood (6.2 to 38.2)")
model_config = {
"json_schema_extra": {
"examples": [
{
"age": 30,
"sex": "M",
"blood_pressure": "HIGH",
"cholesterol": "HIGH",
"na_to_k_ratio": 10
}
]
}
}
# Define a GET endpoint
@router.get("/")
def read_root():
return {"message": "Hello, welcome to the demo"}
@router.get("/get_perams")
def get_perams():
pipe = sio.load("./Model/drug_pipeline.skops", trusted=True)
model = pipe.named_steps["model"]
model_params = model.get_params()
return model_params
@router.post("/predict")
def predict_ml(input_data: PredictDrugInput):
label = predict_drug(input_data.age, input_data.sex, input_data.blood_pressure, input_data.cholesterol, input_data.na_to_k_ratio)
return label |