Carbon-X-SFT weights and tokenizer

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

chat_template.jinja

@@ -0,0 +1,91 @@
+{#- Training variant of the Qwen3 chat template (see qwen3.jinja for the original).
+    Modifications vs the original:
+    - Removed the loop.index0 > ns.last_query_index conditional; always include thinking block.
+      This makes the template prefix-preserving for the [user, assistant] → [user, assistant, tool] transition.
+    - Added {% generation %} / {% endgeneration %} around assistant message output to support
+      assistant-only loss masking in SFT training.
+-#}
+{%- if tools %}
+    {{- '<|im_start|>system\n' }}
+    {%- if messages[0].role == 'system' %}
+        {{- messages[0].content + '\n\n' }}
+    {%- endif %}
+    {{- "# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>" }}
+    {%- for tool in tools %}
+        {{- "\n" }}
+        {{- tool | tojson }}
+    {%- endfor %}
+    {{- "\n</tools>\n\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\n<tool_call>\n{\"name\": <function-name>, \"arguments\": <args-json-object>}\n</tool_call><|im_end|>\n" }}
+{%- else %}
+    {%- if messages[0].role == 'system' %}
+        {{- '<|im_start|>system\n' + messages[0].content + '<|im_end|>\n' }}
+    {%- endif %}
+{%- endif %}
+{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}
+{%- for message in messages[::-1] %}
+    {%- set index = (messages|length - 1) - loop.index0 %}
+    {%- if ns.multi_step_tool and message.role == "user" and message.content is string and not(message.content.startswith('<tool_response>') and message.content.endswith('</tool_response>')) %}
+        {%- set ns.multi_step_tool = false %}
+        {%- set ns.last_query_index = index %}
+    {%- endif %}
+{%- endfor %}
+{%- for message in messages %}
+    {%- if message.content is string %}
+        {%- set content = message.content %}
+    {%- else %}
+        {%- set content = '' %}
+    {%- endif %}
+    {%- if (message.role == "user") or (message.role == "system" and not loop.first) %}
+        {{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }}
+    {%- elif message.role == "assistant" %}
+        {%- set reasoning_content = '' %}
+        {%- if message.reasoning_content is string %}
+            {%- set reasoning_content = message.reasoning_content %}
+        {%- else %}
+            {%- if '</think>' in content %}
+                {%- set reasoning_content = content.split('</think>')[0].rstrip('\n').split('<think>')[-1].lstrip('\n') %}
+                {%- set content = content.split('</think>')[-1].lstrip('\n') %}
+            {%- endif %}
+        {%- endif %}
+        {{- '<|im_start|>' + message.role + '\n' }}
+        {%- generation %}
+        {{- '<think>\n' + reasoning_content.strip('\n') + '\n</think>\n\n' + content.lstrip('\n') }}
+        {%- if message.tool_calls %}
+            {%- for tool_call in message.tool_calls %}
+                {%- if (loop.first and content) or (not loop.first) %}
+                    {{- '\n' }}
+                {%- endif %}
+                {%- if tool_call.function %}
+                    {%- set tool_call = tool_call.function %}
+                {%- endif %}
+                {{- '<tool_call>\n{"name": "' }}
+                {{- tool_call.name }}
+                {{- '", "arguments": ' }}
+                {%- if tool_call.arguments is string %}
+                    {{- tool_call.arguments }}
+                {%- else %}
+                    {{- tool_call.arguments | tojson }}
+                {%- endif %}
+                {{- '}\n</tool_call>' }}
+            {%- endfor %}
+        {%- endif %}
+        {{- '<|im_end|>\n' }}
+        {%- endgeneration %}
+    {%- elif message.role == "tool" %}
+        {%- if loop.first or (messages[loop.index0 - 1].role != "tool") %}
+            {{- '<|im_start|>user' }}
+        {%- endif %}
+        {{- '\n<tool_response>\n' }}
+        {{- content }}
+        {{- '\n</tool_response>' }}
+        {%- if loop.last or (messages[loop.index0 + 1].role != "tool") %}
+            {{- '<|im_end|>\n' }}
+        {%- endif %}
+    {%- endif %}
+{%- endfor %}
+{%- if add_generation_prompt %}
+    {{- '<|im_start|>assistant\n' }}
+    {%- if enable_thinking is defined and enable_thinking is false %}
+        {{- '<think>\n\n</think>\n\n' }}
+    {%- endif %}
+{%- endif %}

config.json

@@ -0,0 +1,75 @@
+{
+  "architectures": [
+    "Qwen3ForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "bos_token_id": null,
+  "dtype": "bfloat16",
+  "eos_token_id": 151645,
+  "head_dim": 128,
+  "hidden_act": "silu",
+  "hidden_size": 5120,
+  "initializer_range": 0.02,
+  "intermediate_size": 17408,
+  "layer_types": [
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention"
+  ],
+  "max_position_embeddings": 40960,
+  "max_window_layers": 40,
+  "model_type": "qwen3",
+  "num_attention_heads": 40,
+  "num_hidden_layers": 40,
+  "num_key_value_heads": 8,
+  "pad_token_id": 151643,
+  "rms_norm_eps": 1e-06,
+  "rope_parameters": {
+    "rope_theta": 1000000,
+    "rope_type": "default"
+  },
+  "sliding_window": null,
+  "tie_word_embeddings": false,
+  "transformers_version": "5.13.0.dev0",
+  "use_cache": false,
+  "use_sliding_window": false,
+  "vocab_size": 157184
+}

dna_config.json

@@ -0,0 +1,12 @@
+{
+  "k": 6,
+  "dna_start_id": 151669,
+  "dna_vocab_size": 5515,
+  "dna_special_tokens": [
+    "<dna>",
+    "</dna>",
+    "<oov>"
+  ],
+  "auto_dna_tags": true,
+  "tail": "pad_n"
+}

generation_config.json

@@ -0,0 +1,13 @@
+{
+  "do_sample": true,
+  "eos_token_id": [
+    151645,
+    151643
+  ],
+  "pad_token_id": 151643,
+  "temperature": 0.6,
+  "top_k": 20,
+  "top_p": 0.95,
+  "transformers_version": "5.13.0.dev0",
+  "trust_remote_code": true
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b44d3ffa2ef0b162f5efb61f4ae12999633201425cfb76a0f691ace04fed48c1
+size 29644145312

tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:be75606093db2094d7cd20f3c2f385c212750648bd6ea4fb2bf507a6a4c55506
+size 11422650

tokenizer.py

@@ -0,0 +1,835 @@
+"""
+HybridDNATokenizer: Combines Qwen3 BPE tokenization with DNA 6-mer tokenization.
+
+DNA spans are detected automatically or supplied in <dna>...</dna> tags and
+tokenized as 6-mers. All remaining text uses Qwen3's BPE tokenization.
+
+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 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
+"""
+
+import os
+import json
+import itertools
+import re
+from typing import List, Optional, Tuple, Dict, Union, Any
+
+from transformers import PreTrainedTokenizer, AutoTokenizer, BatchEncoding
+
+# The BPE vocabulary comes from this tokenizer, and its size sets where the DNA ids start. Pin it so that a change on the
+# Hub cannot silently shift the ids that checkpoints were trained with.
+BASE_TOKENIZER = "Qwen/Qwen3-4B-Base"
+BASE_TOKENIZER_REVISION = "906bfd4b4dc7f14ee4320094d8b41684abff8539"
+
+
+class HybridDNATokenizer(PreTrainedTokenizer):
+    """
+    Hybrid tokenizer combining Qwen3 BPE with DNA 6-mer tokenization.
+
+    Explicit <dna>...</dna> regions are always tokenized as 6-mers. By default,
+    untagged DNA spans embedded in text are detected and tagged automatically;
+    pass auto_dna_tags=False to encode() or __call__() to disable detection for
+    one call.
+    """
+
+    model_input_names = ["input_ids", "attention_mask"]
+
+    # Single-nucleotide alphabet for tail="char": standard bases (A C T G U) +
+    # IUPAC ambiguity codes (R Y S W K M B D H V N). These 16 chars match
+    # microsoft/NatureLM-8x7B nlm_tokenizer.py's single-char DNA set and the
+    # NatureLM paper's stated DNA vocabulary size of 16 (arXiv:2502.07527, S2.1).
+    SINGLE_BASE_ALPHABET = "ACTGURYSWKMBDHVN"
+    IUPAC_DNA_ALPHABET = "ACGTRYSWKMBDHVN"
+    STANDARD_DNA_ALPHABET = frozenset("ACGT")
+    IUPAC_RUN_RE = re.compile(f"[{IUPAC_DNA_ALPHABET}]+", re.IGNORECASE | re.ASCII)
+    SPACED_DNA_RE = re.compile(
+        r"(?<![A-Za-z])(?:[ACGT]{1,5}[ \t]+){5,}[ACGT]{1,5}(?![A-Za-z])"
+    )
+
+    def __init__(
+        self,
+        base_tokenizer_path: Optional[str] = None,
+        k: int = 6,
+        auto_dna_tags: bool = True,
+        tail: str = "pad_a",
+        **kwargs
+    ):
+        if tail not in ("pad_a", "pad_n", "char"):
+            raise ValueError(f"tail must be 'pad_a', 'pad_n' or 'char', got {tail!r}")
+        self.tail = tail
+        self.pad_char = {"pad_a": "A", "pad_n": "N"}.get(tail)
+        self.k = k
+
+        # Load base tokenizer (Qwen3-4B-Base) at the pinned revision
+        self._base_tokenizer = AutoTokenizer.from_pretrained(BASE_TOKENIZER, revision=BASE_TOKENIZER_REVISION)
+
+        # Get base vocabulary
+        self._base_vocab = self._base_tokenizer.get_vocab()
+        self._base_vocab_size = len(self._base_vocab)
+
+        # Initialize DNA vocabulary
+        self._init_dna_vocab()
+
+        # Build combined vocabulary
+        self._build_combined_vocab()
+
+        # Set special tokens
+        self._eos_token = kwargs.pop('eos_token', None) or "<|endoftext|>"
+        self._pad_token = kwargs.pop('pad_token', None) or self._base_tokenizer.pad_token or "<|endoftext|>"
+
+        # Initialize parent class
+        super().__init__(
+            eos_token=self._eos_token,
+            pad_token=self._pad_token,
+            **kwargs
+        )
+
+        # Preserve the base tokenizer's AddedToken metadata. These tokens are
+        # already present in the combined vocabulary at their original IDs, so
+        # this only registers their metadata and does not grow the vocabulary.
+        self.add_tokens(list(self._base_tokenizer.added_tokens_decoder.values()))
+
+        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)."""
+        bases = ['A', 'T', 'C', 'G']
+
+        # DNA special tokens
+        self.dna_special_tokens = ["<dna>", "</dna>", "<oov>"]
+
+        # Generate all k-mer combinations (4^k = 4096 for k=6)
+        self.kmers = [''.join(kmer) for kmer in itertools.product(bases, repeat=self.k)]
+
+        # tail="pad_n": dedicated tokens for the N-padded tail ("CT" -> "CTNNNN"),
+        # one per real prefix of length 1..k-1, i.e. 4^1+...+4^(k-1) = 1364.
+        self.npad_tokens = []
+        if self.tail == "pad_n":
+            for prefix_len in range(1, self.k):
+                for prefix in itertools.product(bases, repeat=prefix_len):
+                    self.npad_tokens.append(''.join(prefix) + 'N' * (self.k - prefix_len))
+
+        # tail="char": single-base tokens for the trailing < k bases (NatureLM
+        # IUPAC alphabet), so N / ambiguity codes fall back to a token, not <oov>.
+        self.base_tokens = []
+        self.dna_base_token = {}         # base -> token, e.g. "A" -> "<b_A>"
+        self.dna_base_token_to_base = {}  # token -> base
+        if self.tail == "char":
+            for b in self.SINGLE_BASE_ALPHABET:
+                tok = f"<b_{b}>"
+                self.base_tokens.append(tok)
+                self.dna_base_token[b] = tok
+                self.dna_base_token_to_base[tok] = b
+
+        # DNA tokens start after base vocabulary
+        self.dna_start_id = self._base_vocab_size
+
+        # All DNA tokens get new IDs (no reuse of base vocab IDs, even for
+        # overlapping tokens like CCCCCC — they have different semantics in
+        # DNA context vs BPE context, per Qiuyi's recommendation). Tail-variant
+        # tokens are listed AFTER the 4^k k-mers so the k-mer ids are unchanged.
+        base_dna_tokens = self.dna_special_tokens + self.kmers + self.npad_tokens + self.base_tokens
+
+        # Calculate padding for 128 alignment
+        total_vocab_unpadded = self._base_vocab_size + len(base_dna_tokens)
+        target_vocab_size = ((total_vocab_unpadded + 127) // 128) * 128
+        num_padding_tokens = target_vocab_size - total_vocab_unpadded
+
+        # Add unused padding tokens
+        self.padding_tokens = [f"<unused_{i}>" for i in range(num_padding_tokens)]
+
+        # Create DNA token mappings — all get sequential new IDs
+        self.dna_token_to_id = {}
+        self.dna_id_to_token = {}
+
+        current_id = self.dna_start_id
+        for token in base_dna_tokens:
+            self.dna_token_to_id[token] = current_id
+            self.dna_id_to_token[current_id] = token
+            current_id += 1
+
+        # Add padding tokens
+        for token in self.padding_tokens:
+            self.dna_token_to_id[token] = current_id
+            self.dna_id_to_token[current_id] = token
+            current_id += 1
+
+        self.dna_vocab_size = len(base_dna_tokens) + len(self.padding_tokens)
+
+        # Set DNA special token IDs
+        self.dna_begin_token_id = self.dna_token_to_id["<dna>"]
+        self.dna_end_token_id = self.dna_token_to_id["</dna>"]
+        self.oov_token_id = self.dna_token_to_id["<oov>"]
+
+    def _build_combined_vocab(self):
+        """Build combined vocabulary (base + DNA)."""
+        self._vocab = self._base_vocab.copy()
+
+        for token, token_id in self.dna_token_to_id.items():
+            if token not in self._vocab:
+                self._vocab[token] = token_id
+
+        self._id_to_token = {v: k for k, v in self._vocab.items()}
+        for token_id, token in self.dna_id_to_token.items():
+            if token_id not in self._id_to_token:
+                self._id_to_token[token_id] = token
+
+    @property
+    def vocab_size(self) -> int:
+        return max(self._vocab.values()) + 1
+
+    def get_vocab(self) -> Dict[str, int]:
+        return self._vocab.copy()
+
+    def __len__(self):
+        # Override default (len(get_vocab())) because get_vocab() deduplicates
+        # CCCCCC which exists as both BPE (ID 91443) and DNA 6-mer (ID 154402).
+        return self.vocab_size
+
+    @classmethod
+    def _auto_tag_text_segment(cls, text: str) -> str:
+        """Tag conservative DNA candidates in a text-only segment.
+
+        Each maximal IUPAC nucleotide run is considered independently. Pure
+        A/C/G/T runs need at least six bases. Runs containing IUPAC ambiguity
+        codes need at least twelve bases and may contain at most 20% ambiguity.
+        Rejected ambiguous runs stay wholly on the BPE path rather than having
+        any A/C/G/T substring tagged.
+        """
+        parts = []
+        previous_end = 0
+
+        for match in cls.IUPAC_RUN_RE.finditer(text):
+            candidate = match.group(0)
+            normalized = candidate.upper()
+            ambiguous_count = sum(base not in cls.STANDARD_DNA_ALPHABET for base in normalized)
+            is_standard_dna = ambiguous_count == 0 and len(normalized) >= 6
+            is_conservative_iupac = (
+                ambiguous_count > 0
+                and len(normalized) >= 12
+                and ambiguous_count / len(normalized) <= 0.2
+            )
+
+            if not (is_standard_dna or is_conservative_iupac):
+                continue
+
+            parts.append(text[previous_end:match.start()])
+            parts.append(f"<dna>{normalized}</dna>")
+            previous_end = match.end()
+
+        if not parts:
+            return text
+
+        parts.append(text[previous_end:])
+        return "".join(parts)
+
+    @classmethod
+    def _normalize_spaced_dna(cls, text: str) -> str:
+        """Join uppercase DNA displays with at least six chunks and twelve bases.
+
+        Chunks contain one to five A/C/G/T bases, separated only by spaces or
+        tabs. Alphabetic boundaries prevent joining fragments of words.
+        """
+        def normalize(match):
+            sequence = "".join(match.group(0).split())
+            return sequence if len(sequence) >= 12 else match.group(0)
+
+        return cls.SPACED_DNA_RE.sub(normalize, text)
+
+    def _add_auto_dna_tags(self, text: str) -> str:
+        """Tag DNA outside existing <dna>...</dna> regions."""
+        return "".join(
+            segment if is_dna else self._auto_tag_text_segment(self._normalize_spaced_dna(segment))
+            for segment, is_dna in self._split_by_dna_tags(text)
+        )
+
+    def _split_by_dna_tags(self, text: str) -> List[Tuple[str, bool]]:
+        segments = []
+        i = 0
+        n = len(text)
+
+        while i < n:
+            start_pos = text.find('<dna>', i)
+            end_pos = text.find('</dna>', i)
+
+            if start_pos == -1 and end_pos == -1:
+                remaining = text[i:]
+                if remaining:
+                    segments.append((remaining, False))
+                break
+
+            if start_pos == -1 and end_pos != -1:
+                dna_region = text[i:end_pos + 6]
+                if dna_region:
+                    segments.append((dna_region, True))
+                i = end_pos + 6
+                continue
+
+            if start_pos != -1 and end_pos == -1:
+                if i < start_pos:
+                    normal_text = text[i:start_pos]
+                    if normal_text:
+                        segments.append((normal_text, False))
+                dna_region = text[start_pos:]
+                if dna_region:
+                    segments.append((dna_region, True))
+                break
+
+            if start_pos < end_pos:
+                if i < start_pos:
+                    normal_text = text[i:start_pos]
+                    if normal_text:
+                        segments.append((normal_text, False))
+                dna_region = text[start_pos:end_pos + 6]
+                if dna_region:
+                    segments.append((dna_region, True))
+                i = end_pos + 6
+            else:
+                dna_region = text[i:end_pos + 6]
+                if dna_region:
+                    segments.append((dna_region, True))
+                i = end_pos + 6
+
+        return segments
+
+    def _parse_dna_region(self, dna_region: str) -> Tuple[str, bool, bool]:
+        if dna_region == '<dna>':
+            return '', True, False
+        elif dna_region == '</dna>':
+            return '', False, True
+
+        has_start = dna_region.startswith('<dna>')
+        has_end = dna_region.endswith('</dna>')
+
+        content = dna_region
+        if has_start:
+            content = content[5:]
+        if has_end and content.endswith('</dna>'):
+            content = content[:-6]
+
+        return content.strip(), has_start, has_end
+
+    def _process_dna_sequence(self, dna_seq: str) -> Dict:
+        k = self.k
+        dna_seq = dna_seq.upper()
+
+        kmer_tokens = []
+        valid_bases = set('ATCG')
+
+        def is_valid_kmer(kmer):
+            return len(kmer) == k and all(base in valid_bases for base in kmer)
+
+        if self.tail == "char":
+            # Greedy NatureLM tokenization (NatureLM, arXiv:2502.07527), ported
+            # from microsoft/NatureLM-8x7B nlm_tokenizer.py::_tokenize_entity:
+            # take a full k-mer when the next k chars form one, else emit a single
+            # base token and advance by 1 (also re-aligns after any non-base char).
+            i = 0
+            n = len(dna_seq)
+            while i < n:
+                window = dna_seq[i:i + k]
+                if is_valid_kmer(window):
+                    kmer_tokens.append(window)
+                    i += k
+                else:
+                    base = dna_seq[i]
+                    kmer_tokens.append(self.dna_base_token.get(base, "<oov>"))
+                    i += 1
+            return {"kmer_tokens": kmer_tokens, "padding_length": 0, "valid_length": k}
+
+        for i in range(0, len(dna_seq) - k + 1, k):
+            kmer = dna_seq[i:i+k]
+            if is_valid_kmer(kmer):
+                kmer_tokens.append(kmer)
+            else:
+                kmer_tokens.append("<oov>")
+
+        processed_length = len(kmer_tokens) * k
+        remaining = dna_seq[processed_length:]
+        padding_length = 0
+        valid_length = k
+
+        if remaining:
+            padding_needed = k - len(remaining)
+            # Right-pad up to a full k-mer: 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. pad_char="A" reuses a real k-mer id ("CT"->"CTAAAA");
+            # "N" maps to a dedicated tail token ("CT"->"CTNNNN").
+            padded = remaining + self.pad_char * padding_needed
+
+            if all(base in valid_bases for base in remaining):
+                kmer_tokens.append(padded)
+            else:
+                kmer_tokens.append("<oov>")
+
+            padding_length = padding_needed
+            valid_length = len(remaining)
+
+        return {
+            "kmer_tokens": kmer_tokens,
+            "padding_length": padding_length,
+            "valid_length": valid_length,
+        }
+
+    def _tokenize(self, text: str, **kwargs) -> List[str]:
+        return list(text)
+
+    def _convert_token_to_id(self, token: str) -> int:
+        if token in self.dna_token_to_id:
+            return self.dna_token_to_id[token]
+        return self._base_vocab.get(token, self._base_tokenizer.unk_token_id or 0)
+
+    def _convert_id_to_token(self, index: int) -> str:
+        if index in self.dna_id_to_token:
+            return self.dna_id_to_token[index]
+        return self._id_to_token.get(index, "<oov>")
+
+    def convert_tokens_to_string(self, tokens: List[str]) -> str:
+        ids = [self._convert_token_to_id(t) for t in tokens]
+        return self.decode(ids, skip_special_tokens=False)
+    def encode(
+        self,
+        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:
+            text = self._add_auto_dna_tags(text)
+
+        segments = self._split_by_dna_tags(text)
+
+        token_ids = []
+        token_mask = [] if return_token_mask else None
+
+        for segment_content, is_dna in segments:
+            if is_dna:
+                dna_content, has_start, has_end = self._parse_dna_region(segment_content)
+
+                if has_start:
+                    token_ids.append(self.dna_begin_token_id)
+                    if return_token_mask:
+                        token_mask.append(0)
+
+                if dna_content:
+                    result = self._process_dna_sequence(dna_content)
+
+                    for idx, kmer in enumerate(result["kmer_tokens"]):
+                        token_id = self.dna_token_to_id.get(kmer, self.oov_token_id)
+                        token_ids.append(token_id)
+
+                        if return_token_mask:
+                            if kmer == "<oov>":
+                                token_mask.append(0)
+                            elif kmer in self.dna_base_token_to_base:
+                                token_mask.append(1)  # char-fallback single base
+                            elif idx == len(result["kmer_tokens"]) - 1 and result["padding_length"] > 0:
+                                token_mask.append(result["valid_length"])
+                            else:
+                                token_mask.append(self.k)
+
+                if has_end:
+                    token_ids.append(self.dna_end_token_id)
+                    if return_token_mask:
+                        token_mask.append(0)
+            else:
+                base_ids = self._base_tokenizer.encode(
+                    segment_content,
+                    add_special_tokens=False
+                )
+                token_ids.extend(base_ids)
+                if return_token_mask:
+                    token_mask.extend([-1] * len(base_ids))
+
+        # 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
+        return token_ids
+
+    def _decode_dna_token(self, token_str: str) -> str:
+        # Surface text for a DNA-vocab token. Baseline (pad_a) returns it
+        # unchanged; char maps "<b_A>"->"A"; pad_n strips the N-pad tail fill
+        # ("CTNNNN"->"CT"). Real k-mers contain no N, so rstrip is a no-op there.
+        if token_str in self.dna_base_token_to_base:
+            return self.dna_base_token_to_base[token_str]
+        if self.tail == "pad_n" and len(token_str) == self.k and token_str[0] in "ATCG":
+            return token_str.rstrip("N")
+        return token_str
+
+    def decode(
+        self,
+        token_ids: Union[int, List[int]],
+        skip_special_tokens: bool = False,
+        **kwargs
+    ) -> str:
+        # A raw tensor/ndarray must become plain ints: dna_id_to_token is keyed by
+        # python ints and a 0-dim tensor scalar hashes by identity, so DNA lookups
+        # silently miss and the bases are dropped (e.g. decode(generate(...)[0])).
+        if hasattr(token_ids, "tolist"):
+            token_ids = token_ids.tolist()
+        if isinstance(token_ids, int):
+            token_ids = [token_ids]
+
+        if skip_special_tokens:
+            special_ids = {self.eos_token_id, self.pad_token_id}
+            token_ids = [tid for tid in token_ids if tid not in special_ids]
+
+        parts = []
+        i = 0
+
+        while i < len(token_ids):
+            tid = token_ids[i]
+
+            if tid == self.dna_begin_token_id:
+                dna_tokens = []
+                i += 1
+
+                while i < len(token_ids) and token_ids[i] != self.dna_end_token_id:
+                    if token_ids[i] in self.dna_id_to_token:
+                        dna_tokens.append(self._decode_dna_token(self.dna_id_to_token[token_ids[i]]))
+                    i += 1
+
+                dna_seq = ''.join(dna_tokens)
+
+                if skip_special_tokens:
+                    parts.append(dna_seq)
+                else:
+                    parts.append(f"<dna>{dna_seq}")
+                    if i < len(token_ids) and token_ids[i] == self.dna_end_token_id:
+                        parts.append("</dna>")
+                        i += 1
+
+            elif tid in self.dna_id_to_token:
+                # 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._decode_dna_token(self.dna_id_to_token[tid]))
+                i += 1
+
+            else:
+                text_ids = []
+                while i < len(token_ids):
+                    curr_id = token_ids[i]
+                    if curr_id in self.dna_id_to_token or curr_id == self.dna_begin_token_id:
+                        break
+                    text_ids.append(curr_id)
+                    i += 1
+
+                if text_ids:
+                    decoded = self._base_tokenizer.decode(text_ids, skip_special_tokens=skip_special_tokens)
+                    parts.append(decoded)
+
+        return ''.join(parts)
+
+    def batch_decode(
+        self,
+        sequences: Union[List[int], List[List[int]], "torch.Tensor"],
+        skip_special_tokens: bool = False,
+        **kwargs
+    ) -> List[str]:
+        return [
+            self.decode(
+                seq.tolist() if hasattr(seq, 'tolist') else list(seq),
+                skip_special_tokens=skip_special_tokens,
+                **kwargs
+            )
+            for seq in sequences
+        ]
+
+    def __call__(
+        self,
+        text: Union[str, List[str]],
+        add_special_tokens: bool = False,
+        padding: bool = False,
+        truncation: bool = False,
+        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)
+        texts = text if is_batch else [text]
+
+        all_ids = []
+        all_masks = [] if return_token_mask else None
+
+        for t in texts:
+            if return_token_mask:
+                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, auto_dna_tags=auto_dna_tags)
+                all_ids.append(ids)
+
+        if padding:
+            max_len = max(len(ids) for ids in all_ids)
+            if max_length:
+                max_len = min(max_len, max_length)
+
+            padded_ids = []
+            attention_masks = []
+            padded_token_masks = [] if return_token_mask else None
+
+            for idx, ids in enumerate(all_ids):
+                pad_len = max_len - len(ids)
+
+                if pad_len > 0:
+                    ids = ids + [self.pad_token_id] * pad_len
+                    attn = [1] * (max_len - pad_len) + [0] * pad_len
+                    if return_token_mask:
+                        mask = all_masks[idx] + [-2] * pad_len
+                else:
+                    ids = ids[:max_len]
+                    attn = [1] * max_len
+                    if return_token_mask:
+                        mask = all_masks[idx][:max_len]
+
+                padded_ids.append(ids)
+                attention_masks.append(attn)
+                if return_token_mask:
+                    padded_token_masks.append(mask)
+
+            all_ids = padded_ids
+            all_masks = padded_token_masks
+        else:
+            attention_masks = [[1] * len(ids) for ids in all_ids]
+
+        result = {
+            "input_ids": all_ids if is_batch else all_ids[0],
+            "attention_mask": attention_masks if is_batch else attention_masks[0],
+        }
+
+        if return_token_mask:
+            result["token_mask"] = all_masks if is_batch else all_masks[0]
+
+        if return_tensors == "pt":
+            import torch
+            if is_batch:
+                result["input_ids"] = torch.tensor(result["input_ids"])
+                result["attention_mask"] = torch.tensor(result["attention_mask"])
+                if return_token_mask:
+                    result["token_mask"] = torch.tensor(result["token_mask"])
+            else:
+                result["input_ids"] = torch.tensor([result["input_ids"]])
+                result["attention_mask"] = torch.tensor([result["attention_mask"]])
+                if return_token_mask:
+                    result["token_mask"] = torch.tensor([result["token_mask"]])
+
+        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|>\n"
+
+    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.
+        """
+        auto_dna_tags = kwargs.pop("auto_dna_tags", None)
+        if auto_dna_tags is not None:
+            tokenizer_kwargs = dict(kwargs.pop("tokenizer_kwargs", None) or {})
+            tokenizer_kwargs["auto_dna_tags"] = auto_dna_tags
+            kwargs["tokenizer_kwargs"] = tokenizer_kwargs
+
+        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,
+            (filename_prefix + "-" if filename_prefix else "") + "vocab.json"
+        )
+
+        with open(vocab_file, "w", encoding="utf-8") as f:
+            json.dump(self._vocab, f, ensure_ascii=False, indent=2)
+
+        return (vocab_file,)
+
+    def save_pretrained(self, save_directory: str, **kwargs):
+        os.makedirs(save_directory, exist_ok=True)
+
+        # Save base tokenizer files with the wrapper's effective chat template.
+        self._base_tokenizer.chat_template = self.chat_template
+        self._base_tokenizer.save_pretrained(save_directory)
+
+        # Save DNA config
+        dna_config = {
+            "k": self.k,
+            "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,
+            "tail": self.tail,
+        }
+
+        dna_config_path = os.path.join(save_directory, "dna_config.json")
+        with open(dna_config_path, "w", encoding="utf-8") as f:
+            json.dump(dna_config, f, indent=2)
+
+        # Update tokenizer_config.json with auto_map
+        config_path = os.path.join(save_directory, "tokenizer_config.json")
+
+        if os.path.exists(config_path):
+            with open(config_path, "r") as f:
+                config = json.load(f)
+        else:
+            config = {}
+
+        config.update({
+            "tokenizer_class": "HybridDNATokenizer",
+            "auto_map": {
+                "AutoTokenizer": ["tokenizer.HybridDNATokenizer", None]
+            },
+            "k": self.k,
+            "auto_dna_tags": self.auto_dna_tags,
+            "tail": self.tail,
+        })
+
+        with open(config_path, "w", encoding="utf-8") as f:
+            json.dump(config, f, indent=2, ensure_ascii=False)
+
+        # Copy this tokenizer.py to save directory
+        import shutil
+        src_py = os.path.abspath(__file__)
+        dst_py = os.path.join(save_directory, "tokenizer.py")
+        if os.path.exists(src_py) and src_py != dst_py:
+            shutil.copy2(src_py, dst_py)
+
+        return (save_directory,)
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        pretrained_model_name_or_path,
+        *args,
+        cache_dir=None,
+        force_download=False,
+        local_files_only=False,
+        token=None,
+        revision="main",
+        trust_remote_code=False,
+        **kwargs,
+    ):
+        # Custom loader: PreTrainedTokenizer.from_pretrained is incompatible with this
+        # tokenizer's __init__ (it re-passes bos_token etc.), so construct directly from
+        # bio_config/dna_config. That bypasses standard metadata loading, so also read
+        # chat_template from tokenizer_config.json (vLLM / apply_chat_template need it).
+        from transformers.utils.hub import cached_file
+
+        proxies = kwargs.pop("proxies", None)
+        subfolder = kwargs.pop("subfolder", None)
+        commit_hash = kwargs.pop("_commit_hash", None)
+        hub_kwargs = {
+            "cache_dir": cache_dir,
+            "force_download": force_download,
+            "local_files_only": local_files_only,
+            "proxies": proxies,
+            "revision": revision,
+            "subfolder": subfolder,
+            "token": token,
+            "_commit_hash": commit_hash,
+            "_raise_exceptions_for_missing_entries": False,
+        }
+
+        cfg = {}
+        for _name in ("bio_config.json", "dna_config.json"):
+            _p = cached_file(pretrained_model_name_or_path, _name, **hub_kwargs)
+            if _p is not None:
+                with open(_p, encoding="utf-8") as _f:
+                    cfg = json.load(_f)
+                break
+
+        _init = {"base_tokenizer_path": pretrained_model_name_or_path}
+        for _key in ("k", "tail", "auto_dna_tags"):
+            if _key in cfg:
+                _init[_key] = cfg[_key]
+
+        tokenizer_config_path = cached_file(pretrained_model_name_or_path, "tokenizer_config.json", **hub_kwargs)
+        if tokenizer_config_path is not None:
+            with open(tokenizer_config_path, encoding="utf-8") as _f:
+                tokenizer_config = json.load(_f)
+            if tokenizer_config.get("chat_template"):
+                _init["chat_template"] = tokenizer_config["chat_template"]
+
+        chat_template_path = cached_file(pretrained_model_name_or_path, "chat_template.jinja", **hub_kwargs)
+        if chat_template_path is not None:
+            with open(chat_template_path, encoding="utf-8") as _f:
+                _init["chat_template"] = _f.read()
+
+        _init.update(kwargs)
+        tokenizer = cls(**_init)
+        # DNA ids start right after the base vocabulary, so a layout that differs from the saved one would misalign every
+        # DNA embedding. Fail instead of tokenizing with shifted ids.
+        for _key in ("dna_start_id", "dna_vocab_size"):
+            if _key in cfg and cfg[_key] != getattr(tokenizer, _key):
+                raise ValueError(
+                    f"{pretrained_model_name_or_path} was saved with {_key}={cfg[_key]}, but the base tokenizer gives "
+                    f"{getattr(tokenizer, _key)}, so its DNA token ids would not match the model."
+                )
+        return tokenizer

tokenizer_config.json

@@ -0,0 +1,39 @@
+{
+  "add_prefix_space": false,
+  "backend": "tokenizers",
+  "bos_token": null,
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|endoftext|>",
+  "errors": "replace",
+  "extra_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|object_ref_start|>",
+    "<|object_ref_end|>",
+    "<|box_start|>",
+    "<|box_end|>",
+    "<|quad_start|>",
+    "<|quad_end|>",
+    "<|vision_start|>",
+    "<|vision_end|>",
+    "<|vision_pad|>",
+    "<|image_pad|>",
+    "<|video_pad|>"
+  ],
+  "is_local": false,
+  "local_files_only": false,
+  "model_max_length": 131072,
+  "pad_token": "<|endoftext|>",
+  "split_special_tokens": false,
+  "tokenizer_class": "HybridDNATokenizer",
+  "unk_token": null,
+  "auto_map": {
+    "AutoTokenizer": [
+      "tokenizer.HybridDNATokenizer",
+      null
+    ]
+  },
+  "k": 6,
+  "auto_dna_tags": true,
+  "tail": "pad_n"
+}