Set DEFAULT_MODEL -> nanbeige-4.1-aiman-merged (by user token)

app.py

@@ -27,9 +27,9 @@ except Exception:
 # ---------------------------------------------------------------------------
 # Config / defaults
 # ---------------------------------------------------------------------------
-DEFAULT_MODEL = "PioTio/nanbeige2.5-nsfw-merged"
+DEFAULT_MODEL = "PioTio/nanbeige-4.1-aiman-merged"
 CPU_DEMO_MODEL = "distilgpt2"  # fast, small CPU-friendly fallback for demos
-DEFAULT_SYSTEM_PROMPT = "You are a helpful, honest inteligent AI chatbot assistant. Answer succinctly unless asked otherwise."
+DEFAULT_SYSTEM_PROMPT = "You are a helpful, honest assistant. Answer succinctly unless asked otherwise."
 
 # globals populated by load_model()
 MODEL = None
@@ -95,11 +95,154 @@ def _diagnose_and_fix_tokenizer_model(tok: AutoTokenizer, mdl: AutoModelForCausa
         pass
 
 
+# Helper: detect Git-LFS pointer files and fetch real tokenizer.model from the Hub
+def _is_lfs_pointer_file(path: str) -> bool:
+    try:
+        with open(path, "rb") as f:
+            start = f.read(128)
+        return b"git-lfs.github.com/spec/v1" in start
+    except Exception:
+        return False
+
+
+def _download_tokenizer_model_from_hub(hf_repo: str, dest_path: str, hf_token: Optional[str] = None) -> bool:
+    """Download tokenizer.model from HF Hub into dest_path. Returns True on success."""
+    try:
+        import urllib.request
+
+        url = f"https://huggingface.co/{hf_repo}/resolve/main/tokenizer.model"
+        req = urllib.request.Request(url, headers={"User-Agent": "spaces-nanbeige-chat/1.0"})
+        if hf_token:
+            req.add_header("Authorization", f"Bearer {hf_token}")
+        with urllib.request.urlopen(req, timeout=30) as r, open(dest_path + ".tmp", "wb") as out:
+            out.write(r.read())
+        os.replace(dest_path + ".tmp", dest_path)
+        return True
+    except Exception as e:
+        print("_download_tokenizer_model_from_hub failed:", e)
+        try:
+            if os.path.exists(dest_path + ".tmp"):
+                os.remove(dest_path + ".tmp")
+        except Exception:
+            pass
+        return False
+
+
+def _ensure_local_tokenizer_model(repo_path: str, hf_token: Optional[str] = None) -> bool:
+    """If tokenizer.model in repo_path is a Git-LFS pointer, try to download the real file from the Hub.
+    Tries to infer a Hub repo id from the local git remote; falls back to `PioTio/<dirname>` for Nanbeige folders.
+    """
+    tm = os.path.join(repo_path, "tokenizer.model")
+    if not os.path.exists(tm):
+        return False
+    if not _is_lfs_pointer_file(tm):
+        return True
+
+    # try to get repo id from git remote origin
+    repo_id = None
+    try:
+        import subprocess
+
+        out = subprocess.check_output(["git", "-C", repo_path, "config", "--get", "remote.origin.url"], text=True).strip()
+        if out and "huggingface.co" in out:
+            # parse https://huggingface.co/owner/repo(.git)
+            parts = out.rstrip(".git").split("/")
+            repo_id = f"{parts[-2]}/{parts[-1]}"
+    except Exception:
+        repo_id = None
+
+    # fallback: guess owner for common Nanbeige folder names
+    if repo_id is None:
+        guessed = os.path.basename(repo_path)
+        if guessed.lower().startswith("nanbeige") or "nanbeige" in guessed.lower():
+            repo_id = f"PioTio/{guessed}"
+
+    if repo_id:
+        return _download_tokenizer_model_from_hub(repo_id, tm, hf_token=hf_token)
+    return False
+
+
+# Helper: upload tokenizer files (from a local tokenizer dir) back to a Hub repo
+def _upload_tokenizer_files_to_hub(repo_id: str, local_tokenizer_dir: str, hf_token: Optional[str] = None) -> bool:
+    """Upload tokenizer files (tokenizer.model, tokenizer_config.json, tokenizer.json, special_tokens_map.json)
+    Returns True if at least one file was uploaded successfully.
+    """
+    try:
+        from huggingface_hub import HfApi
+        api = HfApi()
+        candidates = [
+            "tokenizer.model",
+            "tokenizer_config.json",
+            "tokenizer.json",
+            "special_tokens_map.json",
+            "chat_template.jinja",
+        ]
+        uploaded = 0
+        for fn in candidates:
+            p = os.path.join(local_tokenizer_dir, fn)
+            if not os.path.exists(p):
+                continue
+            try:
+                api.upload_file(
+                    path_or_fileobj=p,
+                    path_in_repo=fn,
+                    repo_id=repo_id,
+                    token=hf_token,
+                    commit_message=f"Auto-fix tokenizer: {fn}",
+                )
+                print(f"_upload_tokenizer_files_to_hub: uploaded {fn} to {repo_id}")
+                uploaded += 1
+            except Exception as e:
+                print(f"_upload_tokenizer_files_to_hub: failed to upload {fn}: {e}")
+        return uploaded > 0
+    except Exception as e:
+        print("_upload_tokenizer_files_to_hub failed:", e)
+        return False
+
+
+def _repair_and_upload_tokenizer(repo_id: str, hf_token: Optional[str] = None) -> bool:
+    """Fetch the correct base tokenizer (Nanbeige4.1 if detected, otherwise DEFAULT_MODEL),
+    then upload tokenizer files to the target repo. Returns True on success.
+    """
+    try:
+        base = "Nanbeige/Nanbeige4.1-3B" if "4.1" in repo_id.lower() else DEFAULT_MODEL
+        from transformers import AutoTokenizer
+        import tempfile, shutil
+        tmp = tempfile.mkdtemp(prefix="tokenizer_fix_")
+        tok = AutoTokenizer.from_pretrained(base, use_fast=False, trust_remote_code=True)
+        tok.save_pretrained(tmp)
+        ok = _upload_tokenizer_files_to_hub(repo_id, tmp, hf_token=hf_token)
+        shutil.rmtree(tmp)
+        return ok
+    except Exception as e:
+        print("_repair_and_upload_tokenizer failed:", e)
+        return False
+
+
+def repair_tokenizer_on_hub(repo_id: str) -> str:
+    """Public helper callable from the UI: attempts to upload a working base tokenizer to `repo_id`.
+    Requires HF_TOKEN in the environment with write access to the target repo.
+    """
+    hf_token = os.environ.get("HF_TOKEN")
+    if not hf_token:
+        return "HF_TOKEN not set — cannot upload tokenizer to Hub. Add HF_TOKEN and retry."
+    try:
+        ok = _repair_and_upload_tokenizer(repo_id, hf_token=hf_token)
+        return "Uploaded tokenizer files to repo" if ok else "Repair attempt failed (see logs)"
+    except Exception as e:
+        return f"Repair failed: {e}"
+
+
 # ----------------------------- Model loading -------------------------------
 
 def load_model(repo_id: str = DEFAULT_MODEL, force_reload: bool = False) -> str:
-    """Load model + tokenizer from the Hub. Graceful fallbacks for CPU/GPU/4-bit.
-    Returns a short status string for the UI.
+    """Load model + tokenizer from the Hub. Graceful fallbacks and HF-token support.
+
+    Changes made:
+    - prefer slow tokenizer (use_fast=False)
+    - accept HF token via env HF_TOKEN for private repos / higher rate limits
+    - fallback to base tokenizer (`PioTio/Nanbeige2.5`) when tokenizer files are missing
+    - pass auth token into from_pretrained calls where supported
     """
     global MODEL, TOKENIZER, MODEL_NAME, DEVICE
 
@@ -112,14 +255,46 @@ def load_model(repo_id: str = DEFAULT_MODEL, force_reload: bool = False) -> str:
         MODEL_NAME = repo_id
 
         DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+        hf_token = os.environ.get("HF_TOKEN")
 
+        # 1) Try to load tokenizer (slow tokenizer is required for Nanbeige family)
         try:
-            # prefer safe (non-fast) tokenizer to avoid SentencePiece piece-id bugs
-            TOKENIZER = AutoTokenizer.from_pretrained(repo_id, use_fast=False, trust_remote_code=True)
-        except Exception:
-            TOKENIZER = AutoTokenizer.from_pretrained(repo_id, trust_remote_code=True)
-
-        # attempt bnb 4-bit if GPU available and BitsAndBytes present
+            TOKENIZER = AutoTokenizer.from_pretrained(
+                repo_id,
+                use_fast=False,
+                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
+            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)
+                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}"
+
+        # 2) Load model (prefer 4-bit on GPU if available)
         if DEVICE == "cuda" and HAS_BNB:
             try:
                 bnb_config = BitsAndBytesConfig(load_in_4bit=True)
@@ -128,31 +303,42 @@ def load_model(repo_id: str = DEFAULT_MODEL, force_reload: bool = False) -> str:
                     device_map="auto",
                     quantization_config=bnb_config,
                     trust_remote_code=True,
+                    use_auth_token=hf_token,
                 )
                 MODEL.eval()
                 _diagnose_and_fix_tokenizer_model(TOKENIZER, MODEL)
                 return f"Loaded {repo_id} (4-bit, device_map=auto)"
             except Exception as e:
-                # fall through to fp16 load
                 print("bnb/4bit load failed - falling back:", e)
 
-        # try fp16 auto device_map if GPU present
+        # 3) FP16 / CPU fallback
         try:
             if DEVICE == "cuda":
-                MODEL = AutoModelForCausalLM.from_pretrained(repo_id, device_map="auto", torch_dtype=torch.float16, trust_remote_code=True)
+                MODEL = AutoModelForCausalLM.from_pretrained(
+                    repo_id,
+                    device_map="auto",
+                    torch_dtype=torch.float16,
+                    trust_remote_code=True,
+                    use_auth_token=hf_token,
+                )
             else:
-                MODEL = AutoModelForCausalLM.from_pretrained(repo_id, low_cpu_mem_usage=True, torch_dtype=torch.float32, trust_remote_code=True)
+                MODEL = AutoModelForCausalLM.from_pretrained(
+                    repo_id,
+                    low_cpu_mem_usage=True,
+                    torch_dtype=torch.float32,
+                    trust_remote_code=True,
+                    use_auth_token=hf_token,
+                )
                 MODEL.to("cpu")
-            MODEL.eval()
 
-            # run the tokenizer/model alignment fix (important for Nanbeige family)
+            MODEL.eval()
             _diagnose_and_fix_tokenizer_model(TOKENIZER, MODEL)
-
             return f"Loaded {repo_id} (@{DEVICE})"
         except Exception as e:
             MODEL = None
             TOKENIZER = None
-            return f"Model load failed: {e}"
+            # provide a helpful diagnostic message
+            return f"Model load failed: {e} (hint: check HF_TOKEN, repo contents and ensure tokenizer.model is present)"
 
 
 # ----------------------------- Prompt building -----------------------------
@@ -361,11 +547,14 @@ def apply_lora_adapter(adapter_repo: str):
     global MODEL
     if MODEL is None:
         return "Load base model first."
+
+    hf_token = os.environ.get("HF_TOKEN")
     try:
-        MODEL = PeftModel.from_pretrained(MODEL, adapter_repo)
+        # allow huggingface auth token for private adapters
+        MODEL = PeftModel.from_pretrained(MODEL, adapter_repo, use_auth_token=hf_token)
         return f"Applied LoRA adapter from {adapter_repo}"
     except Exception as e:
-        return f"Failed to apply adapter: {e}"
+        return f"Failed to apply adapter: {e} (hint: check adapter name and HF_TOKEN)"
 
 
 # ----------------------------- Build UI -----------------------------------
@@ -376,6 +565,7 @@ with gr.Blocks(title="Nanbeige2.5 — Chat UI") as demo:
     with gr.Row():
         model_input = gr.Textbox(value=DEFAULT_MODEL, label="Model repo (HF)", interactive=True)
         load_btn = gr.Button("Load model")
+        repair_btn = gr.Button("Repair tokenizer on Hub")
         model_demo_btn = gr.Button(f"Load fast CPU demo ({CPU_DEMO_MODEL})")
         model_status = gr.Textbox(value="Model not loaded", label="Status", interactive=False)
 
@@ -420,6 +610,7 @@ with gr.Blocks(title="Nanbeige2.5 — Chat UI") as demo:
 
     # Events
     load_btn.click(fn=lambda repo: load_model_ui(repo), inputs=model_input, outputs=model_status)
+    repair_btn.click(fn=repair_tokenizer_on_hub, inputs=model_input, outputs=model_status)
 
     send.click(
         fn=submit_message,