addressed some code rabbit comments

Author: Behnoosh Zamanlooy

src/midst_toolkit/attacks/tf/classification.py

@@ -117,10 +117,12 @@ def evaluate_model(model, x, y):
     optimizer = optim.Adam(regression_model.parameters(), lr=learning_rate)
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
+    has_validation = x_val is not None
     x_train = torch.tensor(x_train, dtype=torch.float32).to(device)
     y_train = torch.tensor(x_train_label, dtype=torch.float32).to(device)
-    x_val = torch.tensor(x_val, dtype=torch.float32).to(device)
-    y_test = torch.tensor(x_val_label, dtype=torch.float32).to(device)
+    if has_validation:
+        x_val = torch.tensor(x_val, dtype=torch.float32).to(device)
+        y_val = torch.tensor(x_val_label, dtype=torch.float32).to(device)
 
     indices = torch.randperm(x_train.size(0))
     x_train, y_train = x_train[indices], y_train[indices]
@@ -137,7 +139,7 @@ def evaluate_model(model, x, y):
         if (epoch + 1) % 100 == 0:
             train_loss, train_tpr = evaluate_model(regression_model, x_train, y_train)
             if x_val is not None:
-                test_loss, test_tpr = evaluate_model(regression_model, x_val, y_test)
+                test_loss, test_tpr = evaluate_model(regression_model, x_val, y_val)
                 if test_tpr > best_tpr:
                     best_tpr = test_tpr
                     save_best_model(regression_model, best_model_path)
@@ -153,7 +155,7 @@ def evaluate_model(model, x, y):
         load_best_model(regression_model, best_model_path, device)
 
     if x_val is not None:
-        test_loss, test_tpr = evaluate_model(regression_model, x_val, y_test)
+        test_loss, test_tpr = evaluate_model(regression_model, x_val, y_val)
         print(f"Final best loss: {test_loss}, best TPR: {test_tpr}")
 
     return regression_model

src/midst_toolkit/attacks/tf/data_utils.py

@@ -176,7 +176,11 @@ def get_tpr_at_fpr(true_membership: list[int], predictions: list[float], max_fpr
     Calculates the best True Positive Rate when the False Positive Rate is at most `max_fpr`.
     """
     fpr, tpr, _ = roc_curve(true_membership, predictions)
-    return max(tpr[fpr < max_fpr])
+
+    valid_tpr = tpr[fpr <= max_fpr]
+    if len(valid_tpr) == 0:
+        raise ValueError("No valid TPR values found for the given max FPR.")
+    return float(max(valid_tpr))
 
 
 def evaluate_attack_performance(

tests/integration/attacks/tf/test_tf_attack.py

@@ -34,7 +34,7 @@ def test_tf_attack_whitebox_small_config():
         "timesteps_list": [5],
         "addt_value_list": [0],
         "predictions_file_format": "test",
-        "results_path": Path("/h/behnzaman/midst-toolkit/tests/integration/attacks/tf/test_tf_attack_results"),
+        "results_path": Path("tests/integration/attacks/tf/test_tf_attack_results"),
         "use_best_checkpoint": True,
         "test_indices": [5],
         "train_indices": [1, 2],
@@ -44,10 +44,13 @@ def test_tf_attack_whitebox_small_config():
     mia_performance_train, mia_performance_val, mia_performance_test = tf_attack(**config)
     tpr_at_fpr_train, roc_auc_train = mia_performance_train.values()
     tpr_at_fpr_val, roc_auc_val = mia_performance_val.values()
-    (
-        tpr_at_fpr_test,
-        roc_auc_test,
-    ) = mia_performance_test.values()
+    tpr_at_fpr_test, roc_auc_test = mia_performance_test.values()
+    tpr_at_fpr_train = mia_performance_train["max_tpr"]
+    roc_auc_train = mia_performance_train["roc_auc"]
+    tpr_at_fpr_val = mia_performance_val["max_tpr"]
+    roc_auc_val = mia_performance_val["roc_auc"]
+    tpr_at_fpr_test = mia_performance_test["max_tpr"]
+    roc_auc_test = mia_performance_test["roc_auc"]
 
     assert roc_auc_train == pytest.approx(0.48133750000000003, abs=1e-8)
     assert tpr_at_fpr_train == pytest.approx(0.125, abs=1e-8)