remove ClassVar

Author: yarden-sony

model_compression_toolkit/target_platform_capabilities/schema/schema_compatability.py

@@ -46,6 +46,18 @@ def _schema_v1_to_v2(tpc: schema_v1.TargetPlatformCapabilities) -> schema_v2.Tar
                                              tpc_platform_type=tpc.tpc_platform_type,
                                              add_metadata=tpc.add_metadata)
 
+def get_conversion_map() -> dict:
+    """
+    Retrieves the schema conversion map.
+    :return: A dictionary where:
+        - Keys representing supported source schema versions.
+        - Values: Callable functions that take tpc in one schema version and return it in the next (higher) version
+    """
+    conversion_map = {
+        1: _schema_v1_to_v2,
+    }
+    return conversion_map
+
 
 def tpc_to_current_schema_version(tpc: Union[all_tpc_types]) -> schema.TargetPlatformCapabilities:
     """
@@ -58,10 +70,8 @@ def tpc_to_current_schema_version(tpc: Union[all_tpc_types]) -> schema.TargetPla
     :param tpc: TargetPlatformCapabilities of some schema version
     :return: TargetPlatformCapabilities with the current MCT schema version
     """
-    conversion_map = {
-        schema_v1.TargetPlatformCapabilities.SCHEMA_VERSION: _schema_v1_to_v2,
-    }
-    while tpc.SCHEMA_VERSION < schema.TargetPlatformCapabilities.SCHEMA_VERSION:
+    conversion_map = get_conversion_map()
+    while not isinstance(tpc, schema.TargetPlatformCapabilities):
         if tpc.SCHEMA_VERSION not in conversion_map:
             raise KeyError(f"TPC using schema version {tpc.SCHEMA_VERSION} which is not in schemas conversion map. "
                            f"Make sure the schema version is supported, or add it in case it's a new schema version")

model_compression_toolkit/target_platform_capabilities/schema/v1.py

@@ -14,7 +14,7 @@
 # ==============================================================================
 import pprint
 from enum import Enum
-from typing import Dict, Any, Union, Tuple, List, Optional, Literal, Annotated, ClassVar
+from typing import Dict, Any, Union, Tuple, List, Optional, Literal, Annotated
 
 from pydantic import BaseModel, Field, root_validator, validator, PositiveInt
 
@@ -647,7 +647,7 @@ class TargetPlatformCapabilities(BaseModel):
     name: Optional[str] = "default_tpc"
     is_simd_padding: bool = False
 
-    SCHEMA_VERSION: ClassVar[int] = 1
+    SCHEMA_VERSION: int = 1
 
     class Config:
         frozen = True

model_compression_toolkit/target_platform_capabilities/schema/v2.py

@@ -14,7 +14,7 @@
 # ==============================================================================
 import pprint
 from enum import Enum
-from typing import Dict, Any, Tuple, Optional, ClassVar
+from typing import Dict, Any, Tuple, Optional
 
 from pydantic import BaseModel, root_validator
 
@@ -60,7 +60,7 @@ class TargetPlatformCapabilities(BaseModel):
     name: Optional[str] = "default_tpc"
     is_simd_padding: bool = False
 
-    SCHEMA_VERSION: ClassVar[int] = 2
+    SCHEMA_VERSION: int = 2
 
     class Config:
         frozen = True

model_compression_toolkit/target_platform_capabilities/tpc_io_handler.py

@@ -68,7 +68,7 @@ def load_target_platform_capabilities(tpc_obj_or_path: Union[tpc_or_str_type]) -
             f"but received type '{type(tpc_obj_or_path).__name__}'."
         )
 
-    if tpc.SCHEMA_VERSION == schema.TargetPlatformCapabilities.SCHEMA_VERSION:
+    if isinstance(tpc.SCHEMA_VERSION, schema.TargetPlatformCapabilities):  # if tpc is of current schema version
         return tpc
     return tpc_to_current_schema_version(tpc)
 

tests_pytest/common_tests/unit_tests/target_platform_capabilities/test_tpc.py

@@ -17,7 +17,8 @@
 import model_compression_toolkit.target_platform_capabilities.schema.mct_current_schema as schema
 from model_compression_toolkit.core.common import BaseNode
 from model_compression_toolkit.target_platform_capabilities.constants import KERNEL_ATTR
-from model_compression_toolkit.target_platform_capabilities.schema.schema_compatability import ALL_SCHEMA_VERSIONS
+from model_compression_toolkit.target_platform_capabilities.schema.schema_compatability import ALL_SCHEMA_VERSIONS, \
+    all_tpc_types, tpc_to_current_schema_version
 from model_compression_toolkit.target_platform_capabilities.schema.schema_functions import get_config_options_by_operators_set, is_opset_in_model
 from model_compression_toolkit.target_platform_capabilities.tpc_io_handler import load_target_platform_capabilities, \
     export_target_platform_capabilities
@@ -107,17 +108,15 @@ def test_valid_model_object(self, tpc):
         result = load_target_platform_capabilities(tpc)
         assert result == tpc
 
-    def test_new_schema(self):
+    def test_new_schema(self, tpc):
         """Tests that current schema is in all schemas list. This test validates new schema was added properly."""
-        current_version = schema.TargetPlatformCapabilities.SCHEMA_VERSION
-        all_supported_versions = [s.TargetPlatformCapabilities.SCHEMA_VERSION for s in ALL_SCHEMA_VERSIONS]
-        assert current_version in all_supported_versions, "Current schema need to be added to ALL_SCHEMA_VERSIONS"
+        assert type(tpc) in all_tpc_types, "Current schema need to be added to ALL_SCHEMA_VERSIONS"
 
     def test_schema_compatibility(self, tpc_by_schema_version):
         """Tests that tpc of any schema version is supported and can be converted into current schema version"""
         tpc_by_schema = tpc_by_schema_version
-        result = load_target_platform_capabilities(tpc_by_schema)
-        assert result.SCHEMA_VERSION == schema.TargetPlatformCapabilities.SCHEMA_VERSION, \
+        result = tpc_to_current_schema_version(tpc_by_schema)
+        assert isinstance(result, schema.TargetPlatformCapabilities), \
             f"Make sure schema version {result.SCHEMA_VERSION} can be converted into current schema version {result.SCHEMA_VERSION}"
 
     def test_invalid_json_parsing(self, tmp_invalid_json):