Switched to raw str_labels and dropping dupes during __add__()

utils/prompts.py

@@ -553,6 +553,14 @@ def __init__(
         self.position = position
         self.label = label
 
+    def drop_duplicates(self):
+        _, indices = torch.unique(
+            self.clean_tokens[torch.arange(len(self.clean_tokens)), self.position], 
+            dim=0, 
+            return_inverse=True
+        )
+        return self.get_subset(indices.tolist())
+
     def __add__(self, other: "CleanCorruptedDataset"):
         assert isinstance(other, CleanCorruptedDataset)
         assert self.tokenizer is not None
@@ -581,7 +589,7 @@ def __add__(self, other: "CleanCorruptedDataset"):
             self.tokenizer,
             None,
             self.label,
-        )
+        ).drop_duplicates()
 
     def get_subset(self, indices: List[int]):
         return CleanCorruptedDataset(

utils/residual_stream.py

@@ -52,8 +52,8 @@ def __init__(
             torch.arange(len(self.prompt_tokens)), self.position
         ].cpu().detach()
         str_tokens = [
-            f"{pos}:{tok}" 
-            for pos, tok in zip(position, model.to_str_tokens(label_tensor))
+            f"{tok}" 
+            for tok in model.to_str_tokens(label_tensor)
         ]
         to_str_check = (
             len(str_tokens) == len(self.prompt_tokens) and