Added support for nightly versions for framework package version (#368)

Framework package versions now list nightly versions of SparseZoo/SparseML/DeepSparse.

Co-authored-by: Mark Kurtz <mark@neuralmagic.com>

Author: kevinaer

src/sparseml/deepsparse/framework/info.py

@@ -140,9 +140,21 @@ def framework_info() -> FrameworkInfo:
     return FrameworkInfo(
         framework=Framework.deepsparse,
         package_versions={
-            "deepsparse": get_version(package_name="deepsparse", raise_on_error=False),
-            "sparsezoo": get_version(package_name="sparsezoo", raise_on_error=False),
-            "sparseml": get_version(package_name="sparseml", raise_on_error=False),
+            "deepsparse": get_version(
+                package_name="deepsparse",
+                raise_on_error=False,
+                alternate_package_names=["deepsparse-nightly"],
+            ),
+            "sparsezoo": get_version(
+                package_name="sparsezoo",
+                raise_on_error=False,
+                alternate_package_names=["sparsezoo-nightly"],
+            ),
+            "sparseml": get_version(
+                package_name="sparseml",
+                raise_on_error=False,
+                alternate_package_names=["sparseml-nightly"],
+            ),
         },
         sparsification=sparsification_info(),
         inference_providers=[cpu_provider],

src/sparseml/keras/framework/info.py

@@ -133,8 +133,16 @@ def framework_info() -> FrameworkInfo:
             "onnx": get_version(package_name="onnx", raise_on_error=False),
             "keras2onnx": get_version(package_name="keras2onnx", raise_on_error=False),
             "tf2onnx": get_version(package_name="tf2onnx", raise_on_error=False),
-            "sparsezoo": get_version(package_name="sparsezoo", raise_on_error=False),
-            "sparseml": get_version(package_name="sparseml", raise_on_error=False),
+            "sparsezoo": get_version(
+                package_name="sparsezoo",
+                raise_on_error=False,
+                alternate_package_names=["sparsezoo-nightly"],
+            ),
+            "sparseml": get_version(
+                package_name="sparseml",
+                raise_on_error=False,
+                alternate_package_names=["sparseml-nightly"],
+            ),
         },
         sparsification=sparsification_info(),
         inference_providers=[cpu_provider, gpu_provider],

src/sparseml/onnx/framework/info.py

@@ -141,8 +141,16 @@ def framework_info() -> FrameworkInfo:
             "onnxruntime": (
                 get_version(package_name="onnxruntime", raise_on_error=False)
             ),
-            "sparsezoo": get_version(package_name="sparsezoo", raise_on_error=False),
-            "sparseml": get_version(package_name="sparseml", raise_on_error=False),
+            "sparsezoo": get_version(
+                package_name="sparsezoo",
+                raise_on_error=False,
+                alternate_package_names=["sparsezoo-nightly"],
+            ),
+            "sparseml": get_version(
+                package_name="sparseml",
+                raise_on_error=False,
+                alternate_package_names=["sparseml-nightly"],
+            ),
         },
         sparsification=sparsification_info(),
         inference_providers=[cpu_provider, gpu_provider],

src/sparseml/pytorch/framework/info.py

@@ -130,8 +130,16 @@ def framework_info() -> FrameworkInfo:
                 get_version(package_name="torchvision", raise_on_error=False)
             ),
             "onnx": get_version(package_name="onnx", raise_on_error=False),
-            "sparsezoo": get_version(package_name="sparsezoo", raise_on_error=False),
-            "sparseml": get_version(package_name="sparseml", raise_on_error=False),
+            "sparsezoo": get_version(
+                package_name="sparsezoo",
+                raise_on_error=False,
+                alternate_package_names=["sparsezoo-nightly"],
+            ),
+            "sparseml": get_version(
+                package_name="sparseml",
+                raise_on_error=False,
+                alternate_package_names=["sparseml-nightly"],
+            ),
         },
         sparsification=sparsification_info(),
         inference_providers=[cpu_provider, gpu_provider],

src/sparseml/tensorflow_v1/framework/info.py

@@ -131,8 +131,16 @@ def framework_info() -> FrameworkInfo:
             ),
             "onnx": get_version(package_name="onnx", raise_on_error=False),
             "tf2onnx": get_version(package_name="tf2onnx", raise_on_error=False),
-            "sparsezoo": get_version(package_name="sparsezoo", raise_on_error=False),
-            "sparseml": get_version(package_name="sparseml", raise_on_error=False),
+            "sparsezoo": get_version(
+                package_name="sparsezoo",
+                raise_on_error=False,
+                alternate_package_names=["sparsezoo-nightly"],
+            ),
+            "sparseml": get_version(
+                package_name="sparseml",
+                raise_on_error=False,
+                alternate_package_names=["sparseml-nightly"],
+            ),
         },
         sparsification=sparsification_info(),
         inference_providers=[cpu_provider, gpu_provider],