✨ Rebuild serialisation protocol to allow versions (#192)

Author: perdy

.pre-commit-config.yaml

@@ -13,7 +13,7 @@ repos:
       - id: name-tests-test
         args:
           - --django
-        exclude: "asserts.py|utils.py"
+        exclude: "_utils/.*"
       - id: pretty-format-json
         args:
           - --autofix

flama/compat.py

@@ -1,6 +1,6 @@
 import sys
 
-__all__ = ["Self", "NotRequired", "StrEnum", "tomllib", "get_annotations"]
+__all__ = ["Self", "NotRequired", "StrEnum", "tomllib", "get_annotations", "bz2", "lzma", "zlib", "zstd"]
 
 # PORT: Remove when stop supporting 3.10
 # Self was added in Python 3.11
@@ -54,3 +54,15 @@ def _generate_next_value_(name, start, count, last_values):
     from annotationlib import get_annotations
 else:
     from typing_extensions import get_annotations
+
+# PORT: Remove when stop supporting 3.13
+# compression was added in Python 3.14
+# https://docs.python.org/3/library/compression.html
+if sys.version_info >= (3, 14):
+    from compression import bz2, lzma, zlib, zstd
+else:
+    import bz2
+    import lzma
+    import zlib
+
+    import zstd

flama/models/base.py

@@ -4,10 +4,10 @@
 if t.TYPE_CHECKING:
     from flama.serialize.data_structures import Artifacts, Metadata
 
-__all__ = ["Model"]
+__all__ = ["BaseModel"]
 
 
-class Model:
+class BaseModel:
     def __init__(self, model: t.Any, meta: "Metadata", artifacts: "Artifacts | None"):
         self.model = model
         self.meta = meta

flama/models/components.py

@@ -1,10 +1,12 @@
 import importlib
 import os
+import pathlib
+import typing as t
 
+from flama import types
 from flama.injection import Component
-from flama.models.base import Model
-from flama.serialize import load
-from flama.serialize.types import Framework
+from flama.models.base import BaseModel
+from flama.serialize.serializer import Serializer
 
 __all__ = ["ModelComponent", "ModelComponentBuilder"]
 
@@ -13,29 +15,38 @@ class ModelComponent(Component):
     def __init__(self, model):
         self.model = model
 
-    def get_model_type(self) -> type[Model]:
+    def get_model_type(self) -> type[BaseModel]:
         return self.model.__class__  # type: ignore[no-any-return]
 
 
 class ModelComponentBuilder:
+    _module_name: t.Final[str] = "flama.models.models.{}"
+    _class_name: t.Final[str] = "Model"
+    _modules: t.Final[dict[types.MLLib, str]] = {
+        "keras": "tensorflow",
+        "sklearn": "sklearn",
+        "tensorflow": "tensorflow",
+        "torch": "pytorch",
+    }
+
     @classmethod
-    def _get_model_class(cls, framework: Framework) -> type[Model]:
+    def _get_model_class(cls, lib: types.MLLib) -> type[BaseModel]:
         try:
-            module, class_name = {
-                Framework.torch: ("pytorch", "PyTorchModel"),
-                Framework.sklearn: ("sklearn", "SKLearnModel"),
-                Framework.tensorflow: ("tensorflow", "TensorFlowModel"),
-                Framework.keras: ("tensorflow", "TensorFlowModel"),
-            }[framework]
+            return getattr(importlib.import_module(cls._module_name.format(cls._modules[lib])), cls._class_name)
         except KeyError:  # pragma: no cover
-            raise ValueError("Wrong framework")
-
-        model_class: type[Model] = getattr(importlib.import_module(f"flama.models.models.{module}"), class_name)
-        return model_class
+            raise ValueError(f"Wrong lib '{lib}'")
+        except ModuleNotFoundError:  # pragma: no cover
+            raise ValueError(f"Module not found '{cls._module_name.format(cls._modules[lib])}'")
+        except AttributeError:  # pragma: no cover
+            raise ValueError(
+                f"Class '{cls._class_name}' not found in module '{cls._module_name.format(cls._modules[lib])}'"
+            )
 
     @classmethod
-    def load(cls, path: str | os.PathLike) -> ModelComponent:
-        load_model = load(path)
+    def load(cls, path: str | os.PathLike | pathlib.Path) -> ModelComponent:
+        with pathlib.Path(str(path)).open("rb") as f:
+            load_model = Serializer.load(f)
+
         parent = cls._get_model_class(load_model.meta.framework.lib)
         model_class = type(parent.__name__, (parent,), {})
         model_obj = model_class(load_model.model, load_model.meta, load_model.artifacts)

flama/models/models/pytorch.py

@@ -1,15 +1,17 @@
 import typing as t
 
 from flama import exceptions
-from flama.models.base import Model
+from flama.models.base import BaseModel
 
 try:
     import torch  # type: ignore
 except Exception:  # pragma: no cover
     torch = None  # type: ignore
 
+__all__ = ["Model"]
 
-class PyTorchModel(Model):
+
+class Model(BaseModel):
     def predict(self, x: list[list[t.Any]]) -> t.Any:
         if torch is None:  # noqa
             raise exceptions.FrameworkNotInstalled("pytorch")

flama/models/models/sklearn.py

@@ -1,15 +1,18 @@
 import typing as t
 
 from flama import exceptions
-from flama.models.base import Model
+from flama.models.base import BaseModel
 
 try:
     import sklearn  # type: ignore
 except Exception:  # pragma: no cover
     sklearn = None
 
 
-class SKLearnModel(Model):
+__all__ = ["Model"]
+
+
+class Model(BaseModel):
     def predict(self, x: list[list[t.Any]]) -> t.Any:
         if sklearn is None:  # noqa
             raise exceptions.FrameworkNotInstalled("scikit-learn")

flama/models/models/tensorflow.py

@@ -1,7 +1,7 @@
 import typing as t
 
 from flama import exceptions
-from flama.models.base import Model
+from flama.models.base import BaseModel
 
 try:
     import numpy as np  # type: ignore
@@ -14,7 +14,10 @@
     tf = None
 
 
-class TensorFlowModel(Model):
+__all__ = ["Model"]
+
+
+class Model(BaseModel):
     def predict(self, x: list[list[t.Any]]) -> t.Any:
         if np is None:  # noqa
             raise exceptions.FrameworkNotInstalled("numpy")

flama/models/resource.py

@@ -10,7 +10,7 @@
 from flama.resources.routing import ResourceRoute
 
 if t.TYPE_CHECKING:
-    from flama.models.base import Model
+    from flama.models.base import BaseModel
     from flama.models.components import ModelComponent
 
 __all__ = ["BaseModelResource", "ModelResource", "InspectMixin", "PredictMixin", "ModelResourceType"]
@@ -21,7 +21,9 @@
 
 class InspectMixin:
     @classmethod
-    def _add_inspect(cls, name: str, verbose_name: str, model_model_type: type["Model"], **kwargs) -> dict[str, t.Any]:
+    def _add_inspect(
+        cls, name: str, verbose_name: str, model_model_type: type["BaseModel"], **kwargs
+    ) -> dict[str, t.Any]:
         @ResourceRoute.method("/", methods=["GET"], name="inspect")
         async def inspect(self, model: model_model_type):  # type: ignore[valid-type]
             return model.inspect()  # type: ignore[attr-defined]
@@ -44,7 +46,9 @@ async def inspect(self, model: model_model_type):  # type: ignore[valid-type]
 
 class PredictMixin:
     @classmethod
-    def _add_predict(cls, name: str, verbose_name: str, model_model_type: type["Model"], **kwargs) -> dict[str, t.Any]:
+    def _add_predict(
+        cls, name: str, verbose_name: str, model_model_type: type["BaseModel"], **kwargs
+    ) -> dict[str, t.Any]:
         @ResourceRoute.method("/predict/", methods=["POST"], name="predict")
         async def predict(
             self,

flama/serialize/__init__.py

@@ -1,5 +1,3 @@
-from flama.serialize.data_structures import *  # noqa
-from flama.serialize.dump import *  # noqa
-from flama.serialize.load import *  # noqa
+from flama.serialize.serializer import dump, load
 
-__all__ = ["dump", "load"]  # noqa
+__all__ = ["dump", "load"]