mirror of
https://github.com/luxfi/torus-ml.git
synced 2026-10-30 15:15:23 +00:00
zeekay
e037e8f7ea
The distribution called itself torus-ml while shipping a package directory named concrete, so the name a user installs and the name they import disagreed — and most of the tree already said torus.ml, which could not resolve against src/concrete. Moved src/concrete to src/torus and pointed the include at it. The dependency was Zama's concrete-python from their index. It is torus-fhe from ours, which is what luxfhe/torus already ships: that frontend declares torus-fhe and ships a torus/fhe module, so every 'from torus.fhe' in this tree was already naming our compiler and only the dependency lagged. Two distributions share one namespace, which is the shape the upstream this forked from used and the reason the split works: torus-fhe gives torus.fhe, torus-ml gives torus.ml. Swept every other spelling — identifiers, enum members, workflow variables, filenames, notebooks. LICENSE and NOTICE keep the upstream attribution a fork owes; the figures carry the old name in JPEG metadata a text sweep cannot reach.
384 lines
13 KiB
Python
384 lines
13 KiB
Python
"""Implement sklearn's Generalized Linear Models (GLM)."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from abc import abstractmethod
|
|
from typing import Any, Dict, Union
|
|
|
|
import numpy
|
|
import sklearn
|
|
import sklearn.linear_model
|
|
|
|
from ..common.debugging.custom_assert import assert_true
|
|
from ..common.utils import FheMode
|
|
from ..onnx.onnx_model_manipulations import clean_graph_after_node_op_type
|
|
from .base import Data, SklearnLinearRegressorMixin
|
|
|
|
|
|
# pylint: disable-next=too-many-instance-attributes
|
|
class _GeneralizedLinearRegressor(SklearnLinearRegressorMixin):
|
|
"""Regression via a penalized Generalized Linear Model (GLM) with FHE.
|
|
|
|
Parameters:
|
|
n_bits (int, Dict[str, int]): Number of bits to quantize the model. If an int is passed
|
|
for n_bits, the value will be used for quantizing inputs and weights. If a dict is
|
|
passed, then it should contain "op_inputs" and "op_weights" as keys with
|
|
corresponding number of quantization bits so that:
|
|
- op_inputs: number of bits to quantize the input values
|
|
- op_weights: number of bits to quantize the learned parameters
|
|
Default to 8.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
n_bits: Union[int, dict] = 8,
|
|
alpha: float = 1.0,
|
|
fit_intercept: bool = True,
|
|
solver: str = "lbfgs",
|
|
max_iter: int = 100,
|
|
tol: float = 1e-4,
|
|
warm_start: bool = False,
|
|
verbose: int = 0,
|
|
):
|
|
# Call SklearnLinearModelMixin's __init__ method
|
|
super().__init__(n_bits=n_bits)
|
|
|
|
self.alpha = alpha
|
|
self.fit_intercept = fit_intercept
|
|
self.solver = solver
|
|
self.max_iter = max_iter
|
|
self.tol = tol
|
|
self.warm_start = warm_start
|
|
self.verbose = verbose
|
|
|
|
def _clean_graph(self) -> None:
|
|
assert self.onnx_model_ is not None, self._is_not_fitted_error_message()
|
|
|
|
# Remove any operators following gemm, as they will be done in the clear in post-processing
|
|
# In particular, this includes the exponential operator
|
|
clean_graph_after_node_op_type(self.onnx_model_, node_op_type="Gemm")
|
|
super()._clean_graph()
|
|
|
|
def post_processing(self, y_preds: numpy.ndarray) -> numpy.ndarray:
|
|
return self._inverse_link(y_preds)
|
|
|
|
def predict(self, X: Data, fhe: Union[FheMode, str] = FheMode.DISABLE) -> numpy.ndarray:
|
|
# Call SklearnLinearModelMixin's predict method
|
|
y_preds = super().predict(X, fhe=fhe)
|
|
|
|
y_preds = self.post_processing(y_preds)
|
|
return y_preds
|
|
|
|
@abstractmethod
|
|
def _inverse_link(self, y_preds: numpy.ndarray) -> numpy.ndarray:
|
|
"""Apply the link function's inverse on the inputs.
|
|
|
|
Args:
|
|
y_preds (numpy.ndarray): The input data.
|
|
"""
|
|
|
|
def dump_dict(self) -> Dict:
|
|
assert self._weight_quantizer is not None, self._is_not_fitted_error_message()
|
|
|
|
metadata: Dict[str, Any] = {}
|
|
|
|
# Torus ML
|
|
metadata["n_bits"] = self.n_bits
|
|
metadata["sklearn_model"] = self.sklearn_model
|
|
metadata["_is_fitted"] = self._is_fitted
|
|
metadata["_is_compiled"] = self._is_compiled
|
|
metadata["input_quantizers"] = self.input_quantizers
|
|
metadata["_weight_quantizer"] = self._weight_quantizer
|
|
metadata["output_quantizers"] = self.output_quantizers
|
|
metadata["onnx_model_"] = self.onnx_model_
|
|
metadata["_q_weights"] = self._q_weights
|
|
metadata["_q_bias"] = self._q_bias
|
|
metadata["post_processing_params"] = self.post_processing_params
|
|
|
|
# Scikit-Learn
|
|
metadata["alpha"] = self.alpha
|
|
metadata["fit_intercept"] = self.fit_intercept
|
|
metadata["solver"] = self.solver
|
|
metadata["max_iter"] = self.max_iter
|
|
metadata["tol"] = self.tol
|
|
metadata["warm_start"] = self.warm_start
|
|
metadata["verbose"] = self.verbose
|
|
|
|
return metadata
|
|
|
|
@classmethod
|
|
def load_dict(cls, metadata: Dict):
|
|
|
|
# Instantiate the model
|
|
obj = cls(n_bits=metadata["n_bits"])
|
|
|
|
# Torus ML
|
|
obj.n_bits = metadata["n_bits"]
|
|
obj.sklearn_model = metadata["sklearn_model"]
|
|
obj.onnx_model_ = metadata["onnx_model_"]
|
|
obj._is_fitted = metadata["_is_fitted"]
|
|
obj._is_compiled = metadata["_is_compiled"]
|
|
obj.input_quantizers = metadata["input_quantizers"]
|
|
obj._weight_quantizer = metadata["_weight_quantizer"]
|
|
obj.output_quantizers = metadata["output_quantizers"]
|
|
obj._q_weights = metadata["_q_weights"]
|
|
obj._q_bias = metadata["_q_bias"]
|
|
obj.post_processing_params = metadata["post_processing_params"]
|
|
|
|
# Scikit-Learn
|
|
obj.alpha = metadata["alpha"]
|
|
obj.fit_intercept = metadata["fit_intercept"]
|
|
obj.solver = metadata["solver"]
|
|
obj.max_iter = metadata["max_iter"]
|
|
obj.tol = metadata["tol"]
|
|
obj.warm_start = metadata["warm_start"]
|
|
obj.verbose = metadata["verbose"]
|
|
|
|
return obj
|
|
|
|
def get_sklearn_params(self, deep: bool = True) -> dict:
|
|
# Here, the `get_params` method is the `BaseEstimator.get_params` method from scikit-learn
|
|
# FIXME: https://github.com/luxfi/torus-ml-internal/issues/3373
|
|
params = super().get_params(deep=deep) # type: ignore[misc]
|
|
|
|
# Remove the parameters added by Torus ML
|
|
params.pop("n_bits", None)
|
|
# Remove sklearn 1.5 parameter when using sklearn 1.1
|
|
if "1.1." in sklearn.__version__:
|
|
params.pop("solver", None) # pragma: no cover
|
|
|
|
return params
|
|
|
|
|
|
class PoissonRegressor(_GeneralizedLinearRegressor):
|
|
"""A Poisson regression model with FHE.
|
|
|
|
Parameters:
|
|
n_bits (int, Dict[str, int]): Number of bits to quantize the model. If an int is passed
|
|
for n_bits, the value will be used for quantizing inputs and weights. If a dict is
|
|
passed, then it should contain "op_inputs" and "op_weights" as keys with
|
|
corresponding number of quantization bits so that:
|
|
- op_inputs : number of bits to quantize the input values
|
|
- op_weights: number of bits to quantize the learned parameters
|
|
Default to 8.
|
|
|
|
For more details on PoissonRegressor please refer to the scikit-learn documentation:
|
|
https://scikit-learn.org/stable/modules/generated/sklearn.linear_model.PoissonRegressor.html
|
|
"""
|
|
|
|
sklearn_model_class = sklearn.linear_model.PoissonRegressor
|
|
_is_a_public_cml_model = True
|
|
|
|
def __init__( # pylint: disable=useless-parent-delegation
|
|
self,
|
|
*,
|
|
n_bits: Union[int, dict] = 8,
|
|
alpha: float = 1.0,
|
|
fit_intercept: bool = True,
|
|
solver="lbfgs",
|
|
max_iter: int = 100,
|
|
tol: float = 1e-4,
|
|
warm_start: bool = False,
|
|
verbose: int = 0,
|
|
):
|
|
super().__init__(
|
|
n_bits=n_bits,
|
|
alpha=alpha,
|
|
fit_intercept=fit_intercept,
|
|
solver=solver,
|
|
max_iter=max_iter,
|
|
tol=tol,
|
|
warm_start=warm_start,
|
|
verbose=verbose,
|
|
)
|
|
|
|
def _inverse_link(self, y_preds: numpy.ndarray) -> numpy.ndarray:
|
|
return numpy.exp(y_preds)
|
|
|
|
|
|
class GammaRegressor(_GeneralizedLinearRegressor):
|
|
"""A Gamma regression model with FHE.
|
|
|
|
Parameters:
|
|
n_bits (int, Dict[str, int]): Number of bits to quantize the model. If an int is passed
|
|
for n_bits, the value will be used for quantizing inputs and weights. If a dict is
|
|
passed, then it should contain "op_inputs" and "op_weights" as keys with
|
|
corresponding number of quantization bits so that:
|
|
- op_inputs : number of bits to quantize the input values
|
|
- op_weights: number of bits to quantize the learned parameters
|
|
Default to 8.
|
|
|
|
For more details on GammaRegressor please refer to the scikit-learn documentation:
|
|
https://scikit-learn.org/stable/modules/generated/sklearn.linear_model.GammaRegressor.html
|
|
"""
|
|
|
|
sklearn_model_class = sklearn.linear_model.GammaRegressor
|
|
_is_a_public_cml_model = True
|
|
|
|
def __init__( # pylint: disable=useless-parent-delegation
|
|
self,
|
|
*,
|
|
n_bits: Union[int, dict] = 8,
|
|
alpha: float = 1.0,
|
|
fit_intercept: bool = True,
|
|
solver="lbfgs",
|
|
max_iter: int = 100,
|
|
tol: float = 1e-4,
|
|
warm_start: bool = False,
|
|
verbose: int = 0,
|
|
):
|
|
super().__init__(
|
|
n_bits=n_bits,
|
|
alpha=alpha,
|
|
fit_intercept=fit_intercept,
|
|
solver=solver,
|
|
max_iter=max_iter,
|
|
tol=tol,
|
|
warm_start=warm_start,
|
|
verbose=verbose,
|
|
)
|
|
|
|
def _inverse_link(self, y_preds: numpy.ndarray) -> numpy.ndarray:
|
|
return numpy.exp(y_preds)
|
|
|
|
|
|
# pylint: disable-next=too-many-instance-attributes
|
|
class TweedieRegressor(_GeneralizedLinearRegressor):
|
|
"""A Tweedie regression model with FHE.
|
|
|
|
Parameters:
|
|
n_bits (int, Dict[str, int]): Number of bits to quantize the model. If an int is passed
|
|
for n_bits, the value will be used for quantizing inputs and weights. If a dict is
|
|
passed, then it should contain "op_inputs" and "op_weights" as keys with
|
|
corresponding number of quantization bits so that:
|
|
- op_inputs : number of bits to quantize the input values
|
|
- op_weights: number of bits to quantize the learned parameters
|
|
Default to 8.
|
|
|
|
For more details on TweedieRegressor please refer to the scikit-learn documentation:
|
|
https://scikit-learn.org/stable/modules/generated/sklearn.linear_model.TweedieRegressor.html
|
|
"""
|
|
|
|
sklearn_model_class = sklearn.linear_model.TweedieRegressor
|
|
_is_a_public_cml_model = True
|
|
|
|
def __init__( # pylint: disable=too-many-arguments
|
|
self,
|
|
*,
|
|
n_bits: Union[int, dict] = 8,
|
|
power: float = 0.0,
|
|
alpha: float = 1.0,
|
|
fit_intercept: bool = True,
|
|
link: str = "auto",
|
|
solver="lbfgs",
|
|
max_iter: int = 100,
|
|
tol: float = 1e-4,
|
|
warm_start: bool = False,
|
|
verbose: int = 0,
|
|
):
|
|
super().__init__(
|
|
n_bits=n_bits,
|
|
alpha=alpha,
|
|
fit_intercept=fit_intercept,
|
|
solver=solver,
|
|
max_iter=max_iter,
|
|
tol=tol,
|
|
warm_start=warm_start,
|
|
verbose=verbose,
|
|
)
|
|
|
|
assert_true(
|
|
link in ["auto", "log", "identity"],
|
|
f"link must be an element of ['auto', 'identity', 'log'], got '{link}'",
|
|
)
|
|
|
|
self.power = power
|
|
self.link = link
|
|
|
|
def _set_post_processing_params(self):
|
|
super()._set_post_processing_params()
|
|
self.post_processing_params.update(
|
|
{
|
|
"link": self.link,
|
|
"power": self.power,
|
|
}
|
|
)
|
|
|
|
def _inverse_link(self, y_preds: numpy.ndarray) -> numpy.ndarray:
|
|
self.check_model_is_fitted()
|
|
|
|
if self.post_processing_params["link"] == "auto":
|
|
|
|
# Identity link
|
|
if self.post_processing_params["power"] <= 0:
|
|
return y_preds
|
|
|
|
# Log link
|
|
return numpy.exp(y_preds)
|
|
|
|
if self.post_processing_params["link"] == "log":
|
|
return numpy.exp(y_preds)
|
|
|
|
return y_preds
|
|
|
|
def dump_dict(self) -> Dict:
|
|
assert self._weight_quantizer is not None, self._is_not_fitted_error_message()
|
|
|
|
metadata: Dict[str, Any] = {}
|
|
|
|
# Torus ML
|
|
metadata["n_bits"] = self.n_bits
|
|
metadata["sklearn_model"] = self.sklearn_model
|
|
metadata["_is_fitted"] = self._is_fitted
|
|
metadata["_is_compiled"] = self._is_compiled
|
|
metadata["input_quantizers"] = self.input_quantizers
|
|
metadata["_weight_quantizer"] = self._weight_quantizer
|
|
metadata["output_quantizers"] = self.output_quantizers
|
|
metadata["onnx_model_"] = self.onnx_model_
|
|
metadata["_q_weights"] = self._q_weights
|
|
metadata["_q_bias"] = self._q_bias
|
|
metadata["post_processing_params"] = self.post_processing_params
|
|
metadata["power"] = self.power
|
|
metadata["link"] = self.link
|
|
|
|
# Scikit-Learn
|
|
metadata["alpha"] = self.alpha
|
|
metadata["fit_intercept"] = self.fit_intercept
|
|
metadata["max_iter"] = self.max_iter
|
|
metadata["tol"] = self.tol
|
|
metadata["warm_start"] = self.warm_start
|
|
metadata["verbose"] = self.verbose
|
|
|
|
return metadata
|
|
|
|
@classmethod
|
|
def load_dict(cls, metadata: Dict):
|
|
|
|
# Instantiate the model
|
|
obj = cls(n_bits=metadata["n_bits"])
|
|
|
|
# Torus ML
|
|
obj.sklearn_model = metadata["sklearn_model"]
|
|
obj.onnx_model_ = metadata["onnx_model_"]
|
|
obj._is_fitted = metadata["_is_fitted"]
|
|
obj._is_compiled = metadata["_is_compiled"]
|
|
obj.input_quantizers = metadata["input_quantizers"]
|
|
obj._weight_quantizer = metadata["_weight_quantizer"]
|
|
obj.output_quantizers = metadata["output_quantizers"]
|
|
obj._q_weights = metadata["_q_weights"]
|
|
obj._q_bias = metadata["_q_bias"]
|
|
obj.power = metadata["power"]
|
|
obj.link = metadata["link"]
|
|
obj.post_processing_params = metadata["post_processing_params"]
|
|
|
|
# Scikit-Learn
|
|
obj.alpha = metadata["alpha"]
|
|
obj.fit_intercept = metadata["fit_intercept"]
|
|
obj.max_iter = metadata["max_iter"]
|
|
obj.tol = metadata["tol"]
|
|
obj.warm_start = metadata["warm_start"]
|
|
obj.verbose = metadata["verbose"]
|
|
|
|
return obj
|