Merge pull request #743 from wearepal/fix-problems

Fix two long-standing problems

ethicml/data/tabular_data/acs.py

@@ -323,7 +323,7 @@ def load(
             target_transform=lambda x: x > self.target_threshold,
             group=self.split,
             preprocess=adult_filter,
-            postprocess=lambda x: np.nan_to_num(x, -1),
+            postprocess=lambda x: np.nan_to_num(x, nan=-1),
         )
 
         dataframe = data_obj._preprocess(dataframe)

ethicml/implementations/adv_debiasing_method.py

@@ -10,11 +10,11 @@
 import numpy as np
 import torch
 
-from ethicml import DataTuple, SoftPrediction, SubgroupTuple
 from ethicml.implementations.adv_debiasing_modules.model import AdvDebiasingClassLearner
-from ethicml.models.inprocess.adv_debiasing import AdvDebArgs
+from ethicml.utility import DataTuple, ModelType, SoftPrediction, SubgroupTuple
 
 if TYPE_CHECKING:
+    from ethicml.models.inprocess.adv_debiasing import AdvDebArgs
     from ethicml.models.inprocess.in_subprocess import InAlgoArgs
 
 
@@ -33,7 +33,7 @@ def fit(train: DataTuple, args: AdvDebArgs, seed: int = 888) -> AdvDebiasingClas
         cost_pred=torch.nn.CrossEntropyLoss(),
         in_shape=len(train.x.columns),
         batch_size=args["batch_size"],
-        model_type=args["model_type"],
+        model_type=ModelType(args["model_type"]),
         num_classes=train.y.nunique(),
         lambda_vec=args["lambda_vec"],
     )

ethicml/implementations/adv_debiasing_modules/model.py

@@ -22,6 +22,7 @@
     LinearModel,
     make_dataset_and_loader,
 )
+from ethicml.utility.data_structures import ModelType
 
 
 class Adversary(nn.Module):
@@ -199,7 +200,7 @@ def __init__(
         cost_pred: nn.Module,
         in_shape: int,
         batch_size: int,
-        model_type: Literal["deep_model", "linear_model"],
+        model_type: ModelType,
         num_classes: int,
         lambda_vec: float,
     ):
@@ -210,9 +211,9 @@ def __init__(
         self.num_classes = num_classes
 
         self.model_type = model_type
-        if self.model_type == "deep_model":
+        if self.model_type is ModelType.deep:
             self.clf: nn.Module = DeepModel(in_shape=in_shape, out_shape=num_classes)
-        elif self.model_type == "linear_model":
+        elif self.model_type is ModelType.linear:
             self.clf = LinearModel(in_shape=in_shape, out_shape=num_classes)
         else:
             raise NotImplementedError
@@ -291,7 +292,7 @@ def __init__(
         cost_pred: nn.Module,
         in_shape: int,
         batch_size: int,
-        model_type: Literal["deep_model", "linear_model"],
+        model_type: ModelType,
         out_shape: int,
         lambda_vec: float,
     ):
@@ -302,9 +303,9 @@ def __init__(
         self.out_shape = out_shape
 
         self.model_type = model_type
-        if self.model_type == "deep_model":
+        if self.model_type is ModelType.deep:
             self.clf: nn.Module = DeepRegModel(in_shape=in_shape, out_shape=out_shape)
-        elif self.model_type == "linear_model":
+        elif self.model_type is ModelType.linear:
             self.clf = LinearModel(in_shape=in_shape, out_shape=out_shape)
         else:
             raise NotImplementedError

ethicml/implementations/fair_dummies_modules/model.py

@@ -12,6 +12,8 @@
 import torch.optim as optim
 from torch.utils.data import DataLoader
 
+from ethicml.utility.data_structures import ModelType
+
 from ... import DataTuple
 from ..pytorch_common import DeepModel, DeepRegModel, LinearModel, PandasDataSet
 from .utility_functions import density_estimation
@@ -450,7 +452,7 @@ def __init__(
         cost_pred: nn.Module,
         in_shape: int,
         batch_size: int,
-        model_type: Literal["deep_model", "linear_model"],
+        model_type: ModelType,
         lambda_vec: float,
         second_moment_scaling: float,
         num_classes: int,
@@ -471,9 +473,9 @@ def __init__(
         self.num_classes = num_classes
 
         self.model_type = model_type
-        if self.model_type == "deep_model":
+        if self.model_type is ModelType.deep:
             self.model: nn.Module = DeepModel(in_shape=in_shape, out_shape=num_classes)
-        elif self.model_type == "linear_model":
+        elif self.model_type is ModelType.linear:
             self.model = LinearModel(in_shape=in_shape, out_shape=num_classes)
         else:
             raise NotImplementedError
@@ -611,7 +613,7 @@ def __init__(
         cost_pred: nn.Module,
         in_shape: int,
         batch_size: int,
-        model_type: Literal["deep_model", "linear_model"],
+        model_type: ModelType,
         lambda_vec: float,
         second_moment_scaling: float,
         out_shape: int,
@@ -625,9 +627,9 @@ def __init__(
         self.use_standardscaler = use_standardscaler
 
         self.model_type = model_type
-        if self.model_type == "deep_model":
+        if self.model_type is ModelType.deep:
             self.model: nn.Module = DeepRegModel(in_shape=in_shape, out_shape=out_shape)
-        elif self.model_type == "linear_model":
+        elif self.model_type is ModelType.linear:
             self.model = LinearModel(in_shape=in_shape, out_shape=out_shape)
         else:
             raise NotImplementedError

ethicml/implementations/fair_dummies_romano.py

@@ -11,7 +11,7 @@
 import torch
 
 from ethicml.implementations.fair_dummies_modules.model import EquiClassLearner
-from ethicml.utility import DataTuple, SoftPrediction, SubgroupTuple, TestTuple
+from ethicml.utility import DataTuple, ModelType, SoftPrediction, SubgroupTuple, TestTuple
 
 if TYPE_CHECKING:
     from ethicml.models.inprocess.fair_dummies import FairDummiesArgs
@@ -35,7 +35,7 @@ def fit(train: DataTuple, args: FairDummiesArgs, seed: int = 888) -> EquiClassLe
         cost_pred=torch.nn.CrossEntropyLoss(),
         in_shape=len(train.x.columns),
         batch_size=args["batch_size"],
-        model_type=args["model_type"],
+        model_type=ModelType(args["model_type"]),
         lambda_vec=args["lambda_vec"],
         second_moment_scaling=args["second_moment_scaling"],
         num_classes=train.y.nunique(),

ethicml/implementations/hgr_method.py

@@ -11,6 +11,7 @@
 
 from ethicml import DataTuple, SoftPrediction, SubgroupTuple
 from ethicml.implementations.hgr_modules.hgr_impl import HgrClassLearner
+from ethicml.utility.data_structures import ModelType
 
 if TYPE_CHECKING:
     from ethicml.models.inprocess.hgr import HgrArgs
@@ -39,7 +40,7 @@ def fit(train: DataTuple, args: HgrArgs, seed: int = 888) -> HgrClassLearner:
         in_shape=len(train.x.columns),
         out_shape=train.y.nunique(),
         batch_size=args["batch_size"],
-        model_type=args["model_type"],
+        model_type=ModelType(args["model_type"]),
     )
     return model.fit(train, seed=seed)
 

ethicml/implementations/hgr_modules/hgr_impl.py

@@ -5,7 +5,6 @@
 # https://github.com/criteo-research/continuous-fairness/tree/master/facl/independence
 from __future__ import annotations
 import random
-from typing import Literal
 from typing_extensions import Self
 
 import numpy as np
@@ -14,8 +13,14 @@
 import torch
 import torch.nn as nn
 
-from ... import DataTuple
-from ..pytorch_common import DeepModel, DeepRegModel, LinearModel, make_dataset_and_loader
+from ethicml.implementations.pytorch_common import (
+    DeepModel,
+    DeepRegModel,
+    LinearModel,
+    make_dataset_and_loader,
+)
+from ethicml.utility.data_structures import DataTuple, ModelType
+
 from .density_estimation import Kde
 from .facl_hgr import chi_2_cond
 
@@ -48,9 +53,9 @@ def __init__(self, lr, epochs, mu, cost_pred, in_shape, out_shape, batch_size, m
         self.batch_size = batch_size
 
         self.out_shape = out_shape
-        if self.model_type == "deep_model":
+        if self.model_type is ModelType.deep:
             self.model: nn.Module = DeepRegModel(in_shape=in_shape, out_shape=out_shape)
-        elif self.model_type == "linear_model":
+        elif self.model_type is ModelType.linear:
             self.model = LinearModel(in_shape=in_shape, out_shape=out_shape)
 
         else:
@@ -132,7 +137,7 @@ def __init__(
         in_shape: int,
         out_shape: int,
         batch_size: int,
-        model_type: Literal["deep_model", "linear_model"],
+        model_type: ModelType,
     ):
 
         self.in_shape = in_shape
@@ -148,9 +153,9 @@ def __init__(
         self.batch_size = batch_size
 
         self.model_type = model_type
-        if self.model_type == "deep_model":
+        if self.model_type is ModelType.deep:
             self.model: nn.Module = DeepModel(in_shape=in_shape, out_shape=out_shape)
-        elif self.model_type == "linear_model":
+        elif self.model_type is ModelType.linear:
             self.model = LinearModel(in_shape=in_shape, out_shape=out_shape)
         else:
             raise NotImplementedError

ethicml/implementations/pytorch_common.py

@@ -7,6 +7,7 @@
 import pandas as pd
 
 from ethicml.utility import DataTuple, TestTuple
+from ethicml.utility.data_structures import ModelType
 
 try:
     import torch

ethicml/metrics/dependence_measures.py

@@ -2,11 +2,11 @@
 from __future__ import annotations
 from abc import ABC
 from dataclasses import dataclass
-from enum import Enum, auto
+from enum import auto
 from typing import ClassVar
 
 import numpy as np
-from ranzen import enum_name_str, implements
+from ranzen import StrEnum, implements
 from sklearn.metrics import normalized_mutual_info_score
 
 from ethicml.utility import EvalTuple, Prediction
@@ -16,8 +16,7 @@
 __all__ = ["DependencyTarget", "NMI", "RenyiCorrelation", "Yanovich"]
 
 
-@enum_name_str
-class DependencyTarget(Enum):
+class DependencyTarget(StrEnum):
     """The variable that is compared to the predictions in order to check how similar they are."""
 
     s = auto()

ethicml/models/inprocess/adv_debiasing.py

@@ -4,12 +4,13 @@
 """
 from __future__ import annotations
 from dataclasses import dataclass
-from typing import Literal, TypedDict
+from typing import TypedDict
 
 from ranzen import implements
 
 from ethicml.models.inprocess.in_algorithm import HyperParamType
 from ethicml.models.inprocess.in_subprocess import InAlgorithmSubprocess
+from ethicml.utility.data_structures import ModelType
 
 __all__ = ["AdvDebiasing", "AdvDebArgs"]
 
@@ -22,7 +23,7 @@ class AdvDebArgs(TypedDict):
     n_adv_epochs: int
     n_epoch_combined: int
     batch_size: int
-    model_type: Literal["deep_model", "linear_model"]
+    model_type: str
     lambda_vec: float
 
 
@@ -35,7 +36,7 @@ class AdvDebiasing(InAlgorithmSubprocess):
     n_adv_epochs = 2
     n_epoch_combined = 40
     batch_size = 32
-    model_type: Literal["deep_model", "linear_model"] = "deep_model"
+    model_type: ModelType = ModelType.deep
     lambda_vec = 0.999999
 
     @implements(InAlgorithmSubprocess)

ethicml/models/inprocess/agarwal_reductions.py

@@ -19,8 +19,8 @@
 class AgarwalArgs(TypedDict):
     """Args for the Agarwal implementation."""
 
-    classifier: int
-    fairness: int
+    classifier: str
+    fairness: str
     eps: float
     iters: int
     C: float
@@ -57,12 +57,12 @@ def _get_flags(self) -> AgarwalArgs:
         chosen_c, chosen_kernel = settings_for_svm_lr(self.classifier, self.C, self.kernel)
         # TODO: replace this with dataclasses.asdict()
         return {
-            "classifier": self.classifier.value,
-            "fairness": self.fairness.value,
+            "classifier": self.classifier,
+            "fairness": self.fairness,
             "eps": self.eps,
             "iters": self.iters,
             "C": chosen_c,
-            "kernel": str(chosen_kernel) if chosen_kernel is not None else "",
+            "kernel": chosen_kernel if chosen_kernel is not None else "",
         }
 
     @property  # type: ignore[misc]

ethicml/models/inprocess/fair_dummies.py

@@ -4,12 +4,13 @@
 """
 from __future__ import annotations
 from dataclasses import dataclass
-from typing import Literal, TypedDict
+from typing import TypedDict
 
 from ranzen import implements
 
 from ethicml.models.inprocess.in_algorithm import HyperParamType
 from ethicml.models.inprocess.in_subprocess import InAlgorithmSubprocess
+from ethicml.utility.data_structures import ModelType
 
 __all__ = ["FairDummies", "FairDummiesArgs"]
 
@@ -24,7 +25,7 @@ class FairDummiesArgs(TypedDict):
     loss_steps: int
     dis_steps: int
     batch_size: int
-    model_type: Literal["deep_model", "linear_model"]
+    model_type: str
     lambda_vec: float
     second_moment_scaling: float
 
@@ -40,7 +41,7 @@ class FairDummies(InAlgorithmSubprocess):
     loss_steps: int = 1
     dis_steps: int = 1
     batch_size: int = 32
-    model_type: str = "deep_model"
+    model_type: ModelType = ModelType.deep
     lambda_vec: float = 0.9
     second_moment_scaling: float = 1e-5
 
@@ -50,9 +51,6 @@ def _get_path_to_script(self) -> list[str]:
 
     @implements(InAlgorithmSubprocess)
     def _get_flags(self) -> FairDummiesArgs:
-        model_type: Literal["deep_model", "linear_model"] = (
-            "deep_model" if self.model_type.lower() == "deep_model" else "linear_model"
-        )
         return {
             "lr": self.lr,
             "pretrain_pred_epochs": self.pretrain_pred_epochs,
@@ -61,7 +59,7 @@ def _get_flags(self) -> FairDummiesArgs:
             "loss_steps": self.loss_steps,
             "dis_steps": self.dis_steps,
             "batch_size": self.batch_size,
-            "model_type": model_type,
+            "model_type": self.model_type,
             "lambda_vec": self.lambda_vec,
             "second_moment_scaling": self.second_moment_scaling,
         }
@@ -85,4 +83,4 @@ def hyperparameters(self) -> HyperParamType:
     @property  # type: ignore[misc]
     @implements(InAlgorithmSubprocess)
     def name(self) -> str:
-        return f"Fair Dummies {self.model_type}"
+        return f"Fair Dummies {self.model_type}_model"