added predict method for a classic ML classifier model (wip)

app.py

@@ -1,47 +1,76 @@
-from typing import Optional
+from typing import Optional, Any
 
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
+from contextlib import asynccontextmanager
+from joblib import load
+
+from models.iris import Iris
 
 
 class Model(BaseModel):
     id: int
     name: str
     param_count: Optional[int] = None
+    _model : Optional[Any] = None
 
 
-app = FastAPI(
-    title="ML Repository API", 
-    description="API for getting predictions from hosted ML models.", 
-    version="0.0.1")
-
-
-models = [Model(id=0, name="CNN"), Model(id=1, name="Transformer")]
-id_2_hosted_models = {
-    model.id : model for model in models
+models = {
+    "0" : Model(id=0, name="CNN"), 
+    "1" : Model(id=1, name="Transformer"),
+    "2" : Model(id=2, name="Iris"),
 }
+id_2_hosted_models = {
+    model.id : model for model in models.values()
+    }
 model_names_2_id = { 
-    model.name.lower() : model.id for model in models
+    model.name.lower() : model.id for model in models.values()
+}
+
+#TODO: fix this mess ^^
+ml_models = {
+    model.name : model for model in models.values()
 }
 
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    # Load the ML model
+    ml_models["Iris"]._model = load('models/iris_v1.joblib')
+    yield
+    # Clean up the ML models and release the resources
+    ml_models.clear()
+
+
+################################################################
+
+
+app = FastAPI(
+    title="ML Repository API", 
+    description="API for getting predictions from hosted ML models.", 
+    version="0.0.1",
+    lifespan=lifespan)
+
 
 @app.get("/")
 def greet_json():
     return {"Hello World": "Welcome to my ML Repository API!"}
 
+
 @app.get("/hosted")
 def list_models():
     "List all the hosted models."
     return models
 
+
 @app.get("/hosted/id/{model_id}")
 def get_by_id(model_id: int):
     "Get a specific model by its ID."
     if model_id not in id_2_hosted_models:
-        raise HTTPException(status_code=404, detail=f"Model with id={model_id} not found")
+        raise HTTPException(status_code=404, detail=f"Model with 'id={model_id}' not found")
 
     return id_2_hosted_models[model_id]
 
+
 @app.get("/hosted/name/{model_name}")
 def get_by_name(model_name: str):
     "Get a specific model by its name."
@@ -50,4 +79,17 @@ def get_by_name(model_name: str):
     if model_name not in model_names_2_id:
         raise HTTPException(status_code=404, detail=f"Model '{model_name}' not found")
     
-    return id_2_hosted_models[model_names_2_id[model_name]]
+    return id_2_hosted_models[model_names_2_id[model_name]]
+
+
+@app.post("/hosted/name/{model_name}/predict", tags=["Predictions"])
+async def get_prediction(model_name: str, iris: Iris):
+    
+    if model_name.lower() != "iris":
+        raise HTTPException(status_code=501, detail="Not implemented yet.")
+    
+    data = dict(iris)['data']
+    prediction = ml_models["Iris"]._model.predict(data).tolist()
+    log_probs = ml_models["Iris"]._model.predict_proba(data).tolist()
+    return {"predictions": prediction,
+            "log_probs": log_probs}

models/iris.py

@@ -0,0 +1,7 @@
+from typing import List
+
+from pydantic import BaseModel, conlist
+
+
+class Iris(BaseModel):
+    data: List[conlist(float, min_length=4, max_length=4)] # type: ignore

models/iris_v1.joblib

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:08b4855f249d786bfaeb015eb7db4ee4d91948434a68a81285bfce06315054e6
+size 3496

models/ml/train.py

@@ -0,0 +1,30 @@
+from typing import Literal, Optional
+
+from joblib import dump
+from sklearn import datasets
+from sklearn.pipeline import Pipeline
+from sklearn.preprocessing import MinMaxScaler
+from sklearn.tree import DecisionTreeClassifier
+
+
+def load_dataset(dataset_name: Literal["iris", "other"]):
+    if dataset_name != "iris":
+        raise NotImplementedError()
+    
+    dataset = datasets.load_iris(return_X_y=True)
+    return dataset[0], dataset[1]
+
+def train_ml_classifier(X, y, output_file: Optional[str] = None):
+    clf_pipeline = [('scaling', MinMaxScaler()), 
+                    ('classifier', DecisionTreeClassifier(random_state=42))]
+    pipeline = Pipeline(clf_pipeline)
+    pipeline.fit(X, y)
+
+    if output_file is not None:
+        dump(pipeline, output_file)
+
+
+if __name__ == '__main__':
+
+    X, y = load_dataset('iris')
+    model = train_ml_classifier(X, y, output_file='./iris_v1.joblib')