diff --git a/tests/utils/test_tags.py b/tests/utils/test_tags.py new file mode 100644 index 0000000..fe92d3b --- /dev/null +++ b/tests/utils/test_tags.py @@ -0,0 +1,55 @@ +"""Tests for torml.utils.get_tags.""" + +from __future__ import annotations + +import torch + +from torml.base import BaseEstimator +from torml.cluster import KMeans +from torml.linear_model import LinearRegression, LogisticRegression +from torml.preprocessing import StandardScaler +from torml.utils import get_tags +from torml.utils._tags import _DEFAULT_TAGS + + +class TestGetTags: + def test_classifier(self): + tags = get_tags(LogisticRegression()) + assert tags["_estimator_type"] == "classifier" + assert tags["require_y"] is True + assert tags["alias"] == "classifier" + + def test_regressor(self): + tags = get_tags(LinearRegression()) + assert tags["_estimator_type"] == "regressor" + assert tags["require_y"] is True + assert tags["alias"] == "regressor" + + def test_clusterer(self): + tags = get_tags(KMeans(n_clusters=2)) + assert tags["_estimator_type"] == "clusterer" + assert tags["pairwise"] is False + + def test_transformer(self): + tags = get_tags(StandardScaler()) + assert tags["_estimator_type"] == "preprocessor" + assert tags["require_y"] is False + + def test_unknown_estimator(self): + tags = get_tags(object()) + assert tags["_estimator_type"] == "unknown" + for key, value in _DEFAULT_TAGS.items(): + assert tags[key] == value + + def test_plain_base_estimator(self): + tags = get_tags(BaseEstimator()) + assert tags["_estimator_type"] == "unknown" + + def test_does_not_mutate_defaults(self): + before = dict(_DEFAULT_TAGS) + get_tags(LogisticRegression()) + assert _DEFAULT_TAGS == before + + def test_torch_free_input(self): + tags = get_tags(torch.nn.Module()) + assert tags["_estimator_type"] == "unknown" diff --git a/torml/utils/_tags.py b/torml/utils/_tags.py index eac4973..7aadc93 100644 --- a/torml/utils/_tags.py +++ b/torml/utils/_tags.py @@ -5,8 +5,6 @@ from __future__ import annotations -from collections import ChainMap - import torch _DEFAULT_TAGS = { @@ -23,140 +21,92 @@ "assume_is_fitted": ["_label_indices"], } +_CLASSIFIER_TAGS = { + "X_types": "1d/2d", + "require_x": True, + "require_y": True, + "transform_output_not_inverted": True, + "alias": "classifier", + "multifit": True, +} -def get_tags(estimator): - """Get tags for training/prediction with data X. - - Parameters - ---------- - estimator : object - Estimator instance. - - Returns - ------- - tags : dict of dicts - Dictionary containing tags. If _get_tags is defined on the estimator, - it is used to obtain tag values. - """ - tags = {} - for estimator_name, all_tags in { - "ClusterMixin": { - "X_types": "1d/2d", - "require_x": True, - "require_y": True, - "pairwise": False, - }, - "ClassifierMixin": { - "X_types": "1d/2d", - "require_x": True, - "require_y": True, - "transform_output_not_inverted": True, - "alias": "classifier", - "multifit": True, - }, - "RegressorMixin": { - "X_types": "1d/2d", - "require_x": True, - "require_y": True, - "transform_output_not_inverted": True, - "alias": "regressor", - "multifit": True, - }, - "TransformerMixin": { - "X_types": "1d/2d", - "require_x": True, - "require_y": False, - "_transform_inv": False, - "transform_output_not_inverted": False, - "alias": "preprocessor", - "_skip_check_finite": True, - "_skip_check_nan": True, - "_skip_check_min_samples": True, - "_skip_check_min_features": True, - "assume_is_fitted": ["_transform_weights"], - }, - }.get(type(estimator).__name__, {}): - combined = ChainMap(_DEFAULT_TAGS, all_tags) - tags = {k: v[v] for k, v in combined.items() if v not in _DEFAULT_TAGS} - tags = _concatenate_tags(tags, ["_transform_inv"], {}) - tags = _add_X_to_tags(tags, estimator_name) - - if hasattr(estimator, "_get_tags"): - tags = _add_custom_X_tags(tags, estimator._get_tags()) - - # Add estimator type - for _mixin_name, estimator_type in { - "ClassifierMixin": "classifier", - "RegressorMixin": "regressor", - "ClusterMixin": "clusterer", - "TransformerMixin": "preprocessor", - "__class__": "unknown", - }.get(type(estimator).__name__, {"__class__": "unknown"}): - _add_estimator_tag(tags, estimator_type) +_REGRESSOR_TAGS = { + "X_types": "1d/2d", + "require_x": True, + "require_y": True, + "transform_output_not_inverted": True, + "alias": "regressor", + "multifit": True, +} - return tags +_CLUSTER_TAGS = { + "X_types": "1d/2d", + "require_x": True, + "require_y": True, + "pairwise": False, +} +_TRANSFORMER_TAGS = { + "X_types": "1d/2d", + "require_x": True, + "require_y": False, + "_transform_inv": False, + "transform_output_not_inverted": False, + "alias": "preprocessor", + "_skip_check_finite": True, + "_skip_check_nan": True, + "_skip_check_min_samples": True, + "_skip_check_min_features": True, + "assume_is_fitted": ["_transform_weights"], +} -def _get_X_tag(tags, dim=1): - """Get tag associated with the 'X' key.""" - return tags["X"] if "X" in tags and len(tags["X"]) else _DEFAULT_TAG_VALUES +def get_tags(estimator): + """Get tags describing an estimator's requirements. -def _add_custom_X_tags(tags, tags_from_estimator): - """Add custom tags that apply to 'X'. + Tags are resolved with :func:`isinstance` against the mixin types, + so subclasses inherit their parents' tags. Priority for + ``_estimator_type`` is classifier, regressor, clusterer, transformer. Parameters ---------- - tags : dict - Default tags. - tags_from_estimator : dict - Custom tags with 'estimator_name' and 'X' keys. + estimator : object + Estimator instance. Returns ------- tags : dict - Updated tags with custom 'X' tags. + Defaults from ``_DEFAULT_TAGS`` overridden by the estimator type, + plus ``_estimator_type`` (``"unknown"`` for non-estimators). """ - for tag_name, tag_value in tags_from_estimator.items(): - if tag_name.startswith("estimator_name_"): - tag_name = tag_name.replace("estimator_name_", "") - if tag_name not in tags.get("X", {}): - tags["X"][tag_name] = tag_value - - # Special case: estimator_name - if "estimator_name" in tags_from_estimator: - if "_estimator_type" not in tags: - tags["X"]["estimator_name"] = tags_from_estimator["estimator_name"] - + from torml.base import ( + ClassifierMixin, + ClusterMixin, + RegressorMixin, + TransformerMixin, + ) + + tags = dict(_DEFAULT_TAGS) + estimator_type = "unknown" + if isinstance(estimator, ClassifierMixin): + tags.update(_CLASSIFIER_TAGS) + estimator_type = "classifier" + elif isinstance(estimator, RegressorMixin): + tags.update(_REGRESSOR_TAGS) + estimator_type = "regressor" + elif isinstance(estimator, ClusterMixin): + tags.update(_CLUSTER_TAGS) + estimator_type = "clusterer" + elif isinstance(estimator, TransformerMixin): + tags.update(_TRANSFORMER_TAGS) + estimator_type = "preprocessor" + tags["_estimator_type"] = estimator_type return tags -def _add_estimator_tag(tags, estimator_type): - """Add a tag associated with estimator type.""" - if "_estimator_type" not in tags: - tags["_estimator_type"] = estimator_type - - -def _add_X_to_tags(tags, estimator_name): - """Add tags that depend on 'X'.""" - if "X" not in tags: - tags["X"] = {} - - for key, value in tags.items(): - if key.startswith("_"): - continue - if isinstance(value, dict): - tags["X"][key] = value - - # Add n_samples_in_X_ if needed - if not {"_estimator_type", "name"} & set(tags.get("X", {}).keys()): - tags["X"]["name"] = f"{estimator_name}X" - - for key in ("_estimator_type", "name"): - if key not in tags["X"]: - tags["X"][key] = None - - return tags +def _get_X_tag(tags, dim=1): + """Get tag associated with the 'X' key.""" + return tags["X"] if "X" in tags and len(tags["X"]) else _DEFAULT_TAG_VALUES # Default tag values @@ -168,13 +118,3 @@ def _add_X_to_tags(tags, estimator_name): "finite": True, "sparse": False, } - - -def _concatenate_tags(tags, tag_names, tags_dict=None): - """Concatenate tags and add to tags_dict.""" - if tags_dict is None: - tags_dict = {} - for tag_name in tag_names: - if tag_name in tags and tag_name not in tags_dict: - tags_dict[tag_name] = tags[tag_name] - return tags_dict