Correct Collator (#1975)

Fixes a mistake in the labels mask, related to #1974.
LAION-AI · Mar 6, 2023 · ba336e3 · ba336e3
1 parent 85fbaf6
commit ba336e3
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 10 deletions.
diff --git a/model/model_training/custom_datasets/dialogue_collator.py b/model/model_training/custom_datasets/dialogue_collator.py
@@ -54,12 +54,13 @@ def __call__(self, features):
                     list(map(lambda x: x[1], flatten_message["offset_mapping"])),
                 )
             )
-            label_mask = np.roll(list(map(lambda x: x % 2 == 1, message_indices)), -1, -1)
-            try:
-                label_mask[[i for i in range(len(message_indices)) if message_indices[i] == -2][0] - 1] = True
-            except IndexError:
-                # due to truncation, we might not have the last termination token
-                label_mask[-1] = False
+            label_mask = np.array(list(map(lambda x: x % 2 == 1, message_indices)))
+            label_mask[-1] = False  # make sure last token is inactive, has an effect only when truncatting
+            # try:
+            #     label_mask[[i for i in range(len(message_indices)) if message_indices[i] == -2][0] - 1] = True
+            # except IndexError:
+            #     # due to truncation, we might not have the last termination token
+            #     label_mask[-1] = False
 
             label_masks.append(label_mask)
             if len(flatten_message["input_ids"]) < self.mix_length_threshold and self.samples_mixing:

diff --git a/model/model_training/custom_datasets/formatting.py b/model/model_training/custom_datasets/formatting.py
@@ -3,10 +3,7 @@
 
 def format_pair(pairs):
     return [
-        "{}{}{}".format(QA_SPECIAL_TOKENS["Question"], pairs[i], QA_SPECIAL_TOKENS["Answer"])
-        if i % 2 == 0
-        else pairs[i]
-        for i in range(len(pairs))
+        "{}{}".format(QA_SPECIAL_TOKENS["Question" if i % 2 == 0 else "Answer"], pairs[i]) for i in range(len(pairs))
     ]