Add tokenizer normalization retry in load_model

app.py

@@ -37,6 +37,10 @@ TOKENIZER = None
 MODEL_NAME = None
 DEVICE = "cpu"
 MODEL_LOCK = threading.Lock()
+# flag: whether a model load is currently in progress (prevents requests)
+MODEL_LOADING = False
+# flag: whether the loaded tokenizer exposes a chat template helper
+USE_CHAT_TEMPLATE = False
 
 # ----------------------------- Utilities ---# ------------------------------
 
@@ -83,8 +87,18 @@ def _diagnose_and_fix_tokenizer_model(tok: AutoTokenizer, mdl: AutoModelForCausa
     # ensure pad token exists and ids/config align
     if getattr(tok, "pad_token", None) is None:
         tok.pad_token = getattr(tok, "eos_token", "[PAD]")
+        # Be defensive: different tokenizer backends expect different arg types
         try:
             tok.add_special_tokens({"pad_token": tok.pad_token})
+        except TypeError as e:
+            # try list form or add_tokens fallback
+            try:
+                tok.add_special_tokens([tok.pad_token])
+            except Exception:
+                try:
+                    tok.add_tokens([tok.pad_token])
+                except Exception:
+                    pass
         except Exception:
             pass
     try:
@@ -235,6 +249,25 @@ def repair_tokenizer_on_hub(repo_id: str) -> str:
 
 # ----------------------------- Model loading -------------------------------
 
+
+def _safe_model_from_pretrained(repo_id, *args, **kwargs):
+    """Call AutoModelForCausalLM.from_pretrained but retry without `use_auth_token`
+    if the called class improperly forwards unexpected kwargs into __init__.
+    """
+    try:
+        return AutoModelForCausalLM.from_pretrained(repo_id, *args, **kwargs)
+    except TypeError as e:
+        msg = str(e)
+        if "use_auth_token" in msg or "unexpected keyword argument" in msg:
+            # retry without auth-token kwargs (some remote `from_pretrained` may leak kwargs)
+            kwargs2 = dict(kwargs)
+            kwargs2.pop("use_auth_token", None)
+            kwargs2.pop("token", None)
+            print(f"_safe_model_from_pretrained: retrying without auth-token due to: {e}")
+            return AutoModelForCausalLM.from_pretrained(repo_id, *args, **kwargs2)
+        raise
+
+
 def load_model(repo_id: str = DEFAULT_MODEL, force_reload: bool = False) -> str:
     """Load model + tokenizer from the Hub. Graceful fallbacks and HF-token support.
 
@@ -250,6 +283,11 @@ def load_model(repo_id: str = DEFAULT_MODEL, force_reload: bool = False) -> str:
         if MODEL is not None and MODEL_NAME == repo_id and not force_reload:
             return f"Model already loaded: {MODEL_NAME} (@ {DEVICE})"
 
+        # mark loading state so UI handlers can guard incoming requests
+        global MODEL_LOADING
+        MODEL_LOADING = True
+        print(f"Model load started: {repo_id}")
+
         MODEL = None
         TOKENIZER = None
         MODEL_NAME = repo_id
@@ -265,40 +303,135 @@ def load_model(repo_id: str = DEFAULT_MODEL, force_reload: bool = False) -> str:
                 trust_remote_code=True,
                 use_auth_token=hf_token,
             )
-        except Exception as e_tok:
-            # If a local repo was cloned without git-lfs, tokenizer.model may be a pointer file — try auto-fetch
+            print(f"Tokenizer loaded from repo: {repo_id}")
+            # detect whether tokenizer supports the Nanbeige chat template API
             try:
-                if os.path.isdir(repo_id) and _ensure_local_tokenizer_model(repo_id, hf_token=hf_token):
-                    print(f"Found LFS pointer at {repo_id}/tokenizer.model — fetched real tokenizer.model; retrying tokenizer load...")
-                    TOKENIZER = AutoTokenizer.from_pretrained(
-                        repo_id,
-                        use_fast=False,
-                        trust_remote_code=True,
-                        use_auth_token=hf_token,
-                    )
-                    # success — continue to model load
-                else:
-                    # fallback: try base model tokenizer (common fix when adapter upload missed tokenizer.model)
-                    print(f"Tokenizer load from {repo_id} failed: {e_tok}. Falling back to base tokenizer PioTio/Nanbeige2.5...")
-                    TOKENIZER = AutoTokenizer.from_pretrained(
-                        DEFAULT_MODEL,
-                        use_fast=False,
-                        trust_remote_code=True,
-                        use_auth_token=hf_token,
-                    )
-            except Exception as e_base:
-                # last-resort: try fast tokenizer (may still fail or produce garbled output)
+                global USE_CHAT_TEMPLATE
+                USE_CHAT_TEMPLATE = hasattr(TOKENIZER, "apply_chat_template")
+                print(f"USE_CHAT_TEMPLATE={USE_CHAT_TEMPLATE}")
+            except Exception:
+                USE_CHAT_TEMPLATE = False
+        except Exception as e_tok:
+            print(f"Tokenizer load from {repo_id} failed: {e_tok}")
+            # specific fix: some tokenizers fail with 'Input must be a List...' when
+            # `special_tokens_map.json` contains dict entries instead of plain strings.
+            # Try an in-memory normalization + local retry before broader fallbacks/repairs.
+            if "Input must be a List" in str(e_tok) or "Input must be a List[Union[str, AddedToken]]" in str(e_tok):
+                try:
+                    print('Detected tokenizer add-tokens type error; attempting in-place normalization and retry...')
+                    # try to download tokenizer files and normalize special_tokens_map.json
+                    try:
+                        from huggingface_hub import hf_hub_download
+                        import json, tempfile, shutil
+
+                        tmp = tempfile.mkdtemp(prefix="tokfix_")
+                        # files we need locally for AutoTokenizer
+                        candidates = ["tokenizer.json", "tokenizer_config.json", "special_tokens_map.json", "tokenizer.model", "added_tokens.json"]
+                        for fn in candidates:
+                            try:
+                                src = hf_hub_download(repo_id=repo_id, filename=fn, token=hf_token)
+                                shutil.copy(src, tmp)
+                            except Exception:
+                                # ignore missing files — AutoTokenizer is tolerant
+                                pass
+
+                        # normalize special_tokens_map.json if present
+                        stm = os.path.join(tmp, "special_tokens_map.json")
+                        if os.path.exists(stm):
+                            try:
+                                with open(stm, "r", encoding="utf-8") as f:
+                                    stm_j = json.load(f)
+                                changed = False
+                                if "additional_special_tokens" in stm_j:
+                                    new = []
+                                    for it in stm_j["additional_special_tokens"]:
+                                        if isinstance(it, dict):
+                                            new.append(it.get("content") or it.get("token") or str(it))
+                                            changed = True
+                                        else:
+                                            new.append(it)
+                                    stm_j["additional_special_tokens"] = new
+                                for k in ["bos_token", "eos_token", "pad_token", "unk_token"]:
+                                    if k in stm_j and isinstance(stm_j[k], dict):
+                                        stm_j[k] = stm_j[k].get("content", stm_j[k])
+                                        changed = True
+                                if changed:
+                                    with open(stm, "w", encoding="utf-8") as f:
+                                        json.dump(stm_j, f, ensure_ascii=False, indent=2)
+                                    print('Normalized special_tokens_map.json in temp dir')
+                            except Exception:
+                                pass
+
+                        # try loading tokenizer from the temporary normalized directory
+                        TOKENIZER = AutoTokenizer.from_pretrained(tmp, use_fast=False, trust_remote_code=True)
+                        print('Tokenizer reloaded from normalized temp copy')
+                        shutil.rmtree(tmp)
+                    except Exception as e_localnorm:
+                        print('In-place normalization retry failed:', e_localnorm)
+                        # fall through to the existing repair path below
+
+                    # as a fallback, attempt to auto-repair the remote repo (if HF token available)
+                    if hf_token:
+                        print('Attempting repo-side auto-repair/upload from base tokenizer...')
+                        _repair_and_upload_tokenizer(repo_id, hf_token=hf_token)
+                        TOKENIZER = AutoTokenizer.from_pretrained(repo_id, use_fast=False, trust_remote_code=True)
+                        print('Tokenizer reloaded after repo repair')
+                    else:
+                        # final fallback will be handled by the outer fallbacks below
+                        raise RuntimeError('Normalization + auto-repair could not proceed (no HF_TOKEN)')
+                except Exception as e_retry:
+                    print('Repair/retry failed:', e_retry)
+                    return f"Tokenizer load failed: {e_retry}"
+            else:
+                # If a local repo was cloned without git-lfs, tokenizer.model may be a pointer file — try auto-fetch
                 try:
-                    print(f"Base tokenizer fallback failed: {e_base}. Trying generic AutoTokenizer...")
-                    TOKENIZER = AutoTokenizer.from_pretrained(repo_id, trust_remote_code=True, use_auth_token=hf_token)
-                except Exception as e_final:
-                    return f"Tokenizer load failed: {e_final}"
+                    if os.path.isdir(repo_id) and _ensure_local_tokenizer_model(repo_id, hf_token=hf_token):
+                        print(f"Found LFS pointer at {repo_id}/tokenizer.model — fetched real tokenizer.model; retrying tokenizer load...")
+                        TOKENIZER = AutoTokenizer.from_pretrained(
+                            repo_id,
+                            use_fast=False,
+                            trust_remote_code=True,
+                            use_auth_token=hf_token,
+                        )
+                        print(f"Tokenizer loaded from local repo after fetching LFS: {repo_id}")
+                    else:
+                        # Local workspace fallback: use bundled Nanbeige4.1 tokenizer if available
+                        local_fallback = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'models', 'Nanbeige4.1-3B'))
+                        if os.path.isdir(local_fallback):
+                            try:
+                                print(f"Attempting local workspace tokenizer fallback: {local_fallback}")
+                                TOKENIZER = AutoTokenizer.from_pretrained(local_fallback, use_fast=False, trust_remote_code=True)
+                                print(f"Tokenizer loaded from local workspace: {local_fallback}")
+                            except Exception as e_local:
+                                print(f"Local tokenizer fallback failed: {e_local}")
+                                raise e_local
+                        else:
+                            # Try known base tokenizer on the Hub (Nanbeige4.1 if repo looks like 4.1)
+                            base = "Nanbeige/Nanbeige4.1-3B" if "4.1" in repo_id.lower() else "PioTio/Nanbeige2.5"
+                            print(f"Falling back to base tokenizer: {base}")
+                            TOKENIZER = AutoTokenizer.from_pretrained(base, use_fast=False, trust_remote_code=True, use_auth_token=hf_token)
+
+                        # If HF token is available, attempt to auto-repair/upload tokenizer files to the target repo
+                        if hf_token:
+                            try:
+                                uploaded = _repair_and_upload_tokenizer(repo_id, hf_token=hf_token)
+                                print(f"Auto-repair attempt to {repo_id}: {'succeeded' if uploaded else 'no-change/failure'}")
+                            except Exception as e_rep:
+                                print(f"Auto-repair attempt failed: {e_rep}")
+                except Exception as e_base:
+                    # last-resort: try fast tokenizer (may still fail or produce garbled output)
+                    try:
+                        print(f"All fallbacks failed: {e_base}. Trying generic AutoTokenizer as last resort...")
+                        TOKENIZER = AutoTokenizer.from_pretrained(repo_id, trust_remote_code=True, use_auth_token=hf_token)
+                    except Exception as e_final:
+                        MODEL_LOADING = False
+                        return f"Tokenizer load failed: {e_final}"
 
         # 2) Load model (prefer 4-bit on GPU if available)
         if DEVICE == "cuda" and HAS_BNB:
             try:
                 bnb_config = BitsAndBytesConfig(load_in_4bit=True)
-                MODEL = AutoModelForCausalLM.from_pretrained(
+                MODEL = _safe_model_from_pretrained(
                     repo_id,
                     device_map="auto",
                     quantization_config=bnb_config,
@@ -307,6 +440,8 @@ def load_model(repo_id: str = DEFAULT_MODEL, force_reload: bool = False) -> str:
                 )
                 MODEL.eval()
                 _diagnose_and_fix_tokenizer_model(TOKENIZER, MODEL)
+                MODEL_LOADING = False
+                print(f"Model load finished (4-bit): {repo_id}")
                 return f"Loaded {repo_id} (4-bit, device_map=auto)"
             except Exception as e:
                 print("bnb/4bit load failed - falling back:", e)
@@ -314,7 +449,7 @@ def load_model(repo_id: str = DEFAULT_MODEL, force_reload: bool = False) -> str:
         # 3) FP16 / CPU fallback
         try:
             if DEVICE == "cuda":
-                MODEL = AutoModelForCausalLM.from_pretrained(
+                MODEL = _safe_model_from_pretrained(
                     repo_id,
                     device_map="auto",
                     torch_dtype=torch.float16,
@@ -322,7 +457,7 @@ def load_model(repo_id: str = DEFAULT_MODEL, force_reload: bool = False) -> str:
                     use_auth_token=hf_token,
                 )
             else:
-                MODEL = AutoModelForCausalLM.from_pretrained(
+                MODEL = _safe_model_from_pretrained(
                     repo_id,
                     low_cpu_mem_usage=True,
                     torch_dtype=torch.float32,
@@ -333,11 +468,15 @@ def load_model(repo_id: str = DEFAULT_MODEL, force_reload: bool = False) -> str:
 
             MODEL.eval()
             _diagnose_and_fix_tokenizer_model(TOKENIZER, MODEL)
+            MODEL_LOADING = False
+            print(f"Model load finished: {repo_id} (@{DEVICE})")
             return f"Loaded {repo_id} (@{DEVICE})"
         except Exception as e:
             MODEL = None
             TOKENIZER = None
-            # provide a helpful diagnostic message
+            # clear loading flag and provide a helpful diagnostic message
+            MODEL_LOADING = False
+            print(f"Model load failed: {repo_id} -> {e}")
             return f"Model load failed: {e} (hint: check HF_TOKEN, repo contents and ensure tokenizer.model is present)"
 
 
@@ -358,6 +497,8 @@ def _normalize_history(raw_history) -> List[Tuple[str, str]]:
     and return a list of (user, assistant) pairs suitable for prompt construction.
 
     Behavior: pairs each user message with the next assistant message (assistant may be "" if not present).
+    NOTE: For chat-first models (Nanbeige4.1) we prefer `tokenizer.apply_chat_template` later
+    so this function only normalizes the history shape.
     """
     if not raw_history:
         return []
@@ -409,6 +550,33 @@ def build_prompt(history, user_input: str, system_prompt: str, max_history: int
     pairs = _normalize_history(history or [])
     pairs = pairs[-max_history:]
 
+    # If tokenizer provides a chat-template helper (Nanbeige4.1), use it.
+    # This avoids instruction-format mismatches that produce garbled output.
+    try:
+        from __main__ import TOKENIZER  # safe access to global TOKENIZER when available
+    except Exception:
+        TOKENIZER = None
+
+    if TOKENIZER is not None and hasattr(TOKENIZER, "apply_chat_template"):
+        # build messages list with optional system prompt first
+        messages = []
+        if system_prompt:
+            messages.append({"role": "system", "content": system_prompt})
+        for u, a in pairs:
+            messages.append({"role": "user", "content": u})
+            if a:
+                messages.append({"role": "assistant", "content": a})
+        # current user turn
+        messages.append({"role": "user", "content": user_input})
+        # use tokenizer's chat template (returns the full prompt string)
+        try:
+            prompt = TOKENIZER.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
+            return prompt
+        except Exception:
+            # fall back to ALPACA format if anything goes wrong
+            pass
+
+    # Default / fallback: ALPACA-style instruction template
     parts: List[str] = [f"System: {system_prompt}"]
     for u, a in pairs:
         # include previous turns as completed instruction/response pairs
@@ -425,7 +593,10 @@ def _generate_text(prompt: str, temperature: float, top_p: float, top_k: int, ma
     if MODEL is None or TOKENIZER is None:
         raise RuntimeError("Model is not loaded. Press 'Load model' first.")
 
-    input_ids = TOKENIZER(prompt, return_tensors="pt", truncation=True, max_length=2048).input_ids.to(next(MODEL.parameters()).device)
+    # When using a chat-template prompt we must avoid adding special tokens again
+    add_special_tokens = False if hasattr(TOKENIZER, "apply_chat_template") else True
+
+    input_ids = TOKENIZER(prompt, return_tensors="pt", truncation=True, max_length=2048, add_special_tokens=add_special_tokens).input_ids.to(next(MODEL.parameters()).device)
 
     gen_kwargs = dict(
         input_ids=input_ids,
@@ -452,7 +623,8 @@ def _generate_stream(prompt: str, temperature: float, top_p: float, top_k: int,
         raise RuntimeError("Model is not loaded. Press 'Load model' first.")
 
     streamer = TextIteratorStreamer(TOKENIZER, skip_prompt=True, skip_special_tokens=True)
-    input_ids = TOKENIZER(prompt, return_tensors="pt", truncation=True, max_length=2048).input_ids.to(next(MODEL.parameters()).device)
+    add_special_tokens = False if hasattr(TOKENIZER, "apply_chat_template") else True
+    input_ids = TOKENIZER(prompt, return_tensors="pt", truncation=True, max_length=2048, add_special_tokens=add_special_tokens).input_ids.to(next(MODEL.parameters()).device)
 
     gen_kwargs = dict(
         input_ids=input_ids,
@@ -491,6 +663,17 @@ def submit_message(user_message: str, history, system_prompt: str, temperature:
     # Append current user turn (assistant reply empty until generated)
     pairs.append((str(user_message or ""), ""))
 
+    # Guard: block generation while model is loading or not loaded
+    if MODEL_LOADING:
+        pairs[-1] = (user_message, "⚠️ Model is still loading — please wait and try again. Check 'Status' for progress.")
+        yield pairs, ""
+        return
+
+    if MODEL is None:
+        pairs[-1] = (user_message, "⚠️ Model is not loaded — click 'Load model' first.")
+        yield pairs, ""
+        return
+
     prompt = build_prompt(pairs[:-1], user_message, system_prompt, max_history)
 
     # If user is running the full Nanbeige model on CPU, warn and suggest options
@@ -538,7 +721,15 @@ def regenerate(history, system_prompt: str, temperature: float, top_p: float, to
 
 def load_model_ui(repo: str):
     status = load_model(repo, force_reload=True)
-    return status
+    try:
+        suffix = " — chat-template detected" if USE_CHAT_TEMPLATE else ""
+    except NameError:
+        suffix = ""
+    # enable the Send button only when the model actually loaded
+    loaded = str(status).lower().startswith("loaded")
+    from gradio import update as gr_update
+    send_state = gr_update(interactive=loaded)
+    return status + suffix, send_state
 
 
 def apply_lora_adapter(adapter_repo: str):
@@ -609,7 +800,7 @@ with gr.Blocks(title="Nanbeige2.5 — Chat UI") as demo:
         apply_adapter = gr.Button("Apply LoRA adapter")
 
     # Events
-    load_btn.click(fn=lambda repo: load_model_ui(repo), inputs=model_input, outputs=model_status)
+    load_btn.click(fn=load_model_ui, inputs=model_input, outputs=[model_status, send])
     repair_btn.click(fn=repair_tokenizer_on_hub, inputs=model_input, outputs=model_status)
 
     send.click(
@@ -635,7 +826,21 @@ with gr.Blocks(title="Nanbeige2.5 — Chat UI") as demo:
 
     # auto-load default model in background (non-blocking)
     def _bg_initial_load():
-        return load_model(DEFAULT_MODEL, force_reload=False)
+        # run load_model in a background thread to warm up model on Space startup
+        def _worker():
+            res = load_model(DEFAULT_MODEL, force_reload=False)
+            try:
+                # update UI Send button when loaded
+                from gradio import update as gr_update
+                interactive = str(res).lower().startswith("loaded")
+                send.update(interactive=interactive)
+            except Exception:
+                pass
+            return res
+
+        t = threading.Thread(target=_worker, daemon=True)
+        t.start()
+        return "Loading model in background..."
 
     # For local smoke tests you can skip automatic model loading by setting
     # environment variable `SKIP_AUTOLOAD=1` so the UI starts without loading
@@ -644,6 +849,11 @@ with gr.Blocks(title="Nanbeige2.5 — Chat UI") as demo:
         model_status.value = "Auto-load skipped (SKIP_AUTOLOAD=1)"
     else:
         model_status.value = _bg_initial_load()
+        # disable Send while background load is in progress
+        try:
+            send.update(interactive=False)
+        except Exception:
+            pass
 
     # CPU warning / demo hint (visible in UI)
     gr.Markdown("""