From d32d27f728c567bc49a5a4a963504a1ef2f5021b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Francesc=20Mart=C3=AD=20Escofet?= Date: Thu, 25 Jul 2024 09:02:42 +0200 Subject: [PATCH] Make methods private --- CHANGELOG.rst | 4 ++-- docs/examples/example_onnx.ipynb | 6 +++--- metalearners/drlearner.py | 10 +++++----- metalearners/metalearner.py | 4 ++-- metalearners/rlearner.py | 10 +++++----- metalearners/slearner.py | 4 ++-- metalearners/tlearner.py | 10 +++++----- metalearners/xlearner.py | 10 +++++----- tests/test_drlearner.py | 2 +- tests/test_metalearner.py | 4 ++-- tests/test_rlearner.py | 2 +- tests/test_tlearner.py | 2 +- tests/test_xlearner.py | 2 +- 13 files changed, 35 insertions(+), 35 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 9af8e28..3ab49a4 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -12,11 +12,11 @@ Changelog **New features** -* Add ``build_onnx`` to :class:`metalearners.MetaLearner` abstract class and implement it +* Add ``_build_onnx`` to :class:`metalearners.MetaLearner` abstract class and implement it for :class:`metalearners.TLearner`, :class:`metalearners.XLearner`, :class:`metalearners.RLearner` and :class:`metalearners.DRLearner`. -* Add ``necessary_onnx_models`` to :class:`metalearners.MetaLearner`. +* Add ``_necessary_onnx_models`` to :class:`metalearners.MetaLearner`. 0.8.0 (2024-07-22) ------------------ diff --git a/docs/examples/example_onnx.ipynb b/docs/examples/example_onnx.ipynb index 0dd2c05..9a1d2c5 100644 --- a/docs/examples/example_onnx.ipynb +++ b/docs/examples/example_onnx.ipynb @@ -147,7 +147,7 @@ "\n", "Before being able to convert the MetaLearner to ONXX we need to manually convert the necessary\n", "base models for the prediction. To get a list of the necessary base models that need to be\n", - "converted we can use :meth:`~metalearners.MetaLearner.necessary_onnx_models`." + "converted we can use :meth:`~metalearners.MetaLearner._necessary_onnx_models`." ] }, { @@ -156,7 +156,7 @@ "metadata": {}, "outputs": [], "source": [ - "xlearner.necessary_onnx_models()" + "xlearner._necessary_onnx_models()" ] }, { @@ -213,7 +213,7 @@ " )\n", " onnx_models[\"treatment_effect_model\"].append(onnx_model)\n", "\n", - "onnx_model = xlearner.build_onnx(onnx_models)" + "onnx_model = xlearner._build_onnx(onnx_models)" ] }, { diff --git a/metalearners/drlearner.py b/metalearners/drlearner.py index e7962cd..7b3db58 100644 --- a/metalearners/drlearner.py +++ b/metalearners/drlearner.py @@ -406,23 +406,23 @@ def _pseudo_outcome( return pseudo_outcome @classmethod - def necessary_onnx_models(cls) -> set[str]: + def _necessary_onnx_models(cls) -> set[str]: return {TREATMENT_MODEL} - @copydoc(MetaLearner.build_onnx, sep="") - def build_onnx(self, models: Mapping[str, Sequence], output_name: str = "tau"): + @copydoc(MetaLearner._build_onnx, sep="") + def _build_onnx(self, models: Mapping[str, Sequence], output_name: str = "tau"): """In the DRLearner case, the necessary models are: * ``"treatment_model"`` """ - warning_experimental_feature("build_onnx") + warning_experimental_feature("_build_onnx") check_spox_installed() import spox.opset.ai.onnx.v21 as op from onnx.checker import check_model from spox import Var, build, inline self._validate_feature_set_none() - self._validate_onnx_models(models, self.necessary_onnx_models()) + self._validate_onnx_models(models, self._necessary_onnx_models()) input_dict = infer_input_dict(models[TREATMENT_MODEL][0]) diff --git a/metalearners/metalearner.py b/metalearners/metalearner.py index 36e4c57..39b0e2a 100644 --- a/metalearners/metalearner.py +++ b/metalearners/metalearner.py @@ -1185,12 +1185,12 @@ def _validate_feature_set_none(self): @classmethod @abstractmethod - def necessary_onnx_models(cls) -> set[str]: + def _necessary_onnx_models(cls) -> set[str]: """Return a set with the necessary models to convert the MetaLearner to ONNX.""" ... @abstractmethod - def build_onnx(self, models: Mapping[str, Sequence], output_name: str = "tau"): + def _build_onnx(self, models: Mapping[str, Sequence], output_name: str = "tau"): """Convert the MetaLearner to an ONNX model. .. warning:: diff --git a/metalearners/rlearner.py b/metalearners/rlearner.py index 28435b7..f5e5283 100644 --- a/metalearners/rlearner.py +++ b/metalearners/rlearner.py @@ -526,23 +526,23 @@ def _pseudo_outcome_and_weights( return pseudo_outcomes, weights @classmethod - def necessary_onnx_models(cls) -> set[str]: + def _necessary_onnx_models(cls) -> set[str]: return {TREATMENT_MODEL} - @copydoc(MetaLearner.build_onnx, sep="") - def build_onnx(self, models: Mapping[str, Sequence], output_name: str = "tau"): + @copydoc(MetaLearner._build_onnx, sep="") + def _build_onnx(self, models: Mapping[str, Sequence], output_name: str = "tau"): """In the RLearner case, the necessary models are: * ``"treatment_model"`` """ - warning_experimental_feature("build_onnx") + warning_experimental_feature("_build_onnx") check_spox_installed() import spox.opset.ai.onnx.v21 as op from onnx.checker import check_model from spox import Var, build, inline self._validate_feature_set_none() - self._validate_onnx_models(models, self.necessary_onnx_models()) + self._validate_onnx_models(models, self._necessary_onnx_models()) input_dict = infer_input_dict(models[TREATMENT_MODEL][0]) diff --git a/metalearners/slearner.py b/metalearners/slearner.py index c058a83..250d1f7 100644 --- a/metalearners/slearner.py +++ b/metalearners/slearner.py @@ -300,12 +300,12 @@ def predict_conditional_average_outcomes( ) @classmethod - def necessary_onnx_models(cls) -> set[str]: + def _necessary_onnx_models(cls) -> set[str]: raise ValueError( "The SLearner does not implement this method. Please refer to comment in the tutorial." ) - def build_onnx(self, models: Mapping[str, Sequence], output_name: str = "tau"): + def _build_onnx(self, models: Mapping[str, Sequence], output_name: str = "tau"): raise ValueError( "The SLearner does not implement this method. Please refer to comment in the tutorial." ) diff --git a/metalearners/tlearner.py b/metalearners/tlearner.py index 3f5c91b..13c330c 100644 --- a/metalearners/tlearner.py +++ b/metalearners/tlearner.py @@ -151,23 +151,23 @@ def evaluate( ) @classmethod - def necessary_onnx_models(cls) -> set[str]: + def _necessary_onnx_models(cls) -> set[str]: return {VARIANT_OUTCOME_MODEL} - @copydoc(MetaLearner.build_onnx, sep="") - def build_onnx(self, models: Mapping[str, Sequence], output_name: str = "tau"): + @copydoc(MetaLearner._build_onnx, sep="") + def _build_onnx(self, models: Mapping[str, Sequence], output_name: str = "tau"): """In the TLearner case, the necessary models are: * ``"variant_outcome_model"`` """ - warning_experimental_feature("build_onnx") + warning_experimental_feature("_build_onnx") check_spox_installed() import spox.opset.ai.onnx.v21 as op from onnx.checker import check_model from spox import build, inline self._validate_feature_set_none() - self._validate_onnx_models(models, self.necessary_onnx_models()) + self._validate_onnx_models(models, self._necessary_onnx_models()) input_dict = infer_input_dict(models[VARIANT_OUTCOME_MODEL][0]) diff --git a/metalearners/xlearner.py b/metalearners/xlearner.py index 9c0b9b7..c9bde02 100644 --- a/metalearners/xlearner.py +++ b/metalearners/xlearner.py @@ -439,25 +439,25 @@ def _pseudo_outcome( return imputed_te_control, imputed_te_treatment @classmethod - def necessary_onnx_models(cls) -> set[str]: + def _necessary_onnx_models(cls) -> set[str]: return {PROPENSITY_MODEL, CONTROL_EFFECT_MODEL, TREATMENT_EFFECT_MODEL} - @copydoc(MetaLearner.build_onnx, sep="") - def build_onnx(self, models: Mapping[str, Sequence], output_name: str = "tau"): + @copydoc(MetaLearner._build_onnx, sep="") + def _build_onnx(self, models: Mapping[str, Sequence], output_name: str = "tau"): """In the XLearner case, the necessary models are: * ``"propensity_model"`` * ``"control_effect_model"`` * ``"treatment_effect_model"`` """ - warning_experimental_feature("build_onnx") + warning_experimental_feature("_build_onnx") check_spox_installed() import spox.opset.ai.onnx.v21 as op from onnx.checker import check_model from spox import Var, build, inline self._validate_feature_set_none() - self._validate_onnx_models(models, self.necessary_onnx_models()) + self._validate_onnx_models(models, self._necessary_onnx_models()) input_dict = infer_input_dict(models[PROPENSITY_MODEL][0]) diff --git a/tests/test_drlearner.py b/tests/test_drlearner.py index 3ebeafc..2ccb9cb 100644 --- a/tests/test_drlearner.py +++ b/tests/test_drlearner.py @@ -110,7 +110,7 @@ def test_drlearner_onnx( ) onnx_models.append(onnx_model) - final = ml.build_onnx({TREATMENT_MODEL: onnx_models}) + final = ml._build_onnx({TREATMENT_MODEL: onnx_models}) sess = rt.InferenceSession( final.SerializeToString(), providers=rt.get_available_providers() ) diff --git a/tests/test_metalearner.py b/tests/test_metalearner.py index e210c15..5e18153 100644 --- a/tests/test_metalearner.py +++ b/tests/test_metalearner.py @@ -112,10 +112,10 @@ def evaluate(self, X, y, w, is_oos, oos_method=None): def predict_conditional_average_outcomes(self, X, is_oos, oos_method=None): return np.zeros((len(X), 2, 1)) - def build_onnx(self, models: Mapping[str, Sequence], output_name: str = "tau"): ... + def _build_onnx(self, models: Mapping[str, Sequence], output_name: str = "tau"): ... @classmethod - def necessary_onnx_models(cls) -> set[str]: + def _necessary_onnx_models(cls) -> set[str]: return set() diff --git a/tests/test_rlearner.py b/tests/test_rlearner.py index 4f3dcb9..c35a508 100644 --- a/tests/test_rlearner.py +++ b/tests/test_rlearner.py @@ -115,7 +115,7 @@ def test_rlearner_onnx( ) onnx_models.append(onnx_model) - final = ml.build_onnx({TREATMENT_MODEL: onnx_models}) + final = ml._build_onnx({TREATMENT_MODEL: onnx_models}) sess = rt.InferenceSession( final.SerializeToString(), providers=rt.get_available_providers() ) diff --git a/tests/test_tlearner.py b/tests/test_tlearner.py index b984fa7..5732673 100644 --- a/tests/test_tlearner.py +++ b/tests/test_tlearner.py @@ -116,7 +116,7 @@ def test_tlearner_onnx( ) onnx_models.append(onnx_model) - final = ml.build_onnx({VARIANT_OUTCOME_MODEL: onnx_models}) + final = ml._build_onnx({VARIANT_OUTCOME_MODEL: onnx_models}) sess = rt.InferenceSession( final.SerializeToString(), providers=rt.get_available_providers() ) diff --git a/tests/test_xlearner.py b/tests/test_xlearner.py index 9d9e672..107310e 100644 --- a/tests/test_xlearner.py +++ b/tests/test_xlearner.py @@ -127,7 +127,7 @@ def test_xlearner_onnx( ) onnx_models[PROPENSITY_MODEL].append(onnx_model) - final = ml.build_onnx(onnx_models) + final = ml._build_onnx(onnx_models) sess = rt.InferenceSession( final.SerializeToString(), providers=rt.get_available_providers()