Upload folder using huggingface_hub

export_to_hf.py

@@ -26,6 +26,46 @@ def infer_special_token_paths(repo_dir):
     return repo_dir / "tokenizer" / "tokenizer.model", subdir_metadata, repo_dir / "tokenizer" / "tokenizer.vocab"
 
 
+def write_clean_tokenizer_metadata(repo_dir, special_tokens):
+    tokenizer_config = {
+        "backend": "custom",
+        "bos_token": "<BOS>",
+        "eos_token": "<EOS>",
+        "unk_token": "<UNK>",
+        "pad_token": "<PAD>",
+        "tokenizer_class": "TaoNetTokenizer",
+        "model_max_length": 1000000000000000019884624838656,
+        "extra_special_tokens": [
+            token
+            for token in special_tokens
+            if token not in {"<UNK>", "<BOS>", "<EOS>", "<PAD>"}
+        ],
+    }
+    special_tokens_map = {
+        "unk_token": "<UNK>",
+        "bos_token": "<BOS>",
+        "eos_token": "<EOS>",
+        "pad_token": "<PAD>",
+        "additional_special_tokens": [
+            token
+            for token in special_tokens
+            if token not in {"<UNK>", "<BOS>", "<EOS>", "<PAD>"}
+        ],
+    }
+
+    with open(repo_dir / "tokenizer_config.json", "w", encoding="utf-8") as handle:
+        json.dump(tokenizer_config, handle, indent=2)
+        handle.write("\n")
+
+    with open(repo_dir / "special_tokens_map.json", "w", encoding="utf-8") as handle:
+        json.dump(special_tokens_map, handle, indent=2)
+        handle.write("\n")
+
+    added_tokens_path = repo_dir / "added_tokens.json"
+    if added_tokens_path.exists():
+        added_tokens_path.unlink()
+
+
 def main():
     import torch
 
@@ -69,6 +109,7 @@ def main():
         special_tokens_file=str(metadata_path),
     )
     tokenizer.save_pretrained(repo_dir)
+    write_clean_tokenizer_metadata(repo_dir, metadata.get("configured_special_tokens", []))
 
     shutil.copyfile(metadata_model_path, repo_dir / "tokenizer.model")
     shutil.copyfile(metadata_path, repo_dir / "tokenizer.special_tokens.json")

modeling_taonet.py

@@ -4,8 +4,12 @@ from torch import nn
 from transformers import GenerationMixin, PreTrainedModel
 from transformers.modeling_outputs import CausalLMOutput
 
-from .configuration_taonet import TaoNetConfig
-from .taonet_model import SimpleLLM, build_runtime_config
+try:
+    from .configuration_taonet import TaoNetConfig
+    from .taonet_model import SimpleLLM, build_runtime_config
+except ImportError:
+    from configuration_taonet import TaoNetConfig
+    from taonet_model import SimpleLLM, build_runtime_config
 
 
 class TaoNetForCausalLM(PreTrainedModel, GenerationMixin):

taonet_model.py

@@ -6,8 +6,12 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
-from .embeddings import FactorizedEmbedding
-from .mla_components import AttentionBlock
+try:
+    from .embeddings import FactorizedEmbedding
+    from .mla_components import AttentionBlock
+except ImportError:
+    from embeddings import FactorizedEmbedding
+    from mla_components import AttentionBlock
 
 
 class SimpleLLM(nn.Module):

tokenization_taonet.py

@@ -43,6 +43,9 @@ class TaoNetTokenizer(PreTrainedTokenizer):
                 str(token): int(token_id) for token, token_id in metadata.get("special_tokens", {}).items()
             }
             configured_special_tokens = [str(token) for token in metadata.get("configured_special_tokens", [])]
+        self.id_to_special_token = {
+            int(token_id): str(token) for token, token_id in self.special_token_ids.items()
+        }
 
         merged_additional_tokens = list(additional_special_tokens or [])
         for token in configured_special_tokens:
@@ -64,6 +67,7 @@ class TaoNetTokenizer(PreTrainedTokenizer):
 
     def get_vocab(self):
         vocab = {self.sp_model.id_to_piece(i): i for i in range(self.vocab_size)}
+        vocab.update(self.special_token_ids)
         vocab.update(self.added_tokens_encoder)
         return vocab
 
@@ -79,6 +83,8 @@ class TaoNetTokenizer(PreTrainedTokenizer):
         return int(piece_id)
 
     def _convert_id_to_token(self, index):
+        if index in self.id_to_special_token:
+            return self.id_to_special_token[index]
         if index in self.added_tokens_decoder:
             return self.added_tokens_decoder[index].content
         return self.sp_model.id_to_piece(int(index))

tokenizer_config.json

@@ -1,90 +1,16 @@
 {
-  "added_tokens_decoder": {
-    "4": {
-      "content": "\n",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "5": {
-      "content": "<think>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "6": {
-      "content": "<user>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "7": {
-      "content": "<assistant>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "8": {
-      "content": "<image>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "8192": {
-      "content": "<BOS>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "8193": {
-      "content": "<EOS>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "8194": {
-      "content": "<UNK>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "8195": {
-      "content": "<PAD>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    }
-  },
   "backend": "custom",
   "bos_token": "<BOS>",
   "eos_token": "<EOS>",
+  "unk_token": "<UNK>",
+  "pad_token": "<PAD>",
+  "tokenizer_class": "TaoNetTokenizer",
+  "model_max_length": 1000000000000000019884624838656,
   "extra_special_tokens": [
     "\n",
     "<think>",
     "<user>",
     "<assistant>",
     "<image>"
-  ],
-  "model_max_length": 1000000000000000019884624838656,
-  "pad_token": "<PAD>",
-  "tokenizer_class": "TaoNetTokenizer",
-  "unk_token": "<UNK>"
+  ]
 }