diff --git a/.github/workflows/surrealml_core_onnx_test.yml b/.github/workflows/surrealml_core_onnx_test.yml new file mode 100644 index 0000000..03ab696 --- /dev/null +++ b/.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 diff --git a/.github/workflows/surrealml_core_tensorflow_test.yml b/.github/workflows/surrealml_core_tensorflow_test.yml new file mode 100644 index 0000000..355e754 --- /dev/null +++ b/.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 diff --git a/.github/workflows/surrealml_core_test.yml b/.github/workflows/surrealml_core_test.yml index 178bcf8..be6e3d9 100644 --- a/.github/workflows/surrealml_core_test.yml +++ b/.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 diff --git a/.github/workflows/surrealml_core_torch_test.yml b/.github/workflows/surrealml_core_torch_test.yml new file mode 100644 index 0000000..01b679f --- /dev/null +++ b/.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 diff --git a/README.md b/README.md index 1300ba4..bf8211e 100644 --- a/README.md +++ b/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 diff --git a/modules/core/Cargo.toml b/modules/core/Cargo.toml index 210ed78..82cb07e 100644 --- a/modules/core/Cargo.toml +++ b/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" diff --git a/modules/core/model_stash/onnx/onnx/linear.onnx b/modules/core/model_stash/onnx/onnx/linear.onnx new file mode 100644 index 0000000..bc1dbf6 Binary files /dev/null and b/modules/core/model_stash/onnx/onnx/linear.onnx differ diff --git a/modules/core/model_stash/onnx/surml/linear.surml b/modules/core/model_stash/onnx/surml/linear.surml new file mode 100644 index 0000000..f092b50 Binary files /dev/null and b/modules/core/model_stash/onnx/surml/linear.surml differ diff --git a/modules/core/model_stash/tensorflow/surml/linear.surml b/modules/core/model_stash/tensorflow/surml/linear.surml index 1670153..87e7cfb 100644 Binary files a/modules/core/model_stash/tensorflow/surml/linear.surml and b/modules/core/model_stash/tensorflow/surml/linear.surml differ diff --git a/modules/core/model_stash/torch/surml/linear.surml b/modules/core/model_stash/torch/surml/linear.surml index ec68ffb..1c00cd9 100644 Binary files a/modules/core/model_stash/torch/surml/linear.surml and b/modules/core/model_stash/torch/surml/linear.surml differ diff --git a/modules/core/src/execution/compute.rs b/modules/core/src/execution/compute.rs index f7fb955..6c796a3 100644 --- a/modules/core/src/execution/compute.rs +++ b/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(); diff --git a/requirements.txt b/requirements.txt index ad5846f..b840da4 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,7 +1,2 @@ -numpy -skl2onnx -scikit-learn -torch -tf2onnx -tensorflow -onnxruntime +onnxruntime==1.17.3 +numpy==1.26.3 \ No newline at end of file diff --git a/surrealml/engine/__init__.py b/surrealml/engine/__init__.py index 9b3f16c..05708ae 100644 --- a/surrealml/engine/__init__.py +++ b/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" diff --git a/surrealml/engine/onnx.py b/surrealml/engine/onnx.py new file mode 100644 index 0000000..d904672 --- /dev/null +++ b/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 + diff --git a/surrealml/engine/sklearn.py b/surrealml/engine/sklearn.py index 3f9de69..089bb57 100644 --- a/surrealml/engine/sklearn.py +++ b/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: diff --git a/surrealml/model_templates/onnx/__init__.py b/surrealml/model_templates/onnx/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/surrealml/model_templates/onnx/onnx_linear.py b/surrealml/model_templates/onnx/onnx_linear.py new file mode 100644 index 0000000..5c10edb --- /dev/null +++ b/surrealml/model_templates/onnx/onnx_linear.py @@ -0,0 +1,44 @@ +""" +Trains a linear regression model using sklearn but keeping the ONNX format for the raw onnx support. +""" +from sklearn.linear_model import LinearRegression + +from surrealml.model_templates.datasets.house_linear import HOUSE_LINEAR + + +def train_model(): + """ + Trains a linear regression model using sklearn and returns the raw ONNX format. + This is a basic model that can be used for testing. + """ + import skl2onnx + model = LinearRegression() + model.fit(HOUSE_LINEAR["inputs"], HOUSE_LINEAR["outputs"]) + return skl2onnx.to_onnx(model, HOUSE_LINEAR["inputs"]) + + +def export_model_onnx(model): + """ + Exports the model to ONNX format. + + :param model: the model to export. + :return: the path to the exported model. + """ + return model + + +def export_model_surml(model): + """ + Exports the model to SURML format. + + :param model: the model to export. + :return: the path to the exported model. + """ + from surrealml import SurMlFile, Engine + file = SurMlFile(model=model, name="linear", inputs=HOUSE_LINEAR["inputs"], engine=Engine.ONNX) + file.add_column("squarefoot") + file.add_column("num_floors") + file.add_normaliser("squarefoot", "z_score", HOUSE_LINEAR["squarefoot"].mean(), HOUSE_LINEAR["squarefoot"].std()) + file.add_normaliser("num_floors", "z_score", HOUSE_LINEAR["num_floors"].mean(), HOUSE_LINEAR["num_floors"].std()) + file.add_output("house_price", "z_score", HOUSE_LINEAR["outputs"].mean(), HOUSE_LINEAR["outputs"].std()) + return file diff --git a/surrealml/surml_file.py b/surrealml/surml_file.py index d8e5a5b..fe89676 100644 --- a/surrealml/surml_file.py +++ b/surrealml/surml_file.py @@ -3,7 +3,7 @@ """ from typing import Optional -from surrealml.engine import Engine, SklearnOnnxAdapter, TorchOnnxAdapter, TensorflowOnnxAdapter +from surrealml.engine import Engine, SklearnOnnxAdapter, TorchOnnxAdapter, TensorflowOnnxAdapter, OnnxAdapter from surrealml.rust_adapter import RustAdapter @@ -54,6 +54,13 @@ def _cache_model(self) -> Optional[str]: model=self.model, inputs=self.inputs ) + # Below doesn't really convert to ONNX, but I want to keep the same structure as the other engines + # (maxwell flitton) + elif self.engine == Engine.ONNX: + raw_file_path: str = OnnxAdapter.save_model_to_onnx( + model=self.model, + inputs=self.inputs + ) else: raise ValueError(f"Engine {self.engine} not supported") return RustAdapter.pass_raw_model_into_rust(raw_file_path) diff --git a/tests/model_builder/__init__.py b/tests/model_builder/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/model_builder/onnx_assets.py b/tests/model_builder/onnx_assets.py new file mode 100644 index 0000000..0a7d124 --- /dev/null +++ b/tests/model_builder/onnx_assets.py @@ -0,0 +1,45 @@ +from tests.model_builder.utils import install_package +# install_package("torch==2.1.2") +install_package("onnx==1.16.0") +install_package("scikit-learn==1.4.0") # use sklearn to train the model +install_package("skl2onnx==1.16.0") # use skl2onnx to convert the model to onnx to test the raw onnx loading +import os +import onnx + +from surrealml.model_templates.onnx.onnx_linear import train_model as linear_onnx_train_model +from surrealml.model_templates.onnx.onnx_linear import export_model_onnx as linear_onnx_export_model_onnx +from surrealml.model_templates.onnx.onnx_linear import export_model_surml as linear_onnx_export_model_surml + +from tests.model_builder.utils import delete_directory, create_directory, MODEL_STASH_DIRECTORY + + +# create the model stash directory if it does not exist +create_directory(dir_path=MODEL_STASH_DIRECTORY) + +# defining directories for onnx +onnx_stash_directory = os.path.join(MODEL_STASH_DIRECTORY, "onnx") +onnx_surml_stash_directory = os.path.join(onnx_stash_directory, "surml") +onnx_onnx_stash_directory = os.path.join(onnx_stash_directory, "onnx") + +# delete the directories if they exist +delete_directory(dir_path=onnx_stash_directory) +delete_directory(dir_path=onnx_surml_stash_directory) +delete_directory(dir_path=onnx_onnx_stash_directory) + +# create directories for the onnx models +create_directory(dir_path=onnx_stash_directory) +create_directory(dir_path=onnx_surml_stash_directory) +create_directory(dir_path=onnx_onnx_stash_directory) + +# train and stash onnx models +onnx_linear_model = linear_onnx_train_model() +onnx_linear_surml_file = linear_onnx_export_model_surml(onnx_linear_model) +onnx_linear_onnx_file = linear_onnx_export_model_onnx(onnx_linear_model) + +onnx_linear_surml_file.save( + path=str(os.path.join(onnx_surml_stash_directory, "linear.surml")) +) +onnx.save( + onnx_linear_onnx_file, + os.path.join(onnx_onnx_stash_directory, "linear.onnx") +) diff --git a/tests/model_builder/sklearn_assets.py b/tests/model_builder/sklearn_assets.py new file mode 100644 index 0000000..5607d6e --- /dev/null +++ b/tests/model_builder/sklearn_assets.py @@ -0,0 +1,44 @@ +""" +This file trains and saves the sklearn linear model to the model stash directory for the core to test against +""" +from tests.model_builder.utils import install_package +install_package("skl2onnx==1.16.0") +install_package("scikit-learn==1.4.0") +import os + +import onnx + +from surrealml.model_templates.sklearn.sklearn_linear import export_model_onnx as linear_sklearn_export_model_onnx +from surrealml.model_templates.sklearn.sklearn_linear import export_model_surml as linear_sklearn_export_model_surml +from surrealml.model_templates.sklearn.sklearn_linear import train_model as linear_sklearn_train_model +from tests.model_builder.utils import delete_directory, create_directory, MODEL_STASH_DIRECTORY + +sklearn_stash_directory = os.path.join(MODEL_STASH_DIRECTORY, "sklearn") +sklearn_surml_stash_directory = os.path.join(sklearn_stash_directory, "surml") +sklearn_onnx_stash_directory = os.path.join(sklearn_stash_directory, "onnx") + +# create the model stash directory if it does not exist +create_directory(dir_path=MODEL_STASH_DIRECTORY) + +# delete the directories if they exist +delete_directory(dir_path=sklearn_stash_directory) +delete_directory(dir_path=sklearn_surml_stash_directory) +delete_directory(dir_path=sklearn_onnx_stash_directory) + +# create directories for the sklearn models +create_directory(sklearn_stash_directory) +create_directory(sklearn_surml_stash_directory) +create_directory(sklearn_onnx_stash_directory) + +# train and stash sklearn models +sklearn_linear_model = linear_sklearn_train_model() +sklearn_linear_surml_file = linear_sklearn_export_model_surml(sklearn_linear_model) +sklearn_linear_onnx_file = linear_sklearn_export_model_onnx(sklearn_linear_model) + +sklearn_linear_surml_file.save( + path=str(os.path.join(sklearn_surml_stash_directory, "linear.surml")) +) +onnx.save( + sklearn_linear_onnx_file, + os.path.join(sklearn_onnx_stash_directory, "linear.onnx") +) diff --git a/tests/model_builder/tensorflow_assets.py b/tests/model_builder/tensorflow_assets.py new file mode 100644 index 0000000..3b30e7b --- /dev/null +++ b/tests/model_builder/tensorflow_assets.py @@ -0,0 +1,37 @@ +from tests.model_builder.utils import install_package +install_package("tf2onnx==1.16.1") +install_package("tensorflow==2.16.1") +import os + +from surrealml.model_templates.tensorflow.tensorflow_linear import train_model as linear_tensorflow_train_model +from surrealml.model_templates.tensorflow.tensorflow_linear import export_model_onnx as linear_tensorflow_export_model_onnx +from surrealml.model_templates.tensorflow.tensorflow_linear import export_model_surml as linear_tensorflow_export_model_surml + +from tests.model_builder.utils import delete_directory, create_directory, MODEL_STASH_DIRECTORY + + +# create the model stash directory if it does not exist +create_directory(dir_path=MODEL_STASH_DIRECTORY) + +tensorflow_stash_directory = os.path.join(MODEL_STASH_DIRECTORY, "tensorflow") +tensorflow_surml_stash_directory = os.path.join(tensorflow_stash_directory, "surml") +tensorflow_onnx_stash_directory = os.path.join(tensorflow_stash_directory, "onnx") + +# delete the directories if they exist +delete_directory(dir_path=tensorflow_stash_directory) +delete_directory(dir_path=tensorflow_surml_stash_directory) +delete_directory(dir_path=tensorflow_onnx_stash_directory) + +# create directories for the tensorflow models +os.mkdir(tensorflow_stash_directory) +os.mkdir(tensorflow_surml_stash_directory) +os.mkdir(tensorflow_onnx_stash_directory) + +# train and stash tensorflow models +tensorflow_linear_model = linear_tensorflow_train_model() +tensorflow_linear_surml_file = linear_tensorflow_export_model_surml(tensorflow_linear_model) +tensorflow_linear_onnx_file = linear_tensorflow_export_model_onnx(tensorflow_linear_model) + +tensorflow_linear_surml_file.save( + path=str(os.path.join(tensorflow_surml_stash_directory, "linear.surml")) +) \ No newline at end of file diff --git a/tests/model_builder/torch_assets.py b/tests/model_builder/torch_assets.py new file mode 100644 index 0000000..5bf1949 --- /dev/null +++ b/tests/model_builder/torch_assets.py @@ -0,0 +1,40 @@ +""" +This file trains and saves the torch linear model to the model stash directory for the core to test against +""" +from tests.model_builder.utils import install_package +install_package("torch==2.1.2") +install_package("onnx==1.16.0") +import os + +from surrealml.model_templates.torch.torch_linear import train_model as linear_torch_train_model +from surrealml.model_templates.torch.torch_linear import export_model_onnx as linear_torch_export_model_onnx +from surrealml.model_templates.torch.torch_linear import export_model_surml as linear_torch_export_model_surml + +from tests.model_builder.utils import delete_directory, create_directory, MODEL_STASH_DIRECTORY + + +# create the model stash directory if it does not exist +create_directory(dir_path=MODEL_STASH_DIRECTORY) + +torch_stash_directory = os.path.join(MODEL_STASH_DIRECTORY, "torch") +torch_surml_stash_directory = os.path.join(torch_stash_directory, "surml") +torch_onnx_stash_directory = os.path.join(torch_stash_directory, "onnx") + +# delete the directories if they exist +delete_directory(dir_path=torch_stash_directory) +delete_directory(dir_path=torch_surml_stash_directory) +delete_directory(dir_path=torch_onnx_stash_directory) + +# create directories for the torch models +create_directory(torch_stash_directory) +create_directory(torch_surml_stash_directory) +create_directory(torch_onnx_stash_directory) + +# train and stash torch models +torch_linear_model, x = linear_torch_train_model() +torch_linear_surml_file = linear_torch_export_model_surml(torch_linear_model) +torch_linear_onnx_file = linear_torch_export_model_onnx(torch_linear_model) + +torch_linear_surml_file.save( + path=str(os.path.join(torch_surml_stash_directory, "linear.surml")) +) diff --git a/tests/model_builder/utils.py b/tests/model_builder/utils.py new file mode 100644 index 0000000..f3d8cc7 --- /dev/null +++ b/tests/model_builder/utils.py @@ -0,0 +1,77 @@ +import os +import shutil +import subprocess +import sys + +import pkg_resources + +SCRIPT_PATH = os.path.abspath(__file__) +SCRIPT_DIRECTORY = os.path.dirname(SCRIPT_PATH) + +TESTS_DIRECTORY = os.path.join(SCRIPT_DIRECTORY, "..") +MAIN_DIRECTORY = os.path.join(TESTS_DIRECTORY, "..") +CORE_DIRECTORY = os.path.join(MAIN_DIRECTORY, "modules", "core") +MODEL_STASH_DIRECTORY = os.path.join(CORE_DIRECTORY, "model_stash") + + +def install_package(package_name): + try: + # Check if the package is installed + pkg_resources.require(package_name) + print(f"{package_name} is already installed.") + except pkg_resources.DistributionNotFound: + # If not installed, install the package using pip + print(f"{package_name} not found, installing...") + subprocess.check_call([sys.executable, "-m", "pip", "install", package_name]) + print(f"{package_name} has been installed.") + + +def delete_directory(dir_path: os.path) -> None: + """ + Checks to see if a directory exists and deletes it if it does. + + :param dir_path: the path to the directory. + """ + if os.path.exists(dir_path): + shutil.rmtree(dir_path) + print(f"Directory '{dir_path}' has been deleted.") + else: + print(f"Directory '{dir_path}' does not exist.") + + +def delete_file(file_path: os.path) -> None: + """ + Checks to see if a file exists and deletes it if it does. + + :param file_path: the path to the file. + """ + if os.path.isfile(file_path): + os.remove(file_path) + print(f"File '{file_path}' has been deleted.") + else: + print(f"File '{file_path}' does not exist.") + + +def write_file(file_path: os.path, model, file_name) -> None: + """ + Writes a file to the specified path. + + :param file_path: the path to write the file to. + :param model: the model to write to the file. + :param file_name: the name of the file to write. + """ + with open(os.path.join(file_path, file_name), "wb") as f: + f.write(model) + + +def create_directory(dir_path: os.path) -> None: + """ + Checks to see if a directory exists and creates it if it does not. + + :param dir_path: the path to the directory. + """ + if not os.path.exists(dir_path): + os.mkdir(dir_path) + print(f"Directory '{dir_path}' has been created.") + else: + print(f"Directory '{dir_path}' already exists.") diff --git a/tests/scripts/build_assets.py b/tests/scripts/build_assets.py index 15199f3..0b54617 100644 --- a/tests/scripts/build_assets.py +++ b/tests/scripts/build_assets.py @@ -21,6 +21,10 @@ from surrealml.model_templates.sklearn.sklearn_linear import export_model_onnx as linear_sklearn_export_model_onnx from surrealml.model_templates.sklearn.sklearn_linear import export_model_surml as linear_sklearn_export_model_surml +from surrealml.model_templates.onnx.onnx_linear import train_model as linear_onnx_train_model +from surrealml.model_templates.onnx.onnx_linear import export_model_onnx as linear_onnx_export_model_onnx +from surrealml.model_templates.onnx.onnx_linear import export_model_surml as linear_onnx_export_model_surml + from surrealml.model_templates.torch.torch_linear import train_model as linear_torch_train_model from surrealml.model_templates.torch.torch_linear import export_model_onnx as linear_torch_export_model_onnx from surrealml.model_templates.torch.torch_linear import export_model_surml as linear_torch_export_model_surml @@ -71,14 +75,22 @@ def write_file(file_path: os.path, model, file_name) -> None: model_stash_directory = os.path.join(core_directory, "model_stash") +# defining directories for sklearn sklearn_stash_directory = os.path.join(model_stash_directory, "sklearn") sklearn_surml_stash_directory = os.path.join(sklearn_stash_directory, "surml") sklearn_onnx_stash_directory = os.path.join(sklearn_stash_directory, "onnx") +# defining directories for onnx +onnx_stash_directory = os.path.join(model_stash_directory, "onnx") +onnx_surml_stash_directory = os.path.join(onnx_stash_directory, "surml") +onnx_onnx_stash_directory = os.path.join(onnx_stash_directory, "onnx") + +# defining directories for torch torch_stash_directory = os.path.join(model_stash_directory, "torch") torch_surml_stash_directory = os.path.join(torch_stash_directory, "surml") torch_onnx_stash_directory = os.path.join(torch_stash_directory, "onnx") +# defining directories for tensorflow tensorflow_stash_directory = os.path.join(model_stash_directory, "tensorflow") tensorflow_surml_stash_directory = os.path.join(tensorflow_stash_directory, "surml") tensorflow_onnx_stash_directory = os.path.join(tensorflow_stash_directory, "onnx") @@ -95,10 +107,15 @@ def main(): os.mkdir(model_stash_directory) + # create directories for the different model types os.mkdir(sklearn_stash_directory) os.mkdir(sklearn_surml_stash_directory) os.mkdir(sklearn_onnx_stash_directory) + os.mkdir(onnx_stash_directory) + os.mkdir(onnx_surml_stash_directory) + os.mkdir(onnx_onnx_stash_directory) + os.mkdir(torch_stash_directory) os.mkdir(torch_surml_stash_directory) os.mkdir(torch_onnx_stash_directory) @@ -120,6 +137,19 @@ def main(): os.path.join(sklearn_onnx_stash_directory, "linear.onnx") ) + # train and stash onnx models + onnx_linear_model = linear_onnx_train_model() + onnx_linear_surml_file = linear_onnx_export_model_surml(onnx_linear_model) + onnx_linear_onnx_file = linear_onnx_export_model_onnx(onnx_linear_model) + + onnx_linear_surml_file.save( + path=str(os.path.join(onnx_surml_stash_directory, "linear.surml")) + ) + onnx.save( + onnx_linear_onnx_file, + os.path.join(onnx_onnx_stash_directory, "linear.onnx") + ) + # train and stash torch models torch_linear_model, x = linear_torch_train_model() torch_linear_surml_file = linear_torch_export_model_surml(torch_linear_model)