api change2

main.py

@@ -1,35 +1,41 @@
 from fastapi import FastAPI
 from pydantic import BaseModel
+from joblib import load
 import numpy as np
-import joblib
 
 app = FastAPI(
-    title="Iris KNN Prediction API",
-    description="API for predicting Iris species using KNN model",
-    version="1.0.0"
+    title="Iris Species Prediction API",
+    description="An API to predict Iris species with class probabilities.",
+    version="1.1.0"
 )
 
+# Load model at startup
 try:
-    model, target_names = joblib.load("iris_knn.pkl")
-except:
+    model = load("iris_random_forest.joblib")
+    target_names = ["setosa", "versicolor", "virginica"]
+except FileNotFoundError:
     model = None
     target_names = []
 
+
 class IrisData(BaseModel):
     sepal_length: float
     sepal_width: float
     petal_length: float
     petal_width: float
 
+
 @app.get("/")
-def root():
-    return {"message": "Iris KNN API Running! Visit /docs to test the API."}
+def read_root():
+    return {"message": "Welcome! Use POST /predict to classify Iris data."}
+
 
 @app.post("/predict")
 def predict_iris(data: IrisData):
     if model is None:
-        return {"error": "Model not found on server"}
+        return {"error": "Model not found on server."}
 
+    # Convert input
     arr = np.array([[
         data.sepal_length,
         data.sepal_width,
@@ -37,9 +43,15 @@ def predict_iris(data: IrisData):
         data.petal_width
     ]])
 
-    pred = model.predict(arr)[0]
+    # Predictions
+    pred_idx = int(model.predict(arr)[0])
+    probas = model.predict_proba(arr)[0]
+
     return {
         "input": data.dict(),
-        "predicted_class": str(target_names[pred])
+        "predicted_class_index": pred_idx,
+        "predicted_class_name": target_names[pred_idx],
+        "class_probabilities": {
+            target_names[i]: float(probas[i]) for i in range(len(target_names))
+        }
     }
-