feat: enable import of fitted linear sklearn models

zama-ai · Sep 21, 2023 · dd89497 · dd89497
1 parent 8ed0eac
commit dd89497
Show file tree

Hide file tree

Showing 14 changed files with 639 additions and 253 deletions.
diff --git a/src/concrete/ml/common/serialization/decoder.py b/src/concrete/ml/common/serialization/decoder.py
@@ -18,7 +18,7 @@
     UniformQuantizationParameters,
     UniformQuantizer,
 )
-from ...sklearn import get_sklearn_models
+from ...sklearn import get_sklearn_all_models
 from . import SUPPORTED_TORCH_ACTIVATIONS, USE_SKOPS
 
 # If USE_SKOPS is False or Skops can't be imported, default to pickle
@@ -49,7 +49,7 @@ def _get_fully_qualified_name(object_class: Type) -> str:
 ]
 
 _TRUSTED_CONCRETE_MODELS = [
-    _get_fully_qualified_name(model_class) for model_class in get_sklearn_models()["all"]
+    _get_fully_qualified_name(model_class) for model_class in get_sklearn_all_models()
 ]
 
 # Define all the trusted types that Skops should consider
@@ -181,7 +181,7 @@ def object_hook(d: Any) -> Any:
         # `dump_dict` and `load_dict` method) if not already done
         if not SERIALIZABLE_CLASSES:
             serializable_classes = (
-                get_sklearn_models()["all"]
+                get_sklearn_all_models()
                 + list(ALL_QUANTIZED_OPS)
                 + [
                     QuantizedArray,