Fix tokenizer persistence using merged source implementation

tokenizer.py

@@ -80,20 +80,28 @@ class HybridDNATokenizer(PreTrainedTokenizer):
         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|>"
+        self._eos_token = kwargs.pop("eos_token", "<|endoftext|>")
+        self._pad_token = kwargs.pop("pad_token", self._base_tokenizer.pad_token or "<|endoftext|>")
 
         # Initialize parent class
         super().__init__(
             eos_token=self._eos_token,
             pad_token=self._pad_token,
+            k=k,
+            tail=tail,
+            auto_dna_tags=auto_dna_tags,
             **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()))
+        # Keep any metadata restored by the parent loader or supplied by the caller.
+        existing_ids = self.added_tokens_decoder
+        self.add_tokens([
+            existing_ids.get(token_id, token)
+            for token_id, token in self._base_tokenizer.added_tokens_decoder.items()
+        ])
 
         self.special_tokens = self.dna_special_tokens + [self._eos_token, self._pad_token]
         self.auto_dna_tags = auto_dna_tags
@@ -711,14 +719,21 @@ class HybridDNATokenizer(PreTrainedTokenizer):
 
         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
+    def _save_pretrained(
+        self,
+        save_directory,
+        file_names,
+        legacy_format=None,
+        filename_prefix=None,
+    ):
+        # The parent saves this wrapper's configuration and special-token metadata.
+        # Delegate only the BPE assets to the underlying tokenizer's save hook.
+        file_names = self._base_tokenizer._save_pretrained(
+            save_directory,
+            file_names,
+            legacy_format=legacy_format,
+            filename_prefix=filename_prefix,
+        )
         dna_config = {
             "k": self.k,
             "dna_start_id": self.dna_start_id,
@@ -727,103 +742,31 @@ class HybridDNATokenizer(PreTrainedTokenizer):
             "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,)
+        return file_names + (dna_config_path,)
 
     @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).
+    def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs):
+        tokenizer = super().from_pretrained(pretrained_model_name_or_path, *args, **kwargs)
+
         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,
+            key: kwargs[key]
+            for key in ("cache_dir", "force_download", "local_files_only", "token", "revision", "subfolder", "_commit_hash")
+            if key in kwargs
         }
-
+        config_path = cached_file(
+            pretrained_model_name_or_path, "dna_config.json",
+            _raise_exceptions_for_missing_entries=False, **hub_kwargs,
+        )
         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
+        if config_path is not None:
+            with open(config_path, encoding="utf-8") as f:
+                cfg = json.load(f)
 
-        _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"):
@@ -833,3 +776,7 @@ class HybridDNATokenizer(PreTrainedTokenizer):
                     f"{getattr(tokenizer, _key)}, so its DNA token ids would not match the model."
                 )
         return tokenizer
+
+
+# Include the custom implementation and AutoTokenizer mapping in saved checkpoints.
+HybridDNATokenizer.register_for_auto_class("AutoTokenizer")