Promote hybrid step-286000 to main (300B CE + 300B FNS, total 600B tokens)

.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

added_tokens.json

@@ -0,0 +1,28 @@
+{
+  "</think>": 151668,
+  "</tool_call>": 151658,
+  "</tool_response>": 151666,
+  "<think>": 151667,
+  "<tool_call>": 151657,
+  "<tool_response>": 151665,
+  "<|box_end|>": 151649,
+  "<|box_start|>": 151648,
+  "<|endoftext|>": 151643,
+  "<|file_sep|>": 151664,
+  "<|fim_middle|>": 151660,
+  "<|fim_pad|>": 151662,
+  "<|fim_prefix|>": 151659,
+  "<|fim_suffix|>": 151661,
+  "<|im_end|>": 151645,
+  "<|im_start|>": 151644,
+  "<|image_pad|>": 151655,
+  "<|object_ref_end|>": 151647,
+  "<|object_ref_start|>": 151646,
+  "<|quad_end|>": 151651,
+  "<|quad_start|>": 151650,
+  "<|repo_name|>": 151663,
+  "<|video_pad|>": 151656,
+  "<|vision_end|>": 151653,
+  "<|vision_pad|>": 151654,
+  "<|vision_start|>": 151652
+}

chat_template.jinja

@@ -0,0 +1,85 @@
+{%- 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 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.role == "user") or (message.role == "system" and not loop.first) %}
+        {{- '<|im_start|>' + message.role + '\n' + message.content + '<|im_end|>' + '\n' }}
+    {%- elif message.role == "assistant" %}
+        {%- set content = message.content %}
+        {%- set reasoning_content = '' %}
+        {%- if message.reasoning_content is defined and message.reasoning_content is not none %}
+            {%- set reasoning_content = message.reasoning_content %}
+        {%- else %}
+            {%- if '</think>' in message.content %}
+                {%- set content = message.content.split('</think>')[-1].lstrip('\n') %}
+                {%- set reasoning_content = message.content.split('</think>')[0].rstrip('\n').split('<think>')[-1].lstrip('\n') %}
+            {%- endif %}
+        {%- endif %}
+        {%- if loop.index0 > ns.last_query_index %}
+            {%- if loop.last or (not loop.last and reasoning_content) %}
+                {{- '<|im_start|>' + message.role + '\n<think>\n' + reasoning_content.strip('\n') + '\n</think>\n\n' + content.lstrip('\n') }}
+            {%- else %}
+                {{- '<|im_start|>' + message.role + '\n' + content }}
+            {%- endif %}
+        {%- else %}
+            {{- '<|im_start|>' + message.role + '\n' + content }}
+        {%- endif %}
+        {%- 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' }}
+    {%- elif message.role == "tool" %}
+        {%- if loop.first or (messages[loop.index0 - 1].role != "tool") %}
+            {{- '<|im_start|>user' }}
+        {%- endif %}
+        {{- '\n<tool_response>\n' }}
+        {{- message.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,29 @@
+{
+  "architectures": [
+    "LlamaForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "bos_token_id": 1,
+  "dtype": "float32",
+  "eos_token_id": 2,
+  "head_dim": 64,
+  "hidden_act": "silu",
+  "hidden_size": 1024,
+  "initializer_range": 0.02,
+  "intermediate_size": 3072,
+  "max_position_embeddings": 8192,
+  "mlp_bias": false,
+  "model_type": "llama",
+  "num_attention_heads": 16,
+  "num_hidden_layers": 28,
+  "num_key_value_heads": 8,
+  "pretraining_tp": 1,
+  "rms_norm_eps": 1e-06,
+  "rope_scaling": null,
+  "rope_theta": 500000.0,
+  "tie_word_embeddings": true,
+  "transformers_version": "4.57.6",
+  "use_cache": true,
+  "vocab_size": 155776
+}

dna_config.json

@@ -0,0 +1,10 @@
+{
+  "k": 6,
+  "dna_start_id": 151669,
+  "dna_vocab_size": 4107,
+  "dna_special_tokens": [
+    "<dna>",
+    "</dna>",
+    "<oov>"
+  ]
+}

generation_config.json

@@ -0,0 +1,6 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 1,
+  "eos_token_id": 2,
+  "transformers_version": "4.57.6"
+}

model.safetensors

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

special_tokens_map.json

@@ -0,0 +1,31 @@
+{
+  "additional_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|>"
+  ],
+  "eos_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer.json

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

tokenizer.py

@@ -0,0 +1,551 @@
+"""
+HybridDNATokenizer: Combines Qwen3 BPE tokenization with DNA 6-mer tokenization.
+
+DNA sequences wrapped in <dna>...</dna> tags are tokenized as 6-mers.
+All other 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 (number of valid bases)
+   6: full 6-mer
+"""
+
+import os
+import json
+import itertools
+from typing import List, Optional, Tuple, Dict, Union, Any
+
+from transformers import PreTrainedTokenizer, AutoTokenizer, BatchEncoding
+
+
+class HybridDNATokenizer(PreTrainedTokenizer):
+    """
+    Hybrid tokenizer combining Qwen3 BPE with DNA 6-mer tokenization.
+
+    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.
+    """
+
+    model_input_names = ["input_ids", "attention_mask"]
+
+    def __init__(
+        self,
+        base_tokenizer_path: Optional[str] = None,
+        k: int = 6,
+        **kwargs
+    ):
+        self.k = k
+
+        # Load base tokenizer (Qwen3-4B-Base)
+        self._base_tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-4B-Base")
+
+        # 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
+        )
+
+        self.special_tokens = self.dna_special_tokens + [self._eos_token, self._pad_token]
+
+    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)]
+
+        # 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)
+        base_dna_tokens = self.dna_special_tokens + self.kmers
+
+        # 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
+
+    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)
+
+        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)
+            padded = remaining + 'A' * padding_needed
+
+            if is_valid_kmer(padded):
+                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:
+        return "".join(tokens)
+
+    def encode(
+        self,
+        text: str,
+        add_special_tokens: bool = False,
+        return_token_mask: bool = False,
+        **kwargs
+    ) -> Union[List[int], Tuple[List[int], List[int]]]:
+        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 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))
+
+        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)
+
+        if return_token_mask:
+            return token_ids, token_mask
+        return token_ids
+
+    def decode(
+        self,
+        token_ids: Union[int, List[int]],
+        skip_special_tokens: bool = False,
+        **kwargs
+    ) -> str:
+        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.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:
+                if not skip_special_tokens:
+                    parts.append(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,
+        **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)
+                all_ids.append(ids)
+                all_masks.append(mask)
+            else:
+                ids = self.encode(t, add_special_tokens=add_special_tokens, return_token_mask=False)
+                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)
+
+    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
+        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,
+        }
+
+        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,
+        })
+
+        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: str, **kwargs):
+        dna_config_path = os.path.join(pretrained_model_name_or_path, "dna_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)

tokenizer_config.json

@@ -0,0 +1,246 @@
+{
+  "add_bos_token": false,
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "151643": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151644": {
+      "content": "<|im_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151645": {
+      "content": "<|im_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151646": {
+      "content": "<|object_ref_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151647": {
+      "content": "<|object_ref_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151648": {
+      "content": "<|box_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151649": {
+      "content": "<|box_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151650": {
+      "content": "<|quad_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151651": {
+      "content": "<|quad_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151652": {
+      "content": "<|vision_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151653": {
+      "content": "<|vision_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151654": {
+      "content": "<|vision_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151655": {
+      "content": "<|image_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151656": {
+      "content": "<|video_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151657": {
+      "content": "<tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151658": {
+      "content": "</tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151659": {
+      "content": "<|fim_prefix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151660": {
+      "content": "<|fim_middle|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151661": {
+      "content": "<|fim_suffix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151662": {
+      "content": "<|fim_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151663": {
+      "content": "<|repo_name|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151664": {
+      "content": "<|file_sep|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151665": {
+      "content": "<tool_response>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151666": {
+      "content": "</tool_response>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151667": {
+      "content": "<think>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151668": {
+      "content": "</think>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    }
+  },
+  "additional_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|>"
+  ],
+  "bos_token": null,
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|endoftext|>",
+  "errors": "replace",
+  "extra_special_tokens": {},
+  "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
+}