adding support for raw ONNX models

.github/workflows/surrealml_core_onnx_test.yml

@@ -0,0 +1,49 @@
+name: Run Onnx Tests
+
+on:
+  pull_request:
+    types: [opened, reopened, synchronize]
+
+jobs:
+  test_core:
+    runs-on: ubuntu-latest
+
+    steps:
+    - uses: actions/checkout@v3
+      with:
+        fetch-depth: 0
+
+    - name: Set up Rust
+      uses: actions-rs/toolchain@v1
+      with:
+        profile: minimal
+        toolchain: stable
+        override: true
+
+    - name: Set up Python
+      uses: actions/setup-python@v2
+      with:
+        python-version: '3.11'
+
+    - name: Pre-test Setup
+      run: |
+        python3 -m venv venv
+        source venv/bin/activate
+        pip install --upgrade pip
+        pip install -r requirements.txt
+
+        # build the local version of the core module to be loaded into python
+        echo "Building local version of core module"
+
+        pip install .
+        export PYTHONPATH="."
+
+        python ./tests/scripts/ci_local_build.py
+        echo "Local build complete"
+
+        # train the models for the tests
+        python ./tests/model_builder/onnx_assets.py
+        deactivate
+
+    - name: Run Core Unit Tests
+      run: cd modules/core && cargo test --features onnx-tests

.github/workflows/surrealml_core_tensorflow_test.yml

@@ -0,0 +1,49 @@
+name: Run Tensorflow Tests
+
+on:
+  pull_request:
+    types: [opened, reopened, synchronize]
+
+jobs:
+  test_core:
+    runs-on: ubuntu-latest
+
+    steps:
+    - uses: actions/checkout@v3
+      with:
+        fetch-depth: 0
+
+    - name: Set up Rust
+      uses: actions-rs/toolchain@v1
+      with:
+        profile: minimal
+        toolchain: stable
+        override: true
+
+    - name: Set up Python
+      uses: actions/setup-python@v2
+      with:
+        python-version: '3.11'
+
+    - name: Pre-test Setup
+      run: |
+        python3 -m venv venv
+        source venv/bin/activate
+        pip install --upgrade pip
+        pip install -r requirements.txt
+
+        # build the local version of the core module to be loaded into python
+        echo "Building local version of core module"
+
+        pip install .
+        export PYTHONPATH="."
+
+        python ./tests/scripts/ci_local_build.py
+        echo "Local build complete"
+
+        # train the models for the tests
+        python ./tests/model_builder/tensorflow_assets.py
+        deactivate
+
+    - name: Run Core Unit Tests
+      run: cd modules/core && cargo test --features tensorflow-tests

.github/workflows/surrealml_core_test.yml

@@ -1,4 +1,4 @@
-name: Run tests on Pull Request
+name: Run SkLearn Tests
 
 on:
   pull_request:
@@ -36,22 +36,24 @@ jobs:
         echo "Building local version of core module"
 
         pip install .
+        export PYTHONPATH="."
 
         python ./tests/scripts/ci_local_build.py
         echo "Local build complete"
 
         # train the models for the tests
-        python ./tests/scripts/build_assets.py
+        python ./tests/model_builder/sklearn_assets.py
         deactivate
 
     - name: Run Python Unit Tests
       run: |
         source venv/bin/activate
-        python -m unittest discover
+        export PYTHONPATH="."
+        python tests/unit_tests/engine/test_sklearn.py
         deactivate
 
     - name: Run Core Unit Tests
-      run: cd modules/core && cargo test
+      run: cd modules/core && cargo test --features sklearn-tests
 
     - name: Run HTTP Transfer Tests
       run: cargo test

.github/workflows/surrealml_core_torch_test.yml

@@ -0,0 +1,58 @@
+name: Run Torch Tests
+
+on:
+  pull_request:
+    types: [opened, reopened, synchronize]
+
+jobs:
+  test_core:
+    runs-on: ubuntu-latest
+
+    steps:
+    - uses: actions/checkout@v3
+      with:
+        fetch-depth: 0
+
+    - name: Set up Rust
+      uses: actions-rs/toolchain@v1
+      with:
+        profile: minimal
+        toolchain: stable
+        override: true
+
+    - name: Set up Python
+      uses: actions/setup-python@v2
+      with:
+        python-version: '3.11'
+
+    - name: Pre-test Setup
+      run: |
+        python3 -m venv venv
+        source venv/bin/activate
+        pip install --upgrade pip
+        pip install -r requirements.txt
+
+        # build the local version of the core module to be loaded into python
+        echo "Building local version of core module"
+
+        pip install .
+        export PYTHONPATH="."
+
+        python ./tests/scripts/ci_local_build.py
+        echo "Local build complete"
+
+        # train the models for the tests
+        python ./tests/model_builder/torch_assets.py
+        deactivate
+
+    - name: Run Python Unit Tests
+      run: |
+        source venv/bin/activate
+        export PYTHONPATH="."
+        python tests/unit_tests/engine/test_torch.py
+        python tests/unit_tests/test_rust_adapter.py
+        python tests/unit_tests/test_surml_file.py
+        deactivate
+
+    - name: Run Core Unit Tests
+      run: cd modules/core && cargo test --features torch-tests

README.md

@@ -96,6 +96,15 @@ print(new_file.buffered_compute(value_map={"squarefoot": 5, "num_floors": 6}))
 print(new_file.raw_compute(input_vector=[5, 6]))
 ```
 
+## Raw ONNX models
+
+You may not have a model that is supported by the `surrealml` library. However, if you can convert the model into ONNX
+format by yourself, you can merely use the `ONNX` engine when saving your model with the following code:
+
+```python
+file = SurMlFile(model=raw_onnx_model, name="linear", inputs=HOUSE_LINEAR["inputs"], engine=Engine.ONNX)
+```
+
 ## Python tutorial using Pytorch
 
 First we need to have one script where we create and store the model. In this example we will merely do a linear regression model

modules/core/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "surrealml-core"
-version = "0.1.2"
+version = "0.1.3"
 edition = "2021"
 build = "./build.rs"
 description = "The core machine learning library for SurrealML that enables SurrealDB to store and load ML models"
@@ -12,6 +12,11 @@ license-file = "LICENSE"
 axum-feature = ["axum"]
 actix-feature = ["actix-web"]
 default = []
+# below are the features for testing different engines
+sklearn-tests = []
+onnx-tests = []
+torch-tests = []
+tensorflow-tests = []
 
 [dependencies]
 regex = "1.9.3"

modules/core/src/execution/compute.rs

@@ -160,6 +160,7 @@ mod tests {
 
     use super::*;
 
+    #[cfg(feature = "sklearn-tests")]
     #[test]
     fn test_raw_compute_linear_sklearn() {
         let mut file = SurMlFile::from_file("./model_stash/sklearn/surml/linear.surml").unwrap();
@@ -178,7 +179,7 @@ mod tests {
         assert_eq!(output[0], 985.57745);
     }
 
-
+    #[cfg(feature = "sklearn-tests")]
     #[test]
     fn test_buffered_compute_linear_sklearn() {
         let mut file = SurMlFile::from_file("./model_stash/sklearn/surml/linear.surml").unwrap();
@@ -194,6 +195,42 @@ mod tests {
         assert_eq!(output.len(), 1);
     }
 
+    #[cfg(feature = "onnx-tests")]
+    #[test]
+    fn test_raw_compute_linear_onnx() {
+        let mut file = SurMlFile::from_file("./model_stash/onnx/surml/linear.surml").unwrap();
+        let model_computation = ModelComputation {
+            surml_file: &mut file,
+        };
+
+        let mut input_values = HashMap::new();
+        input_values.insert(String::from("squarefoot"), 1000.0);
+        input_values.insert(String::from("num_floors"), 2.0);
+
+        let raw_input = model_computation.input_tensor_from_key_bindings(input_values).unwrap();
+
+        let output = model_computation.raw_compute(raw_input, Some((1, 2))).unwrap();
+        assert_eq!(output.len(), 1);
+        assert_eq!(output[0], 985.57745);
+    }
+
+    #[cfg(feature = "onnx-tests")]
+    #[test]
+    fn test_buffered_compute_linear_onnx() {
+        let mut file = SurMlFile::from_file("./model_stash/onnx/surml/linear.surml").unwrap();
+        let model_computation = ModelComputation {
+            surml_file: &mut file,
+        };
+
+        let mut input_values = HashMap::new();
+        input_values.insert(String::from("squarefoot"), 1000.0);
+        input_values.insert(String::from("num_floors"), 2.0);
+
+        let output = model_computation.buffered_compute(&mut input_values).unwrap();
+        assert_eq!(output.len(), 1);
+    }
+
+    #[cfg(feature = "torch-tests")]
     #[test]
     fn test_raw_compute_linear_torch() {
         let mut file = SurMlFile::from_file("./model_stash/torch/surml/linear.surml").unwrap();
@@ -211,6 +248,7 @@ mod tests {
         assert_eq!(output.len(), 1);
     }
 
+    #[cfg(feature = "torch-tests")]
     #[test]
     fn test_buffered_compute_linear_torch() {
         let mut file = SurMlFile::from_file("./model_stash/torch/surml/linear.surml").unwrap();
@@ -226,6 +264,7 @@ mod tests {
         assert_eq!(output.len(), 1);
     }
 
+    #[cfg(feature = "tensorflow-tests")]
     #[test]
     fn test_raw_compute_linear_tensorflow() {
         let mut file = SurMlFile::from_file("./model_stash/tensorflow/surml/linear.surml").unwrap();
@@ -243,6 +282,7 @@ mod tests {
         assert_eq!(output.len(), 1);
     }
 
+    #[cfg(feature = "tensorflow-tests")]
     #[test]
     fn test_buffered_compute_linear_tensorflow() {
         let mut file = SurMlFile::from_file("./model_stash/tensorflow/surml/linear.surml").unwrap();

requirements.txt

@@ -1,7 +1,2 @@
-numpy
-skl2onnx
-scikit-learn
-torch
-tf2onnx
-tensorflow
-onnxruntime
+onnxruntime==1.17.3
+numpy==1.26.3

surrealml/engine/__init__.py

@@ -3,6 +3,7 @@
 from surrealml.engine.sklearn import SklearnOnnxAdapter
 from surrealml.engine.torch import TorchOnnxAdapter
 from surrealml.engine.tensorflow import TensorflowOnnxAdapter
+from surrealml.engine.onnx import OnnxAdapter
 
 
 class Engine(Enum):
@@ -14,8 +15,10 @@ class Engine(Enum):
         NATIVE: The native engine which will be native rust and linfa.
         SKLEARN: The sklearn engine which will be sklearn and ONNX
         TENSOFRLOW: The TensorFlow engine which will be TensorFlow and ONNX
+        ONNX: The ONNX engine which bypasses the conversion to ONNX.
     """
     PYTORCH = "pytorch"
     NATIVE = "native"
     SKLEARN = "sklearn"
     TENSORFLOW = "tensorflow"
+    ONNX = "onnx"

surrealml/engine/onnx.py

@@ -0,0 +1,26 @@
+"""
+This file defines the adapter for the ONNX file format. This adapter does not convert anything as the input
+model is already in the ONNX format. It simply saves the model to a file. However, I have added this adapter
+to keep the same structure as the other adapters for different engines (maxwell flitton).
+"""
+from surrealml.engine.utils import create_file_cache_path
+
+
+class OnnxAdapter:
+
+    @staticmethod
+    def save_model_to_onnx(model, inputs) -> str:
+        """
+        Saves a model to an onnx file.
+
+        :param model: the raw ONNX model to directly save
+        :param inputs: the inputs to the model needed to trace the model
+        :return: the path to the cache created with a unique id to prevent collisions.
+        """
+        file_path = create_file_cache_path()
+
+        with open(file_path, "wb") as f:
+            f.write(model.SerializeToString())
+
+        return file_path
+

surrealml/engine/sklearn.py

@@ -1,3 +1,6 @@
+"""
+This file defines the adapter that converts an sklearn model to an onnx model and saves the onnx model to a file.
+"""
 try:
     import skl2onnx
 except ImportError:
@@ -7,6 +10,9 @@
 
 
 class SklearnOnnxAdapter:
+    """
+    Converts and saves sklearn models to onnx format.
+    """
 
     @staticmethod
     def check_dependency() -> None: