[Export Refactor] Feature Branch (#1858)

* initial commit

* respond to PR comments

* [Export Refactor][Image Classification] `create_model` function (#1878)

* initial commit

* looking good, time to cleanup

* Delete src/sparseml/export/helpers.py

* Delete tests/sparseml/export/test_helpers.py

* ready for review

* improve design

* tests pass

* reuse _validate_dataset_num_classes

* [Export Refactor][Image Classification] `create_dummy_input` function (#1880)

* initial commit

* looking good, time to cleanup

* Delete src/sparseml/export/helpers.py

* Delete tests/sparseml/export/test_helpers.py

* ready for review

* improve design

* tests pass

* reuse _validate_dataset_num_classes

* initial commit

* Update src/sparseml/pytorch/image_classification/integration_helper_functions.py

* Update src/sparseml/pytorch/image_classification/integration_helper_functions.py

* ready for review

* Update src/sparseml/export/export.py

* Update src/sparseml/integration_helper_functions.py

* [Export Refactor][Image Classification] `export_model` function (#1883)

* initial commit

* looking good, time to cleanup

* Delete src/sparseml/export/helpers.py

* Delete tests/sparseml/export/test_helpers.py

* ready for review

* improve design

* tests pass

* reuse _validate_dataset_num_classes

* initial commit

* Update src/sparseml/pytorch/image_classification/integration_helper_functions.py

* Update src/sparseml/pytorch/image_classification/integration_helper_functions.py

* ready for review

* Update src/sparseml/export/export.py

* Update src/sparseml/integration_helper_functions.py

* initial commit

* fixes

* ready for review

* nit

* add return

* make export function more general

* [Export Refactor][Image Classification] `apply_optimizations` function (#1884)

* initial commit

* looking good, time to cleanup

* Delete src/sparseml/export/helpers.py

* Delete tests/sparseml/export/test_helpers.py

* ready for review

* improve design

* tests pass

* reuse _validate_dataset_num_classes

* initial commit

* Update src/sparseml/pytorch/image_classification/integration_helper_functions.py

* Update src/sparseml/pytorch/image_classification/integration_helper_functions.py

* ready for review

* Update src/sparseml/export/export.py

* Update src/sparseml/integration_helper_functions.py

* initial commit

* fixes

* ready for review

* nit

* add return

* initial commit

* [Export Refactor][Image Classification] `export_sample_inputs_outputs` function (#1888)

* initial commit

* looking good, time to cleanup

* Delete src/sparseml/export/helpers.py

* Delete tests/sparseml/export/test_helpers.py

* ready for review

* improve design

* tests pass

* reuse _validate_dataset_num_classes

* initial commit

* Update src/sparseml/pytorch/image_classification/integration_helper_functions.py

* Update src/sparseml/pytorch/image_classification/integration_helper_functions.py

* ready for review

* Update src/sparseml/export/export.py

* Update src/sparseml/integration_helper_functions.py

* initial commit

* fixes

* ready for review

* nit

* add return

* initial commit

* initial commit

* PR comments

* beautification

* remove duplicated function

* [Export Refactor][Image Classification] `create_deployment_folder` function (#1889)

* initial commit

* looking good, time to cleanup

* Delete src/sparseml/export/helpers.py

* Delete tests/sparseml/export/test_helpers.py

* ready for review

* improve design

* tests pass

* reuse _validate_dataset_num_classes

* initial commit

* Update src/sparseml/pytorch/image_classification/integration_helper_functions.py

* Update src/sparseml/pytorch/image_classification/integration_helper_functions.py

* ready for review

* Update src/sparseml/export/export.py

* Update src/sparseml/integration_helper_functions.py

* initial commit

* fixes

* ready for review

* nit

* add return

* initial commit

* initial commit

* initial commit

* fix rebase, tests_work

* ready to push

* [Export Refactor][Image Classification] `validate_correctness` function (#1890)

* initial commit

* looking good, time to cleanup

* Delete src/sparseml/export/helpers.py

* Delete tests/sparseml/export/test_helpers.py

* ready for review

* improve design

* tests pass

* reuse _validate_dataset_num_classes

* initial commit

* Update src/sparseml/pytorch/image_classification/integration_helper_functions.py

* Update src/sparseml/pytorch/image_classification/integration_helper_functions.py

* ready for review

* Update src/sparseml/export/export.py

* Update src/sparseml/integration_helper_functions.py

* initial commit

* fixes

* ready for review

* nit

* add return

* initial commit

* initial commit

* initial commit

* initial commit

* Delete tests/sparseml/test_integration_helper_functions.py

* ready to merge

* [Export Refactor] End to end testing (#1898)

* initial commit

* looking good, time to cleanup

* Delete src/sparseml/export/helpers.py

* Delete tests/sparseml/export/test_helpers.py

* ready for review

* improve design

* tests pass

* reuse _validate_dataset_num_classes

* initial commit

* Update src/sparseml/pytorch/image_classification/integration_helper_functions.py

* Update src/sparseml/pytorch/image_classification/integration_helper_functions.py

* ready for review

* Update src/sparseml/export/export.py

* Update src/sparseml/integration_helper_functions.py

* initial commit

* fixes

* ready for review

* nit

* add return

* initial commit

* initial commit

* initial commit

* initial commit

* Delete tests/sparseml/test_integration_helper_functions.py

* ready to merge

* add structure validator

* ready for review

* Delete tests/sparseml/export/model.onnx

* Delete tests/sparseml/export/image_classification/model.onnx

* Delete tests/sparseml/export/image_classification/conftest.py

* PR comments

* remove onnx

* [Export Refactor] Prepare the module to be more general (before including `transformers`) (#1908)

* adapt the export script to handle transformers

* Update src/sparseml/pytorch/image_classification/integration_helper_functions.py

* Delete tests/sparseml/export/transformers/__init__.py

* Delete tests/sparseml/export/transformers/test_generative_transformers.py

* Delete tests/sparseml/export/transformers/test_transformers.py

* Update src/sparseml/export/export.py

Co-authored-by: Benjamin Fineran <bfineran@users.noreply.github.com>

* addressing review comments

* [Export Refactor] Export `transformers` (#1909)

* cleanup

* Delete src/sparseml/transformers/integration_helper_functions_generative.py

* Delete src/sparseml/transformers/utils/optimizations.py

* Delete tests/sparseml/export/transformers/test_generative_transformers.py

* Delete tests/sparseml/transformers/test_integration_helper_functions_generative.py

* addressing PR reviews

* [Export Refactor] Export generative transformers(#1910)

* make tests green, remove using task to resolve the integration type

* fix all the tests after the merge, make integration resolution independent of the task name

* fold generative transformers into transformer helper functions

* complete tests for export_data.py

* Update src/sparseml/export/export.py

* add tests that confirms that kv cache injection has been added

* move applying optimizations into integration helper functions

---------

Co-authored-by: Benjamin Fineran <bfineran@users.noreply.github.com>

* [Export Refactor][Transformers] Enable loading SparseModels (#1921)

* initial commit

* adressing review comments

* Fix the tests

* fix tests with help from sara

* [Export][Transformers] Enable loading `text-generation` datasets (#1938)

* add suport for past_key_values in sample-outputs

* [Export][Transformers] Implementation of correctness validation (#1935)

* fix tests with help from sara

* Update src/sparseml/transformers/utils/initializers.py

* swap sparsezoo validator for custom one (top k match)

* add more informative error message

* add correctness validation for LLMs

* remove past_key_values from outputs

* remove past_key_values from outputs (2)

* small note comment for the future

* tests fixed

* fix test

* [Export refactor] final manual testing fixes (#1948)

* [Export refactor] final manual testing fixes

* review

---------

Co-authored-by: Benjamin Fineran <bfineran@users.noreply.github.com>

.github/workflows/test-check.yaml

@@ -18,6 +18,7 @@ jobs:
       deepsparse: ${{ steps.deepsparse-check.outputs.output }}
       onnx: ${{ steps.onnx-check.outputs.output }}
       pytorch: ${{ steps.pytorch-check.outputs.output }}
+      export: ${{ steps.export-check.outputs.output }}
       transformers: ${{ steps.transformers-check.outputs.output }}
     steps:
       - uses: actions/checkout@v2
@@ -54,6 +55,10 @@ jobs:
           ((git diff --name-only origin/main HEAD | grep -E "[src|tests]/sparseml/pytorch|setup.py|.github")
           || (echo $GITHUB_REF | grep -E "refs/heads/[release/|main]"))
           && echo "::set-output name=output::1" || echo "::set-output name=output::0"
+      - name: "Checking if sparseml.export was changed"
+        id: export-check
+        run: >
+          ((git diff --name-only origin/main HEAD | grep -E "[src|tests]/sparseml/export|setup.py|.github")
       - name: "Checking if sparseml.transformers was changed"
         id: transformers-check
         run: >
@@ -228,3 +233,28 @@ jobs:
         run: pip3 install .[dev,torch,transformers]
       - name: "🔬 Running transformers tests"
         run: make test TARGETS=transformers
+  export-tests:
+    runs-on: ubuntu-22.04
+    env:
+      SPARSEZOO_TEST_MODE: "true"
+    needs: test-setup
+    if: ${{needs.test-setup.outputs.export == 1}}
+    steps:
+      - uses: actions/setup-python@v4
+        with:
+          python-version: '3.11'
+      - uses: actions/checkout@v2
+      - uses: actions/checkout@v2
+        with:
+          repository: "neuralmagic/sparsezoo"
+          path: "sparsezoo"
+          ref: ${{needs.test-setup.outputs.branch}}
+      - name: "⚙️ Install sparsezoo dependencies"
+        run: pip3 install -U pip && pip3 install setuptools sparsezoo/
+      - name: "Clean sparsezoo directory"
+        run: rm -r sparsezoo/
+      - name: "⚙️ Install dependencies"
+        run: pip3 install .[dev,torch,transformers,torchvision,onnxruntime]
+      - name: "🔬 Running export tests"
+        run: make test TARGETS=export
+

Makefile

@@ -9,7 +9,7 @@ MDCHECKFILES := CODE_OF_CONDUCT.md CONTRIBUTING.md DEVELOPING.md README.md
 SPARSEZOO_TEST_MODE := "true"
 
 BUILD_ARGS :=  # set nightly to build nightly release
-TARGETS := ""  # targets for running pytests: deepsparse,keras,onnx,pytorch,pytorch_models,pytorch_datasets,tensorflow_v1,tensorflow_v1_models,tensorflow_v1_datasets
+TARGETS := ""  # targets for running pytests: deepsparse,keras,onnx,pytorch,pytorch_models,export,pytorch_datasets,tensorflow_v1,tensorflow_v1_models,tensorflow_v1_datasets
 PYTEST_ARGS ?= ""
 PYTEST_INTEG_ARGS ?= ""
 ifneq ($(findstring deepsparse,$(TARGETS)),deepsparse)
@@ -18,6 +18,9 @@ endif
 ifneq ($(findstring transformers,$(TARGETS)),transformers)
     PYTEST_ARGS := $(PYTEST_ARGS) --ignore tests/sparseml/transformers
 endif
+ifneq ($(findstring export,$(TARGETS)),export)
+    PYTEST_ARGS := $(PYTEST_ARGS) --ignore tests/sparseml/export
+endif
 ifneq ($(findstring keras,$(TARGETS)),keras)
     PYTEST_ARGS := $(PYTEST_ARGS) --ignore tests/sparseml/keras
 endif

setup.py

@@ -170,6 +170,8 @@ def _setup_extras() -> Dict:
 def _setup_entry_points() -> Dict:
     entry_points = {
         "console_scripts": [
+            # export
+            "sparseml.export=sparseml.export.export:main",
             # sparsification
             "sparseml.framework=sparseml.framework.info:_main",
             "sparseml.sparsification=sparseml.sparsification.info:_main",

src/sparseml/core/session.py

@@ -100,7 +100,7 @@ def pre_initialize_structure(
         This will run the pre-initialize structure method for each modifier in the
         session's lifecycle. This will also set the session's state to the
         pre-initialized state. Takes care of cases when the model(s) structure
-        has been previosuly modified by a modifier.
+        has been previously modified by a modifier.
 
         :param model: the model to pre-initialize the structure for
         :param recipe: the recipe to use for the sparsification, can be a path to a

src/sparseml/export/__init__.py

@@ -0,0 +1,17 @@
+# Copyright (c) 2021 - present / Neuralmagic, Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# flake8: noqa
+
+from .export import *