tokenizer: fix EOS append bug, decode skip_special_tokens=True, add auto_dna_tags

tokenizer.py

@@ -8,7 +8,9 @@ Supports token_mask for Fine-grained Nucleotide Supervision (FNS):
   -2: padding token
   -1: text token (BPE)
    0: DNA special token (<dna>, </dna>, <oov>)
-  1-5: partial 6-mer (number of valid bases)
+  1-5: partial 6-mer token — valid_length real bases at positions [0, valid_length),
+       right-padded with 'A' at positions [valid_length, k) so loss can supervise
+       positions 0..valid_len-1 via pos_mask = (valid_len > pos)
    6: full 6-mer
 """
 
@@ -26,6 +28,12 @@ class HybridDNATokenizer(PreTrainedTokenizer):
 
     DNA regions must be wrapped in <dna>...</dna> tags to be tokenized as 6-mers.
     Without tags, DNA sequences are tokenized as regular BPE text.
+
+    For pure-DNA input (no metadata tokens), pass auto_dna_tags=True to have
+    <dna>...</dna> tags added automatically when they are absent.  Do NOT set
+    this if the input may contain BPE metadata such as species tags
+    (<fungi_species> etc.) — those must appear outside <dna>...</dna> and would
+    be incorrectly k-mer encoded if auto-wrapping fired.
     """
 
     model_input_names = ["input_ids", "attention_mask"]
@@ -34,6 +42,7 @@ class HybridDNATokenizer(PreTrainedTokenizer):
         self,
         base_tokenizer_path: Optional[str] = None,
         k: int = 6,
+        auto_dna_tags: bool = False,
         **kwargs
     ):
         self.k = k
@@ -63,6 +72,7 @@ class HybridDNATokenizer(PreTrainedTokenizer):
         )
 
         self.special_tokens = self.dna_special_tokens + [self._eos_token, self._pad_token]
+        self.auto_dna_tags = auto_dna_tags
 
     def _init_dna_vocab(self):
         """Initialize DNA vocabulary (special tokens + k-mers + padding for 128 alignment)."""
@@ -228,6 +238,10 @@ class HybridDNATokenizer(PreTrainedTokenizer):
 
         if remaining:
             padding_needed = k - len(remaining)
+            # Right-pad with A: real bases occupy positions [0, valid_length).
+            # The hybrid BP loss supervises positions 0..valid_len-1 via
+            #   pos_mask = (valid_len > pos)
+            # so padding must be at the END, not the start.
             padded = remaining + 'A' * padding_needed
 
             if is_valid_kmer(padded):
@@ -265,8 +279,13 @@ class HybridDNATokenizer(PreTrainedTokenizer):
         text: str,
         add_special_tokens: bool = False,
         return_token_mask: bool = False,
+        auto_dna_tags: Optional[bool] = None,
         **kwargs
     ) -> Union[List[int], Tuple[List[int], List[int]]]:
+        use_auto = self.auto_dna_tags if auto_dna_tags is None else auto_dna_tags
+        if use_auto and '<dna>' not in text:
+            text = f'<dna>{text}</dna>'
+
         segments = self._split_by_dna_tags(text)
 
         token_ids = []
@@ -309,10 +328,11 @@ class HybridDNATokenizer(PreTrainedTokenizer):
                 if return_token_mask:
                     token_mask.extend([-1] * len(base_ids))
 
-        if add_special_tokens and self.eos_token_id is not None:
-            token_ids.append(self.eos_token_id)
-            if return_token_mask:
-                token_mask.append(-1)
+        # Do NOT append EOS when add_special_tokens=True. Qwen3 doesn't add
+        # BOS/EOS either, and appending EOS here breaks lighteval's
+        # tok_encode_pair: it relies on
+        #   len(encode(ctx)) + len(encode(answer)) == len(encode(ctx + answer))
+        # which the extra EOS violates by shifting the split by 1.
 
         if return_token_mask:
             return token_ids, token_mask
@@ -357,7 +377,14 @@ class HybridDNATokenizer(PreTrainedTokenizer):
                         i += 1
 
             elif tid in self.dna_id_to_token:
-                if not skip_special_tokens:
+                # This branch handles k-mer tokens that appear without a <dna>
+                # wrapper — the common generation case where <dna> was in the
+                # prompt but only the generated portion is being decoded.
+                # K-mer tokens are content, not special tokens, so always decode
+                # them.  Only drop true DNA special tokens (<dna>, </dna>, <oov>)
+                # when skip_special_tokens=True.
+                is_dna_special = tid in (self.dna_begin_token_id, self.dna_end_token_id, self.oov_token_id)
+                if not (skip_special_tokens and is_dna_special):
                     parts.append(self.dna_id_to_token[tid])
                 i += 1
 
@@ -400,6 +427,7 @@ class HybridDNATokenizer(PreTrainedTokenizer):
         max_length: Optional[int] = None,
         return_tensors: Optional[str] = None,
         return_token_mask: bool = False,
+        auto_dna_tags: Optional[bool] = None,
         **kwargs
     ) -> Dict[str, Any]:
         is_batch = isinstance(text, list)
@@ -410,11 +438,11 @@ class HybridDNATokenizer(PreTrainedTokenizer):
 
         for t in texts:
             if return_token_mask:
-                ids, mask = self.encode(t, add_special_tokens=add_special_tokens, return_token_mask=True)
+                ids, mask = self.encode(t, add_special_tokens=add_special_tokens, return_token_mask=True, auto_dna_tags=auto_dna_tags)
                 all_ids.append(ids)
                 all_masks.append(mask)
             else:
-                ids = self.encode(t, add_special_tokens=add_special_tokens, return_token_mask=False)
+                ids = self.encode(t, add_special_tokens=add_special_tokens, return_token_mask=False, auto_dna_tags=auto_dna_tags)
                 all_ids.append(ids)
 
         if padding:
@@ -496,6 +524,7 @@ class HybridDNATokenizer(PreTrainedTokenizer):
             "dna_start_id": self.dna_start_id,
             "dna_vocab_size": self.dna_vocab_size,
             "dna_special_tokens": self.dna_special_tokens,
+            "auto_dna_tags": self.auto_dna_tags,
         }
 
         dna_config_path = os.path.join(save_directory, "dna_config.json")
@@ -517,6 +546,7 @@ class HybridDNATokenizer(PreTrainedTokenizer):
                 "AutoTokenizer": ["tokenizer.HybridDNATokenizer", None]
             },
             "k": self.k,
+            "auto_dna_tags": self.auto_dna_tags,
         })
 
         with open(config_path, "w", encoding="utf-8") as f:
@@ -533,19 +563,21 @@ class HybridDNATokenizer(PreTrainedTokenizer):
 
     @classmethod
     def from_pretrained(cls, pretrained_model_name_or_path: str, **kwargs):
+        k = 6
+        auto_dna_tags = False
+
         dna_config_path = os.path.join(pretrained_model_name_or_path, "dna_config.json")
+        tok_config_path = os.path.join(pretrained_model_name_or_path, "tokenizer_config.json")
 
         if os.path.exists(dna_config_path):
             with open(dna_config_path, "r") as f:
                 dna_config = json.load(f)
             k = dna_config.get("k", 6)
-        else:
-            config_path = os.path.join(pretrained_model_name_or_path, "tokenizer_config.json")
-            if os.path.exists(config_path):
-                with open(config_path, "r") as f:
-                    config = json.load(f)
-                k = config.get("k", 6)
-            else:
-                k = 6
-
-        return cls(base_tokenizer_path=pretrained_model_name_or_path, k=k, **kwargs)
+            auto_dna_tags = dna_config.get("auto_dna_tags", False)
+        elif os.path.exists(tok_config_path):
+            with open(tok_config_path, "r") as f:
+                tok_config = json.load(f)
+            k = tok_config.get("k", 6)
+            auto_dna_tags = tok_config.get("auto_dna_tags", False)
+
+        return cls(base_tokenizer_path=pretrained_model_name_or_path, k=k, auto_dna_tags=auto_dna_tags, **kwargs)