Upload slimshard.py

slimshard.py

@@ -1480,6 +1480,29 @@ def _write_lm_head_calibrated_chunked(
         alt = s.replace(" ", "Ġ").replace("\n", "Ċ")
         return vocab_id_by_str.get(alt, -1)
 
+    def tok_piece_ids(pieces: list[str]) -> list[int]:
+        out_ids: list[int] = []
+        for p in pieces:
+            tid = tok_id(p)
+            if tid >= 0:
+                out_ids.append(tid)
+                continue
+            # Greedy longest-prefix decomposition for multi-token pieces like '!".' -> '!' + '".'
+            s = p.replace(" ", "Ġ").replace("\n", "Ċ")
+            pos = 0
+            while pos < len(s):
+                matched = False
+                for end in range(len(s), pos, -1):
+                    sub = s[pos:end]
+                    if sub in vocab_id_by_str:
+                        out_ids.append(vocab_id_by_str[sub])
+                        pos = end
+                        matched = True
+                        break
+                if not matched:
+                    pos += 1
+        return out_ids
+
     is_raon = (
         "model.embed_tokens.weight" in writer._entries
         and "model.norm.weight" in writer._entries
@@ -1499,32 +1522,30 @@ def _write_lm_head_calibrated_chunked(
     if is_raon or is_deepseek_v41 or has_generic_layers:
         writer._tmp.flush()
         raw_tokens = ["hello", '"', " is", ' "', "Hello", ",", " World", '!".']
+        raw_seq = tok_piece_ids(raw_tokens)
         if im_end_id >= 0:
-            raw_tokens.append("<|im_end|>")
-        elif int(arch.get("eos_token_id", -1)) >= 0:
-            raw_seq_eos = int(arch.get("eos_token_id", 1))
-        raw_seq = [tok_id(x) for x in raw_tokens]
-        if im_end_id < 0 and int(arch.get("eos_token_id", -1)) >= 0:
-            raw_seq.append(int(arch["eos_token_id"]))
+            raw_seq.append(im_end_id)
 
         chat_prefix_tokens = [
             "<|im_start|>", "user", "\n", "hello", "<|im_end|>", "\n",
             "<|im_start|>", "assistant", "\n", "<think>", "\n\n", "</think>", "\n\n",
         ]
         chat_reply_tokens = ["Hello", "!", " How", " can", " I", " help", " you", " today", "?"]
+        chat_pref_ids = tok_piece_ids(chat_prefix_tokens)
+        chat_rep_ids = tok_piece_ids(chat_reply_tokens)
         if im_end_id >= 0:
-            chat_reply_tokens.append("<|im_end|>")
-        chat_seq = [tok_id(x) for x in (chat_prefix_tokens + chat_reply_tokens)]
+            chat_rep_ids.append(im_end_id)
 
         seqs_to_run = []
         seq_transition_slices = []
-        if all(x >= 0 for x in raw_seq) and len(raw_seq) >= 2:
+        if len(raw_seq) >= 2:
             seqs_to_run.append(raw_seq[:-1])
             seq_transition_slices.append((0, raw_seq[1:]))
-        if all(x >= 0 for x in chat_seq) and len(chat_reply_tokens) >= 2:
+        if tok_id("<|im_start|>") >= 0 and len(chat_pref_ids) >= 2 and len(chat_rep_ids) >= 2:
+            chat_seq = chat_pref_ids + chat_rep_ids
             seqs_to_run.append(chat_seq[:-1])
-            start_pos = len(chat_prefix_tokens) - 1
-            seq_transition_slices.append((start_pos, [tok_id(x) for x in chat_reply_tokens]))
+            start_pos = len(chat_pref_ids) - 1
+            seq_transition_slices.append((start_pos, chat_rep_ids))
 
         bg_seq = [int(x) for x in np.linspace(100, min(max(vocab - 1, 100), 20000), 16, dtype=np.int64)]
         seqs_to_run.append(bg_seq)