Upload folder using huggingface_hub

src/taoTrain/inference/inferencer.py

@@ -464,8 +464,10 @@ class Inferencer:
         else:
             input_ids = self._encode_prompt(prompt)
 
+        prompt_token_ids = input_ids.squeeze(0).tolist()
         prompt_length = input_ids.shape[1]
-        special_token_id_values = set(self._get_special_token_ids().values())
+        special_token_ids = self._get_special_token_ids()
+        special_token_id_values = set(special_token_ids.values())
         generated_token_ids: list[int] = []
         stop_reason = "max_length"
 

src/taoTrain/models/taonet.py

@@ -137,7 +137,7 @@ class SimpleLLM(BaseModel):
         trainable_params = sum(p.numel() for p in self.parameters() if p.requires_grad)
         
         print(f"\n{'='*70}")
-        print("MODEL ARCHITECTURE - TAОNET (DeepSeek MLA + RoPE)")
+        print("MODEL ARCHITECTURE - TAONET (DeepSeek MLA + RoPE)")
         print(f"{'='*70}")
         print(f"Embedding:")
         if self.use_factorized_embedding:
@@ -153,10 +153,10 @@ class SimpleLLM(BaseModel):
         
         output_params = self.d_model * self.vocab_size
         print(f"Output Head:       {output_params/1e6:>8.2f}M")
-        print(f"Attention Blocks:  {len(self.blocks):>10} layers × AttentionBlock")
-        print(f"{'─'*70}")
-        print(f"Total Parameters:  {total_params/1e6:>8.2f}M (trainable: {trainable_params/1e6:.2f}M)")
-        print(f"{'─'*70}")
+        print(f"Attention Blocks:  {len(self.blocks):>10} layers x AttentionBlock")
+        print(f"{'-'*70}")
+        print(f"Total Parameters:  {total_params/1e6:>8.2f}M (trainable: {trainable_params/1e6:.2f}M)")
+        print(f"{'-'*70}")
         print(f"Configuration:")
         print(f"  Model dimension (d_model):          {self.d_model}")
         print(f"  KV latent dimension (d_latent_kv): {self.d_latent_kv}")

tokenization_taonet.py

@@ -30,6 +30,9 @@ class TaoNetTokenizer(PreTrainedTokenizer):
         except ImportError as exc:
             raise ImportError("TaoNetTokenizer requires sentencepiece to be installed") from exc
 
+        # Newer Transformers versions may round-trip this field from tokenizer_config.json
+        # as either a dict or a list. TaoNet only needs the token strings here.
+        extra_special_tokens = kwargs.pop("extra_special_tokens", None)
         self.vocab_file = vocab_file
         self.special_tokens_file = special_tokens_file
         self.sp_model = spm.SentencePieceProcessor()
@@ -50,9 +53,14 @@ class TaoNetTokenizer(PreTrainedTokenizer):
         }
 
         merged_additional_tokens = list(additional_special_tokens or [])
+        if isinstance(extra_special_tokens, dict):
+            merged_additional_tokens.extend(str(token) for token in extra_special_tokens.keys())
+        elif isinstance(extra_special_tokens, (list, tuple)):
+            merged_additional_tokens.extend(str(token) for token in extra_special_tokens)
         for token in configured_special_tokens:
             if token not in {bos_token, eos_token, unk_token, pad_token} and token not in merged_additional_tokens:
                 merged_additional_tokens.append(token)
+        merged_additional_tokens = list(dict.fromkeys(merged_additional_tokens))
 
         super().__init__(
             bos_token=bos_token,
@@ -125,6 +133,26 @@ class TaoNetTokenizer(PreTrainedTokenizer):
             return ""
         return self.sp_model.decode_pieces(tokens)
 
+    def decode(self, token_ids, skip_special_tokens=False, **kwargs):
+        del kwargs
+
+        if hasattr(token_ids, "tolist"):
+            token_ids = token_ids.tolist()
+
+        if isinstance(token_ids, (list, tuple)) and token_ids and isinstance(token_ids[0], (list, tuple)):
+            token_ids = token_ids[0]
+
+        if not isinstance(token_ids, list):
+            token_ids = [int(token_ids)]
+        else:
+            token_ids = [int(token_id) for token_id in token_ids]
+
+        if skip_special_tokens:
+            special_token_ids = {int(token_id) for token_id in self.special_token_ids.values()}
+            token_ids = [token_id for token_id in token_ids if token_id not in special_token_ids]
+
+        return self.sp_model.decode(token_ids)
+
     def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None):
         if token_ids_1 is None:
             return list(token_ids_0)

verify_tokenizer_match.py

@@ -0,0 +1,99 @@
+"""Verify that the exported HF tokenizer matches the training-time tokenizer wrapper."""
+
+from pathlib import Path
+import sys
+
+import sentencepiece as spm
+
+from tokenization_taonet import TaoNetTokenizer
+
+
+SAMPLES = [
+    "Explain why compact language models can still be useful.",
+    "Fruit is now expensive so we should",
+    "Hello world",
+    "<user>",
+    "<assistant>",
+    "\n",
+]
+
+
+def main():
+    repo_dir = Path(__file__).resolve().parent
+    sys.path.insert(0, str(repo_dir / "src"))
+    from taoTrain.data.tokenizer import SentencePieceTokenizerWrapper, load_special_token_metadata
+
+    tokenizer_model = repo_dir / "tokenizer" / "tokenizer.model"
+    if not tokenizer_model.exists():
+        tokenizer_model = repo_dir / "tokenizer.model"
+
+    sp = spm.SentencePieceProcessor()
+    sp.Load(str(tokenizer_model))
+    special_token_ids = load_special_token_metadata(tokenizer_model)
+
+    train_tokenizer = SentencePieceTokenizerWrapper(sp, special_token_ids=special_token_ids)
+    hf_tokenizer = TaoNetTokenizer.from_pretrained(str(repo_dir))
+
+    print(f"train vocab_size: {train_tokenizer.vocab_size}")
+    print(f"hf vocab_size:    {hf_tokenizer.vocab_size}")
+
+    for token in ["<UNK>", "<BOS>", "<EOS>", "<PAD>", "<think>", "<user>", "<assistant>", "<image>", "\n"]:
+        train_id = train_tokenizer.get_special_token_id(token)
+        hf_id = hf_tokenizer.get_special_token_id(token)
+        print(f"{token!r}: train={train_id}, hf={hf_id}")
+        if train_id != hf_id:
+            raise SystemExit(f"Special token mismatch for {token}: train={train_id}, hf={hf_id}")
+
+    print("\nChecking ID -> token mapping...")
+    for token_id in range(sp.vocab_size()):
+        train_piece = sp.id_to_piece(token_id)
+        hf_piece = hf_tokenizer._convert_id_to_token(token_id)
+        if token_id in special_token_ids.values():
+            expected = next(token for token, value in special_token_ids.items() if value == token_id)
+            if hf_piece != expected:
+                raise SystemExit(
+                    f"HF id->token mismatch at id={token_id}: expected special token {expected!r}, got {hf_piece!r}"
+                )
+        else:
+            if hf_piece != train_piece:
+                raise SystemExit(
+                    f"HF id->token mismatch at id={token_id}: train={train_piece!r}, hf={hf_piece!r}"
+                )
+
+    print("ID -> token mapping matches.")
+
+    print("\nChecking sample encodes/decodes...")
+    for sample in SAMPLES:
+        train_ids = train_tokenizer(sample, return_attention_mask=True)
+        hf_ids = hf_tokenizer(sample, return_attention_mask=True)
+        print(f"sample: {sample!r}")
+        print(f"  train ids: {train_ids['input_ids']}")
+        print(f"  hf ids:    {hf_ids['input_ids']}")
+        if train_ids["input_ids"] != hf_ids["input_ids"]:
+            raise SystemExit(f"Encoding mismatch for sample {sample!r}")
+
+        train_decoded = train_tokenizer.decode(train_ids["input_ids"], skip_special_tokens=True)
+        hf_decoded = hf_tokenizer.decode(hf_ids["input_ids"], skip_special_tokens=True)
+        print(f"  train decode: {train_decoded!r}")
+        print(f"  hf decode:    {hf_decoded!r}")
+        if train_decoded != hf_decoded:
+            raise SystemExit(f"Decode mismatch for sample {sample!r}")
+
+    prompt = "Explain why compact language models can still be useful."
+    chat_ids = [
+        train_tokenizer.get_special_token_id("<user>"),
+        *train_tokenizer(prompt)["input_ids"],
+        train_tokenizer.get_special_token_id("<assistant>"),
+    ]
+    hf_chat = hf_tokenizer.build_chat_inputs(prompt)
+    print("\nChecking chat prompt construction...")
+    print(f"  train-style chat ids: {chat_ids}")
+    print(f"  hf chat ids:          {hf_chat['input_ids']}")
+    if chat_ids != hf_chat["input_ids"]:
+        raise SystemExit("Chat prompt IDs do not match training-time construction.")
+
+    print("\nTokenizer verification passed.")
+
+
+if __name__ == "__main__":
+    main()