fix masking error

Q-Future · Jan 21, 2024 · 4007d89 · 4007d89
1 parent f79ed74
commit 4007d89
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/q_align/train/train_mem.py b/q_align/train/train_mem.py
@@ -370,7 +370,7 @@ def preprocess_v1(
 
             if has_image:
                 round_len = len(tokenizer_image_token(rou, tokenizer))
-                instruction_len = len(tokenizer_image_token(parts[0], tokenizer)) - 2
+                instruction_len = len(tokenizer_image_token(parts[0], tokenizer)) - 3
             else:
                 round_len = len(tokenizer(rou).input_ids)
                 instruction_len = len(tokenizer(parts[0]).input_ids) - 2
@@ -387,7 +387,6 @@ def preprocess_v1(
                     f"WARNING: tokenization mismatch: {cur_len} vs. {total_len}."
                     f" (ignored)"
                 )
-
     return dict(
         input_ids=input_ids,
         labels=targets,