Support return_assistant_tokens_mask on slow HybridDNATokenizer for completion-only SFT (apply_chat_template token-boundary mask override)

tokenizer.py

@@ -519,6 +519,61 @@ class HybridDNATokenizer(PreTrainedTokenizer):
 
         return BatchEncoding(result, tensor_type=return_tensors)
 
+    # Markers delimiting an assistant turn, used to build the assistant-token mask.
+    # Qwen chat format; override on the instance for other bases (e.g. Gemma uses
+    # "<start_of_turn>model\n" / "<end_of_turn>").
+    ASSISTANT_HEADER = "<|im_start|>assistant\n"
+    TURN_END = "<|im_end|>"
+
+    def apply_chat_template(self, conversation, *args, return_assistant_tokens_mask: bool = False, **kwargs):
+        """Completion-only assistant mask WITHOUT a fast tokenizer.
+
+        HF derives the assistant mask from ``{% generation %}`` char-spans via
+        ``BatchEncoding.char_to_token()``, which only fast (Rust) tokenizers implement.
+        HybridDNATokenizer is a slow Python tokenizer (the DNA k-mer logic can't be a
+        ``tokenizer.json``), so that path raises
+        ``char_to_token() is not available when using Python based tokenizers``.
+
+        Instead we tokenize normally and mark the tokens of every assistant turn -- from
+        just after ``ASSISTANT_HEADER`` up to and including the next ``TURN_END``. This
+        makes ``return_assistant_tokens_mask=True`` (and thus TRL ``assistant_only_loss``)
+        work unchanged, and needs no ``{% generation %}`` markers in the template.
+        """
+        if not return_assistant_tokens_mask:
+            return super().apply_chat_template(conversation, *args, **kwargs)
+        kwargs.pop("return_assistant_tokens_mask", None)
+        want_tensors = kwargs.pop("return_tensors", None)
+        kwargs["tokenize"] = True
+        kwargs["return_dict"] = True
+        enc = super().apply_chat_template(conversation, *args, return_tensors=None, **kwargs)
+        seqs = enc["input_ids"]
+        batched = len(seqs) > 0 and isinstance(seqs[0], (list, tuple))
+        batch = seqs if batched else [seqs]
+        hdr = self(self.ASSISTANT_HEADER, add_special_tokens=False)["input_ids"]
+        end = self(self.TURN_END, add_special_tokens=False)["input_ids"]
+        masks = []
+        for ids in batch:
+            ids = list(ids)
+            m = [0] * len(ids)
+            i = 0
+            while i <= len(ids) - len(hdr):
+                if ids[i:i + len(hdr)] == hdr:
+                    j = i + len(hdr)
+                    while j < len(ids) and ids[j:j + len(end)] != end:
+                        m[j] = 1
+                        j += 1
+                    for k in range(len(end)):          # include the turn-end token(s)
+                        if j + k < len(ids):
+                            m[j + k] = 1
+                    i = j + len(end)
+                else:
+                    i += 1
+            masks.append(m)
+        enc["assistant_masks"] = masks if batched else masks[0]
+        if want_tensors is not None:
+            enc = enc.convert_to_tensors(tensor_type=want_tensors)
+        return enc
+
     def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
         vocab_file = os.path.join(
             save_directory,