Harden vocab loader: accept tuple/dict/list and rebuild missing sides; EOS-tolerant decode

bot_inference.py

@@ -162,72 +162,90 @@ def _load_ckpt():
 
 _checkpoint = _load_ckpt()
 
-# ---- Build tokenizer (ultra-robust) ----
+# ---- Build tokenizer (ultra-robust, reconstruct missing sides) ----
 def _build_vocab(checkpoint):
     v = checkpoint.get("vocab", None)
     stoi = None
     itos = None
 
-    # 1) tuple or list of length 2: (stoi, itos)
-    if isinstance(v, (tuple, list)) and len(v) == 2:
-        s_raw, i_raw = v
-        # normalize stoi
-        if isinstance(s_raw, dict) and s_raw:
-            stoi = s_raw
-        elif s_raw is None and isinstance(i_raw, (dict, list)) and i_raw:
-            pass  # we'll build stoi from itos below
-        # normalize itos
-        if isinstance(i_raw, dict) and i_raw:
-            itos = i_raw
-        elif isinstance(i_raw, list) and i_raw:
-            itos = {i: ch for i, ch in enumerate(i_raw)}
-        elif i_raw is None and isinstance(s_raw, dict) and s_raw:
-            pass  # we'll build itos from stoi below
-
-    # 2) dict form inside 'vocab'
-    elif isinstance(v, dict) and v:
-        if "stoi" in v: stoi = v["stoi"]
+    def dictify_itos(x):
+        # accept dict[int->str] or list[str]
+        if isinstance(x, dict) and x and all(isinstance(k, int) for k in x.keys()):
+            return x
+        if isinstance(x, list) and x:
+            return {i: ch for i, ch in enumerate(x)}
+        return None
+
+    def dictify_stoi(x):
+        # accept dict[str->int]
+        if isinstance(x, dict) and x and all(isinstance(k, str) for k in x.keys()):
+            return x
+        return None
+
+    # Case A: vocab is a tuple/list (any length >= 2)
+    if isinstance(v, (tuple, list)) and len(v) >= 2:
+        s_raw, i_raw = v[0], v[1]
+        s_norm = dictify_stoi(s_raw)
+        i_norm = dictify_itos(i_raw)
+
+        # If itos is None but stoi exists, build itos from stoi
+        if i_norm is None and isinstance(s_raw, dict) and s_raw:
+            itos = {i: ch for ch, i in s_raw.items()}
+        else:
+            itos = i_norm
+
+        # If stoi is None but itos exists, build stoi from itos
+        if s_norm is None and isinstance(itos, dict) and itos:
+            stoi = {ch: i for i, ch in itos.items()}
+        else:
+            stoi = s_norm
+
+    # Case B: vocab is a dict (may contain keys 'stoi' and/or 'itos', or be pure stoi/itos)
+    elif isinstance(v, dict) and v is not None:
+        if "stoi" in v:
+            stoi = dictify_stoi(v["stoi"])
         if "itos" in v:
-            itos_raw = v["itos"]
-            if isinstance(itos_raw, dict) and itos_raw:
-                itos = itos_raw
-            elif isinstance(itos_raw, list) and itos_raw:
-                itos = {i: ch for i, ch in enumerate(itos_raw)}
-        # pure itos or pure stoi in v
+            itos = dictify_itos(v["itos"])
+        # pure itos (id->char dict)
         if itos is None and all(isinstance(k, int) for k in v.keys()):
-            itos = v  # id->char dict
+            itos = v
+        # pure stoi (char->id dict)
         if stoi is None and all(isinstance(k, str) for k in v.keys()):
-            stoi = v  # char->id dict
-
-    # 3) legacy separate keys at top-level
+            stoi = v
+        # reconstruct missing side
+        if itos is None and isinstance(stoi, dict) and stoi:
+            itos = {i: ch for ch, i in stoi.items()}
+        if stoi is None and isinstance(itos, dict) and itos:
+            stoi = {ch: i for i, ch in itos.items()}
+
+    # Case C: legacy top-level keys
     else:
-        if "stoi" in checkpoint: stoi = checkpoint["stoi"]
+        if "stoi" in checkpoint:
+            stoi = dictify_stoi(checkpoint["stoi"])
         if "itos" in checkpoint:
-            i_raw = checkpoint["itos"]
-            if isinstance(i_raw, dict) and i_raw:
-                itos = i_raw
-            elif isinstance(i_raw, list) and i_raw:
-                itos = {i: ch for i, ch in enumerate(i_raw)}
-
-    # 4) Rebuild the missing side if possible
-    if itos is None and isinstance(stoi, dict) and stoi:
-        itos = {i: ch for ch, i in stoi.items()}
-    if stoi is None and isinstance(itos, dict) and itos:
+            itos = dictify_itos(checkpoint["itos"])
+        if itos is None and isinstance(stoi, dict) and stoi:
+            itos = {i: ch for ch, i in stoi.items()}
+        if stoi is None and isinstance(itos, dict) and itos:
+            stoi = {ch: i for i, ch in itos.items()}
+
+    # Final fallback: if vocab is a flat list/tuple of characters
+    if (not stoi or not itos) and isinstance(v, (list, tuple)) and v and all(isinstance(x, str) for x in v):
+        itos = {i: ch for i, ch in enumerate(v)}
         stoi = {ch: i for i, ch in itos.items()}
 
-    # 5) Final sanity
-    if not isinstance(stoi, dict) or not stoi or not isinstance(itos, dict) or not itos:
+    if not (isinstance(stoi, dict) and stoi and isinstance(itos, dict) and itos):
         raise KeyError(
-            "Checkpoint vocab not usable. Expected 'vocab'=(stoi, itos), "
-            "'vocab'={'stoi','itos'}, pure stoi/itos, or legacy keys. "
-            f"Found keys: {list(checkpoint.keys())}; type(vocab)={type(v)}"
+            "Checkpoint vocab not usable. Expected 'vocab'=(stoi, itos) or dict with 'stoi'/'itos', "
+            "or pure stoi/itos; can reconstruct if one side present. "
+            f"Found keys: {list(checkpoint.keys())}; type(vocab)={type(v)}; value sample={repr(v)[:200]}"
         )
     return stoi, itos
 
 stoi, itos = _build_vocab(_checkpoint)
 
 EOS = "<EOS>"
-EOS_ID = stoi.get(EOS, None)  # ok if None
+EOS_ID = stoi.get(EOS, None)  # if missing, we'll just decode to max_new_tokens
 
 
 block_size = _checkpoint.get("block_size", None)