chore: import upstream snapshot with attribution

tests/test_frameworks.py

@@ -0,0 +1,143 @@
+"""
+Scripts to test cleanlab usage with various ML frameworks:
+pytorch, skorch
+"""
+
+import pytest
+import warnings
+
+# pytest.mark.filterwarnings is unable to catch filterbuffers library DeprecationWarning
+warnings.filterwarnings(action="ignore", category=DeprecationWarning)
+
+import os
+from copy import deepcopy
+import random
+import numpy as np
+
+import torch
+import skorch
+import sklearn
+
+from cleanlab.classification import CleanLearning
+
+# Version check for sklearn compatibility
+uses_sklearn_ge_1_6_0 = tuple(map(int, sklearn.__version__.split(".")[:2])) >= (1, 6)
+
+
+def dataset_w_errors():
+    num_classes = 2
+    num_features = 3
+    n = 50
+    margin = 5
+    X = np.vstack(
+        [np.random.random((n, num_features)), np.random.random((n, num_features)) + margin]
+    )
+    X = (X - X.mean(axis=0)) / X.std(axis=0)  # normalize columns
+    y = np.array([0] * n + [1] * n)
+    y_og = np.array(y)
+    # Introduce label errors
+    error_indices = [n - 3, n - 2, n - 1, n, n + 1, n + 2]
+    for idx in error_indices:
+        y[idx] = 1 - y[idx]  # Flip label
+
+    if os.name == "nt":  # running on Windows
+        # numpy converts to int32 instead of int64 on Windows, incompatible with neural nets
+        # https://github.com/numpy/numpy/issues/17640
+        X = np.float64(X)
+        y = np.int64(y)
+        y_og = np.int64(y_og)
+
+    return {
+        "X": X,
+        "y": y,
+        "y_og": y_og,
+        "error_indices": error_indices,
+        "num_classes": num_classes,
+        "num_features": num_features,
+    }
+
+
+def make_rare_label(data):
+    """Makes one label really rare in the dataset."""
+    data = deepcopy(data)
+    y = data["y"]
+    class0_inds = np.where(y == 0)[0]
+    if len(class0_inds) < 1:
+        raise ValueError("Class 0 too rare already")
+    class0_inds_remove = class0_inds[1:]
+    if len(class0_inds_remove) > 0:
+        y[class0_inds_remove] = 1
+    data["y"] = y
+    return data
+
+
+SEED = 1
+np.random.seed(SEED)
+random.seed(SEED)
+torch.manual_seed(SEED)
+torch.backends.cudnn.deterministic = True
+torch.backends.cudnn.benchmark = False
+torch.cuda.manual_seed_all(SEED)
+
+DATA = dataset_w_errors()
+DATA_RARE_LABEL = make_rare_label(DATA)
+
+
+def test_torch(data=DATA, hidden_units=128):
+    dataset = torch.utils.data.TensorDataset(
+        torch.from_numpy(data["X"]).float(), torch.from_numpy(data["y"])
+    )
+
+    class TorchNetwork(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.layers = torch.nn.Sequential(
+                torch.nn.Linear(data["num_features"], hidden_units),
+                torch.nn.ReLU(),
+                torch.nn.Linear(hidden_units, data["num_classes"]),
+            )
+
+        def forward(self, X):
+            return self.layers(X)
+
+    # Test base model works:
+    skorch_config = {"criterion": torch.nn.CrossEntropyLoss, "optimizer": torch.optim.Adam}
+    net = skorch.NeuralNet(TorchNetwork, **skorch_config)
+    net.fit(dataset, data["y"], epochs=2)
+    preds_base = net.predict(dataset)
+
+    # Test Cleanlearning performs well:
+    net = skorch.NeuralNet(TorchNetwork, **skorch_config)
+    cl = CleanLearning(net)
+    cl.fit(dataset, data["y"], clf_kwargs={"epochs": 30}, clf_final_kwargs={"epochs": 60})
+
+    preds = cl.predict(dataset).argmax(axis=1)
+    err = np.sum(preds != data["y_og"]) / len(data["y_og"])
+    issue_indices = list(cl.label_issues_df[cl.label_issues_df["is_label_issue"]].index.values)
+    assert issue_indices == data["error_indices"]
+    assert err < 1e-3
+
+
+@pytest.mark.filterwarnings("ignore")
+def test_torch_rarelabel(data=DATA_RARE_LABEL, hidden_units=8):
+    dataset = torch.utils.data.TensorDataset(
+        torch.from_numpy(data["X"]).float(), torch.from_numpy(data["y"])
+    )
+
+    class TorchNetwork(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.layers = torch.nn.Sequential(
+                torch.nn.Linear(data["num_features"], hidden_units),
+                torch.nn.ReLU(),
+                torch.nn.Linear(hidden_units, data["num_classes"]),
+            )
+
+        def forward(self, X):
+            return self.layers(X)
+
+    # Test Cleanlearning works:
+    net = skorch.NeuralNet(TorchNetwork, criterion=torch.nn.CrossEntropyLoss)
+    cl = CleanLearning(net)
+    cl.fit(dataset, data["y"], clf_kwargs={"epochs": 2})
+    pred_probs = cl.predict(dataset)