use sparsezoo mixin

Author: dsikka

src/deepsparse/v2/operators/registry.py

@@ -12,67 +12,40 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import List, Optional, Type
+from typing import Type
 
 from deepsparse.v2.task import SupportedTasks, dynamic_import_task
+from sparsezoo.utils.registry import (
+    RegistryMixin,
+    get_from_registry,
+    register,
+    registered_names,
+)
 
 
-_REGISTERED_OPERATORS = {}
-
 __all__ = ["OperatorRegistry"]
 
 
-class OperatorRegistry:
-    def register(task: str, task_aliases: Optional[List[str]] = None):
+class OperatorRegistry(RegistryMixin):
+    """
+    Register operators with given task name(s). Leverages the RegistryMixin
+    functionality.
+    """
+
+    @classmethod
+    def register_value(cls, operator, name):
         from deepsparse.v2.operators import Operator
 
-        """
-        Decorator to register an operator with its task name and its aliases. The
-        registered names can be used to load the operator through `Operator.create()`
+        if not isinstance(name, list):
+            name = [name]
 
-        Multiple operators may not have the same task name. An error will
-        be raised if two different operators attempt to register the same task name
+        for task_name in name:
+            register(Operator, operator, task_name, require_subclass=True)
 
-        :param task: main task name of this operator
-        :param task_aliases: list of extra task names that may be used to reference
-            this operator. Default is None
-        """
-        task_names = [task]
-        if task_aliases:
-            task_names.extend(task_aliases)
-
-        task_names = [task_name.lower().replace("-", "_") for task_name in task_names]
-
-        def _register_task(task_name, operator):
-            if task_name in _REGISTERED_OPERATORS and (
-                operator is not _REGISTERED_OPERATORS[task_name]
-            ):
-                raise RuntimeError(
-                    f"task {task_name} already registered by OperatorRegistry.register "
-                    f"attempting to register operator: {operator}, but"
-                    f"operator: {_REGISTERED_OPERATORS[task_name]}, already registered"
-                )
-            _REGISTERED_OPERATORS[task_name] = operator
-
-        def _register_operator(operator: Operator):
-            if not issubclass(operator, Operator):
-                raise RuntimeError(
-                    f"Attempting to register operator {operator}. "
-                    f"Registered operators must inherit from {Operator}"
-                )
-            for task_name in task_names:
-                _register_task(task_name, operator)
-
-            # set task and task_aliases as class level property
-            operator.task = task
-            operator.task_aliases = task_aliases
-
-            return operator
-
-        return _register_operator
-
-    @staticmethod
-    def get_task_constructor(task: str) -> Type["Operator"]:
+        return operator
+
+    @classmethod
+    def get_task_constructor(cls, task: str) -> Type["Operator"]:  # noqa: F821
         """
         This function retrieves the class previously registered via
         `OperatorRegistry.register` for `task`.
@@ -86,6 +59,8 @@ def get_task_constructor(task: str) -> Type["Operator"]:
         :return: The class registered to `task`
         :raises ValueError: if `task` was not registered via `OperatorRegistry.register`
         """
+        from deepsparse.v2.operators import Operator
+
         if task.startswith("import:"):
             # dynamically import the task from a file
             task = dynamic_import_task(module_or_path=task.replace("import:", ""))
@@ -95,14 +70,7 @@ def get_task_constructor(task: str) -> Type["Operator"]:
         else:
             task = task.lower().replace("-", "_")
 
+        tasks = registered_names(Operator)
         # step needed to import relevant files required to load the operator
-        SupportedTasks.check_register_task(task, _REGISTERED_OPERATORS.keys())
-
-        if task not in _REGISTERED_OPERATORS:
-            raise ValueError(
-                f"Unknown Operator task {task}. Operators tasks should be "
-                "must be declared with the OperatorRegistry.register decorator. "
-                f"Currently registered operators: {list(_REGISTERED_OPERATORS.keys())}"
-            )
-
-        return _REGISTERED_OPERATORS[task]
+        SupportedTasks.check_register_task(task, tasks)
+        return get_from_registry(Operator, task)

src/deepsparse/v2/text_generation/pipeline.py

@@ -45,7 +45,7 @@
 _LOGGER = logging.getLogger(__name__)
 
 
-@OperatorRegistry.register(task="text_generation")
+@OperatorRegistry.register(name="text_generation")
 class TextGenerationPipeline(Pipeline):
     def __init__(
         self,