Move evaluator registry (#411)

Author: rahul-tuli

src/sparsezoo/evaluation/__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 .registry import *

src/sparsezoo/evaluation/registry.py

@@ -0,0 +1,28 @@
+# 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.
+"""
+Implementation of a registry for evaluation functions
+"""
+
+from sparsezoo.utils.registry import RegistryMixin
+
+
+__all__ = ["EvaluationRegistry"]
+
+
+class EvaluationRegistry(RegistryMixin):
+    """
+    Extends the RegistryMixin to enable registering
+    and loading of evaluation functions.
+    """

src/sparsezoo/evaluation/results.py

@@ -0,0 +1,142 @@
+# 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.
+
+from typing import Any, List, Optional, Union
+
+import numpy
+import yaml
+from pydantic import BaseModel, Field
+
+
+__all__ = [
+    "Metric",
+    "Dataset",
+    "EvalSample",
+    "Evaluation",
+    "Result",
+    "save_result",
+]
+
+
+def prep_for_serialization(
+    data: Union[BaseModel, numpy.ndarray, list]
+) -> Union[BaseModel, list]:
+    """
+    Prepares input data for JSON serialization by converting any numpy array
+    field to a list. For large numpy arrays, this operation will take a while to run.
+
+    :param data: data to that is to be processed before
+        serialization. Nested objects are supported.
+    :return: Pipeline_outputs with potential numpy arrays
+        converted to lists
+    """
+    if isinstance(data, BaseModel):
+        for field_name in data.__fields__.keys():
+            field_value = getattr(data, field_name)
+            if isinstance(field_value, (numpy.ndarray, BaseModel, list)):
+                setattr(
+                    data,
+                    field_name,
+                    prep_for_serialization(field_value),
+                )
+
+    elif isinstance(data, numpy.ndarray):
+        data = data.tolist()
+
+    elif isinstance(data, list):
+        for i, value in enumerate(data):
+            data[i] = prep_for_serialization(value)
+
+    elif isinstance(data, dict):
+        for key, value in data.items():
+            data[key] = prep_for_serialization(value)
+
+    return data
+
+
+class Metric(BaseModel):
+    name: str = Field(description="Name of the metric")
+    value: float = Field(description="Value of the metric")
+
+
+class Dataset(BaseModel):
+    type: Optional[str] = Field(description="Type of dataset")
+    name: str = Field(description="Name of the dataset")
+    config: Any = Field(description="Configuration for the dataset")
+    split: Optional[str] = Field(description="Split of the dataset")
+
+
+class EvalSample(BaseModel):
+    input: Any = Field(description="Sample input to the model")
+    output: Any = Field(description="Sample output from the model")
+
+
+class Evaluation(BaseModel):
+    task: str = Field(
+        description="Name of the evaluation integration "
+        "that the evaluation was performed on"
+    )
+    dataset: Dataset = Field(description="Dataset that the evaluation was performed on")
+    metrics: List[Metric] = Field(description="List of metrics for the evaluation")
+    samples: Optional[List[EvalSample]] = Field(
+        description="List of samples for the evaluation"
+    )
+
+
+class Result(BaseModel):
+    formatted: List[Evaluation] = Field(
+        description="Evaluation result represented in the unified, structured format"
+    )
+    raw: Any = Field(
+        description="Evaluation result represented in the raw format "
+        "(characteristic for the specific evaluation integration)"
+    )
+
+
+def save_result(
+    result: Result,
+    save_path: str,
+    save_format: str = "json",
+):
+    """
+    Saves a list of Evaluation objects to a file in the specified format.
+
+    :param result: Result object to save
+    :param save_path: Path to save the evaluations to.
+    :param save_format: Format to save the evaluations in.
+    :return: The serialized evaluations
+    """
+    # prepare the Result object for serialization
+    result: Result = prep_for_serialization(result)
+    if save_format == "json":
+        _save_to_json(result, save_path)
+    elif save_format == "yaml":
+        _save_to_yaml(result, save_path)
+    else:
+        NotImplementedError("Currently only json and yaml formats are supported")
+
+
+def _save_to_json(result: Result, save_path: str):
+    _save(result.json(), save_path, expected_ext=".json")
+
+
+def _save_to_yaml(result: Result, save_path: str):
+    _save(yaml.dump(result.dict()), save_path, expected_ext=".yaml")
+
+
+def _save(data: str, save_path: str, expected_ext: str):
+    if not save_path.endswith(expected_ext):
+        raise ValueError(f"save_path must end with extension: {expected_ext}")
+    with open(save_path, "w") as f:
+        f.write(data)

tests/sparsezoo/evaluation/__init__.py

@@ -0,0 +1,13 @@
+# 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.

tests/sparsezoo/evaluation/test_registry

@@ -0,0 +1,51 @@
+# 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.
+
+import pytest
+from sparsezoo.evaluation.registry import EvaluationRegistry
+
+
+@pytest.fixture
+def registry_with_foo():
+    class Registry(EvaluationRegistry):
+        pass
+
+    @Registry.register()
+    def foo(*args, **kwargs):
+        return "foo"
+
+    return Registry
+
+
+@pytest.fixture
+def registry_with_buzz():
+    class Registry(EvaluationRegistry):
+        pass
+
+    @Registry.register(name=["buzz", "buzzer"])
+    def buzz(*args, **kwargs):
+        return "buzz"
+
+    return Registry
+
+
+def test_get_foo_from_registry(registry_with_foo):
+    eval_function = registry_with_foo.load_from_registry("foo")
+    assert eval_function() == "foo"
+
+
+def test_get_multiple_buzz_from_registry(registry_with_buzz):
+    eval_function_1 = registry_with_buzz.load_from_registry("buzz")
+    eval_function_2 = registry_with_buzz.load_from_registry("buzzer")
+    assert eval_function_1() == eval_function_2() == "buzz"

tests/sparsezoo/evaluation/test_results.py

@@ -0,0 +1,78 @@
+# 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.
+
+import json
+
+import numpy as np
+import pytest
+import yaml
+
+from sparsezoo.evaluation.results import (
+    Dataset,
+    EvalSample,
+    Evaluation,
+    Metric,
+    Result,
+    save_result,
+)
+
+
+@pytest.fixture()
+def evaluations():
+    return [
+        Evaluation(
+            task="task_1",
+            dataset=Dataset(
+                type="type_1", name="name_1", config="config_1", split="split_1"
+            ),
+            metrics=[Metric(name="metric_name_1", value=1.0)],
+            samples=[EvalSample(input=np.array([[5]]), output=5)],
+        ),
+        Evaluation(
+            task="task_2",
+            dataset=Dataset(
+                type="type_2", name="name_2", config="config_2", split="split_2"
+            ),
+            metrics=[
+                Metric(name="metric_name_2", value=2.0),
+                Metric(name="metric_name_3", value=3.0),
+            ],
+            samples=[
+                EvalSample(input=np.array([[10.0]]), output=10.0),
+                EvalSample(input=np.array([[20.0]]), output=20.0),
+            ],
+        ),
+    ]
+
+
+@pytest.fixture()
+def result(evaluations):
+    return Result(formatted=evaluations, raw="dummy_raw_evaluation")
+
+
+def test_serialize_result_json(tmp_path, result):
+    path_to_file = tmp_path / "result.json"
+    save_result(result=result, save_format="json", save_path=path_to_file.as_posix())
+
+    with open(path_to_file.as_posix(), "r") as f:
+        reloaded_results = json.load(f)
+    assert reloaded_results == result.dict()
+
+
+def test_serialize_result_yaml(tmp_path, result):
+    path_to_file = tmp_path / "result.yaml"
+    save_result(result=result, save_format="yaml", save_path=path_to_file.as_posix())
+    with open(path_to_file.as_posix(), "r") as f:
+        reloaded_results = yaml.safe_load(f)
+    assert reloaded_results == result.dict()