unlearning done

sets/Data.py

@@ -1,5 +1,5 @@
 from torchvision import datasets, transforms
-from torch.utils.data import Dataset, DataLoader, Subset
+from torch.utils.data import Dataset, DataLoader, Subset, ConcatDataset
 import torch
 import numpy as np
 import os
@@ -181,4 +181,59 @@ def get_unlearning_loaders(dataset: Dataset, forget_class_idx: int, batch_size:
     
     print(f"[Data Split] Local Class {forget_class_idx}: {len(forget_subset)} samples | Remaining Classes: {len(retain_subset)} samples.")
     
-    return forget_loader, retain_loader
+    return retain_loader, forget_loader
+
+
+def vertical_split(dataset, batch_size,num_classes):
+        """
+        Executes a class-wise vertical split. 
+        Divides the samples of every single identity class exactly in half:
+        50% of each class goes to the Retain Set, 50% goes to the Forget Set.
+        """
+        
+        # 1. Group dataset indices by their respective ground-truth classes
+        class_to_indices = {c: [] for c in range(num_classes)}
+        
+        print("  [Vertical Split] Tracking class indices across the combined dataset...")
+        for idx in range(len(dataset)):
+            # Extract the label cleanly from the underlying dataset structure
+            _, label = dataset[idx]
+            if label in class_to_indices:
+                class_to_indices[label].append(idx)
+
+        retain_indices = []
+        forget_indices = []
+
+        # 2. Slice each class identity vertically (exactly 50/50)
+        for c, indices in class_to_indices.items():
+            if len(indices) < 2:
+                print(f"  Warning: Class {c} has fewer than 2 samples. Cannot split vertically.")
+                retain_indices.extend(indices)
+                continue
+                
+            # Deterministic shuffle per class to ensure honest distribution before splitting
+            np.random.shuffle(indices)
+            
+            mid = len(indices) // 2
+            forget_indices.extend(indices[:mid])   # First half assigned to unlearning
+            retain_indices.extend(indices[mid:])   # Second half assigned to retention
+
+        print(f"  Vertical split complete: Retain Index Size = {len(retain_indices)} | Forget Index Size = {len(forget_indices)}")
+
+        # 3. Construct lightweight PyTorch Subsets using our sliced index maps
+        retain_subset = Subset(dataset, retain_indices)
+        forget_subset = Subset(dataset, forget_indices)
+
+        # 4. Return pristine, shuffled DataLoaders mirroring your environment's batch specifications
+        retain_loader = DataLoader(retain_subset, batch_size=batch_size, shuffle=True)
+        forget_loader = DataLoader(forget_subset, batch_size=batch_size, shuffle=True)
+
+        return retain_loader, forget_loader
+
+def _combine_set(loader_one, loader_two):
+    full_train_dataset = ConcatDataset([loader_one.dataset, loader_two.dataset])
+    return DataLoader(
+        full_train_dataset, 
+        batch_size=loader_one.batch_size, 
+        shuffle=True
+    )