feat: torchscript_and_onnx

main.py

@@ -1,4 +1,4 @@
-## This is for my local path
+# This is for my local path
 # from pathlib import Path
 # import sys
 
@@ -8,7 +8,8 @@
 
 import uvicorn
 from fastapi import FastAPI
-from src.handlers.predict import model_post_handler, graph_post_handler
+from src.handlers.predict_sklearn import sklearn_post_handler
+from src.handlers.predict_pyg import graph_post_handler
 from src.entities.prediction_request import PredictionRequestPydantic
 from fastapi.responses import JSONResponse
 from src.loggers.log_middleware import LogMiddleware
@@ -25,11 +26,9 @@ def health_check():
 @app.post("/predict/")
 def predict(req: PredictionRequestPydantic):
     if req.model["type"] == "SKLEARN":
-        return JSONResponse(content=model_post_handler(req))
-    elif req.model["type"] == "TORCH":
-        return JSONResponse(content=graph_post_handler(req))
+        return JSONResponse(content=sklearn_post_handler(req))
     else:
-        raise ValueError("Only SKLEARN and TORCH models are supported")
+        return JSONResponse(content=graph_post_handler(req))
 
 
 if __name__ == "__main__":

requirements.txt

@@ -2,5 +2,5 @@ fastapi==0.111.0
 pydantic==2.7.1
 uvicorn==0.29.0
 starlette~=0.37.2
-jaqpotpy==6.4.0
+jaqpotpy==6.5.0
 

src/handlers/predict_pyg.py

@@ -0,0 +1,98 @@
+from ..entities.prediction_request import PredictionRequestPydantic
+import base64
+import onnxruntime
+import torch
+import io
+import numpy as np
+import torch.nn.functional as F
+from jaqpotpy.descriptors.graph.graph_featurizer import SmilesGraphFeaturizer
+
+
+def graph_post_handler(request: PredictionRequestPydantic):
+
+    feat_config = request.extraConfig["torchConfig"]["featurizerConfig"]
+    featurizer = _load_featurizer(feat_config)
+    target_name = request.model["dependentFeatures"][0]["name"]
+    model_task = request.model["task"]
+    smiles = request.dataset["input"][0]["SMILES"]
+    data = featurizer.featurize(smiles)
+    raw_model = request.model["rawModel"]
+    if request.model["type"] == "TORCH_ONNX":
+        model_output = onnx_post_handler(raw_model, data)
+        return check_model_task(model_task, target_name, model_output)
+    elif request.model["type"] == "TORCHSCRIPT":
+        model_output = torchscript_post_handler(raw_model, data)
+        return check_model_task(model_task, target_name, model_output)
+
+
+def onnx_post_handler(raw_model, data):
+    onnx_model = base64.b64decode(raw_model)
+    ort_session = onnxruntime.InferenceSession(onnx_model)
+    ort_inputs = {
+        ort_session.get_inputs()[0].name: _to_numpy(data.x),
+        ort_session.get_inputs()[1].name: _to_numpy(data.edge_index),
+        ort_session.get_inputs()[2].name: _to_numpy(
+            torch.zeros(data.x.shape[0], dtype=torch.int64)
+        ),
+    }
+    ort_outs = torch.tensor(np.array(ort_session.run(None, ort_inputs)))
+    return ort_outs
+
+
+def torchscript_post_handler(raw_model, data):
+    torchscript_model = base64.b64decode(raw_model)
+    model_buffer = io.BytesIO(torchscript_model)
+    model_buffer.seek(0)
+    torchscript_model = torch.jit.load(model_buffer)
+    torchscript_model.eval()
+    with torch.no_grad():
+        if data.edge_attr.shape[1] == 0:
+            out = torchscript_model(data.x, data.edge_index, data.batch)
+        else:
+            out = torchscript_model(data.x, data.edge_index, data.batch, data.edge_attr)
+    return out
+
+
+def _to_numpy(tensor):
+    return (
+        tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()
+    )
+
+
+def _load_featurizer(config):
+
+    featurizer = SmilesGraphFeaturizer()
+    featurizer.load_dict(config)
+    featurizer.sort_allowable_sets()
+    return featurizer
+
+
+def graph_regression(target_name, output):
+    preds = [output.squeeze().tolist()]
+    results = {}
+    results[target_name] = [str(pred) for pred in preds]
+    final_all = {"predictions": [dict(zip(results, t)) for t in zip(*results.values())]}
+    return final_all
+
+
+def graph_binary_classification(target_name, output):
+    probs = [F.sigmoid(output).squeeze().tolist()]
+    preds = [int(prob > 0.5) for prob in probs]
+    # UI Results
+    results = {}
+    results["Probabilities"] = [str(prob) for prob in probs]
+    results[target_name] = [str(pred) for pred in preds]
+    final_all = {"predictions": [dict(zip(results, t)) for t in zip(*results.values())]}
+    return final_all
+
+
+def check_model_task(model_task, target_name, out):
+
+    if model_task == "BINARY_CLASSIFICATION":
+        return graph_binary_classification(target_name, out)
+    elif model_task == "REGRESSION":
+        return graph_regression(target_name, out)
+    else:
+        raise ValueError(
+            "Only BINARY_CLASSIFICATION and REGRESSION tasks are supported"
+        )

src/handlers/predict_sklearn.py

@@ -0,0 +1,30 @@
+from ..entities.prediction_request import PredictionRequestPydantic
+from ..helpers import model_decoder, json_to_predreq
+from ..helpers.predict_methods import predict_onnx, predict_proba_onnx
+
+
+def sklearn_post_handler(request: PredictionRequestPydantic):
+    model = model_decoder.decode(request.model["rawModel"])
+    data_entry_all = json_to_predreq.decode(request)
+    prediction = predict_onnx(model, data_entry_all, request)
+    task = request.model["task"].lower()
+    if task == "binary_classification" or task == "multiclass_classification":
+        probabilities = predict_proba_onnx(model, data_entry_all, request)
+    else:
+        probabilities = [None for _ in range(len(prediction))]
+
+    results = {}
+    for i, feature in enumerate(request.model["dependentFeatures"]):
+        key = feature["key"]
+        values = [
+            str(item[i]) if len(request.model["dependentFeatures"]) > 1 else str(item)
+            for item in prediction
+        ]
+        results[key] = values
+
+    results["Probabilities"] = [str(prob) for prob in probabilities]
+    results["AD"] = [None for _ in range(len(prediction))]
+
+    final_all = {"predictions": [dict(zip(results, t)) for t in zip(*results.values())]}
+
+    return final_all

src/helpers/predict_methods.py

@@ -1,6 +1,6 @@
 import numpy as np
 from onnxruntime import InferenceSession
-from jaqpotpy.datasets.molecular_datasets import JaqpotpyDataset
+from jaqpotpy.datasets import JaqpotpyDataset
 from src.helpers.recreate_preprocessor import recreate_preprocessor