unlearning LF

Data.py

@@ -1,8 +1,10 @@
 from torchvision import datasets, transforms, models
+from torch.utils.data import Dataset, DataLoader, Subset
 import torch
 import numpy as np
 
 
+
 # train set transform 
 def train_transform(res):
     return transforms.Compose([
@@ -101,3 +103,39 @@ def get_indices(dataset, identities, split_at, size = 30):
         test_indices.extend(indices[split_at:size])
 
     return train_indices, test_indices
+
+
+
+def get_forget_retain_loaders(dataset: Dataset, forget_class_idx: int, batch_size: int = 32) -> tuple[DataLoader, DataLoader]:
+    """
+    Splits an IdentitySubset or standard Dataset into forget and retain sets 
+    based on a remapped target class index.
+    """
+    # 1. Safely extract targets whether it's a standard dataset or a Subset wrapper
+    if hasattr(dataset, 'targets'):
+        targets = dataset.targets
+    elif hasattr(dataset, 'identity'):  # Raw CelebA support
+        targets = dataset.identity
+    else:
+        # If it's an IdentitySubset or standard Subset, extract mapped targets sequentially
+        # This guarantees we get the 0 -> (n-1) remapped labels
+        targets = [dataset[i][1] for i in range(len(dataset))]
+        
+    if not isinstance(targets, torch.Tensor):
+        targets = torch.tensor(targets)
+        
+    # 2. Generate mask indices local to this subset
+    forget_indices = torch.where(targets == forget_class_idx)[0].tolist()
+    retain_indices = torch.where(targets != forget_class_idx)[0].tolist()
+    
+    # 3. Create PyTorch Subsets
+    forget_subset = Subset(dataset, forget_indices)
+    retain_subset = Subset(dataset, retain_indices)
+    
+    # 4. Wrap into clean DataLoaders
+    forget_loader = DataLoader(forget_subset, batch_size=batch_size, shuffle=False)
+    retain_loader = DataLoader(retain_subset, batch_size=batch_size, shuffle=True)
+    
+    print(f"[Data Split] Local Class {forget_class_idx}: {len(forget_subset)} samples | Remaining Classes: {len(retain_subset)} samples.")
+    
+    return forget_loader, retain_loader