Update app.py

app.py

@@ -2,23 +2,26 @@ import os
 import re
 import shutil
 
-import torch
-import torch.nn.functional as F
-import pandas as pd
-import gradio as gr
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
-
 # ============================================================
-# ENV (set BEFORE loading models)
+# ENV (set BEFORE transformers/hub usage)
 # ============================================================
-# Use a predictable cache location (helps avoid reusing a corrupt home cache)
 os.environ.setdefault("HF_HOME", "/tmp/hf")
 os.environ.setdefault("HUGGINGFACE_HUB_CACHE", "/tmp/hf/hub")
 os.environ.setdefault("TRANSFORMERS_CACHE", "/tmp/hf/transformers")
-# Disable Xet (helps avoid partial/corrupt downloads in some environments)
-os.environ.setdefault("HF_HUB_DISABLE_XET", "1")
+os.environ.setdefault("HF_HUB_DISABLE_XET", "1")  # disable hf-xet if present
 os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
 
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import pandas as pd
+import gradio as gr
+
+from huggingface_hub import hf_hub_download
+from transformers import AutoConfig, AutoTokenizer, AutoModel
+from safetensors.torch import load_file
+
+
 # -----------------------------
 # MODEL INITIALIZATION
 # -----------------------------
@@ -26,10 +29,17 @@ MODEL_NAME = "desklib/ai-text-detector-v1.01"
 tokenizer = None
 model = None
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
 THRESHOLD = 0.59
 
 
+def _build_error_card(msg: str) -> str:
+    return (
+        "<div style='color:#b80d0d; padding:14px; border:1px solid #b80d0d; "
+        "border-radius:10px; background:rgba(184,13,13,0.06);'>"
+        f"{msg}</div>"
+    )
+
+
 def wipe_model_cache(model_id: str) -> int:
     """
     Delete cached files for this model from common HF cache locations.
@@ -49,81 +59,104 @@ def wipe_model_cache(model_id: str) -> int:
     removed = 0
     for path in candidates:
         if os.path.exists(path):
-            try:
-                shutil.rmtree(path, ignore_errors=True)
-                removed += 1
-            except Exception:
-                # ignore deletion errors (permissions etc.)
-                pass
+            shutil.rmtree(path, ignore_errors=True)
+            removed += 1
     return removed
 
 
-def _build_error_card(msg: str) -> str:
-    return (
-        "<div style='color:#b80d0d; padding:14px; border:1px solid #b80d0d; "
-        "border-radius:10px; background:rgba(184,13,13,0.06);'>"
-        f"{msg}</div>"
-    )
+class DesklibAIDetectionModel(nn.Module):
+    """
+    Matches the architecture described by desklib:
+    base transformer + mean pooling + linear classifier to 1 logit.
+    The repo config lists "architectures": ["DesklibAIDetectionModel"]. :contentReference[oaicite:1]{index=1}
+    """
+    def __init__(self, config):
+        super().__init__()
+        self.backbone = AutoModel.from_config(config)
+        self.classifier = nn.Linear(config.hidden_size, 1)
+
+    def forward(self, input_ids, attention_mask=None):
+        outputs = self.backbone(input_ids=input_ids, attention_mask=attention_mask)
+        last_hidden = outputs.last_hidden_state  # (B, T, H)
+
+        if attention_mask is None:
+            pooled = last_hidden.mean(dim=1)
+        else:
+            mask = attention_mask.unsqueeze(-1).expand(last_hidden.size()).float()
+            summed = torch.sum(last_hidden * mask, dim=1)
+            denom = torch.clamp(mask.sum(dim=1), min=1e-9)
+            pooled = summed / denom
+
+        logits = self.classifier(pooled)  # (B, 1)
+        return logits
 
 
-def get_model(force_redownload: bool = False):
+def load_desklib_model(force_redownload: bool = False):
     """
-    Normal load uses cache (fast).
-    If force_redownload=True (from the Nuke button), we wipe cache + re-download.
+    Robust loader:
+    - downloads config/tokenizer normally
+    - downloads model.safetensors explicitly
+    - loads safetensors via safetensors.torch.load_file
+    - loads into our matching PyTorch module with strict=False
     """
     global tokenizer, model
 
-    if (not force_redownload) and model is not None and tokenizer is not None:
+    if (not force_redownload) and tokenizer is not None and model is not None:
         return tokenizer, model
 
     if force_redownload:
-        print("💣 NUKE requested: wiping cache + forcing re-download...")
+        print("💣 NUKE requested: wiping cache + forcing fresh downloads...")
         removed = wipe_model_cache(MODEL_NAME)
         print(f"🧹 Cache dirs removed: {removed}")
         tokenizer = None
         model = None
 
-    print(f"🚀 Loading Model: {MODEL_NAME} on {device}")
+    print(f"🚀 Loading tokenizer/config: {MODEL_NAME}")
+    config = AutoConfig.from_pretrained(MODEL_NAME, force_download=force_redownload)
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, force_download=force_redownload)
 
-    # Tokenizer
-    tokenizer = AutoTokenizer.from_pretrained(
-        MODEL_NAME,
+    print("⬇️ Downloading model.safetensors explicitly...")
+    weights_path = hf_hub_download(
+        repo_id=MODEL_NAME,
+        filename="model.safetensors",
         force_download=force_redownload,
     )
 
-    # Model (prefer safetensors)
-    model = AutoModelForSequenceClassification.from_pretrained(
-        MODEL_NAME,
-        use_safetensors=True,
-        ignore_mismatched_sizes=True,
-        low_cpu_mem_usage=True,
-        force_download=force_redownload,
-    ).to(device).eval()
+    size_gb = os.path.getsize(weights_path) / (1024**3)
+    print(f"✅ model.safetensors path: {weights_path}")
+    print(f"✅ model.safetensors size: {size_gb:.2f} GB")
+
+    # Build model + load weights
+    print("🧠 Building DesklibAIDetectionModel + loading weights...")
+    m = DesklibAIDetectionModel(config)
+    state = load_file(weights_path)  # this will throw if file is truly corrupt
+    missing, unexpected = m.load_state_dict(state, strict=False)
+
+    # Helpful debug (won't crash)
+    if missing:
+        print(f"⚠️ Missing keys (first 20): {missing[:20]}")
+    if unexpected:
+        print(f"⚠️ Unexpected keys (first 20): {unexpected[:20]}")
 
+    model = m.to(device).eval()
     return tokenizer, model
 
 
 # -----------------------------
 # UTILITIES
 # -----------------------------
-ABBR = [
-    "e.g", "i.e", "mr", "mrs", "ms", "dr", "prof", "vs", "etc",
-    "fig", "al", "jr", "sr", "st", "inc", "ltd", "u.s", "u.k"
-]
+ABBR = ["e.g", "i.e", "mr", "mrs", "ms", "dr", "prof", "vs", "etc", "fig", "al", "jr", "sr", "st", "inc", "ltd", "u.s", "u.k"]
 ABBR_REGEX = re.compile(r"\b(" + "|".join(map(re.escape, ABBR)) + r")\.", re.IGNORECASE)
 
-
 def _protect(text):
     text = text.replace("...", "⟨ELLIPSIS⟩")
     text = re.sub(r"(?<=\d)\.(?=\d)", "⟨DECIMAL⟩", text)
     text = ABBR_REGEX.sub(r"\1⟨ABBRDOT⟩", text)
     return text
 
-
 def _restore(text):
     return text.replace("⟨ABBRDOT⟩", ".").replace("⟨DECIMAL⟩", ".").replace("⟨ELLIPSIS⟩", "...")
 
-
 def split_preserving_structure(text):
     blocks = re.split(r"(\n+)", text)
     final_blocks = []
@@ -157,14 +190,11 @@ def analyze(text):
 
     word_count = len(text.split())
     if word_count < 250:
-        warning_msg = (
-            f"⚠️ <b>Insufficient Text:</b> Your input has {word_count} words. "
-            f"Please enter at least 250 words for accurate results."
-        )
+        warning_msg = f"⚠️ <b>Insufficient Text:</b> Your input has {word_count} words. Please enter at least 250 words for accurate results."
         return "Too Short", "N/A", _build_error_card(warning_msg), None, ""
 
     try:
-        tok, mod = get_model(force_redownload=False)
+        tok, mod = load_desklib_model(force_redownload=False)
     except Exception as e:
         return "ERROR", "0%", _build_error_card(f"<b>Failed to load model:</b><br>{str(e)}"), None, ""
 
@@ -186,19 +216,16 @@ def analyze(text):
     for i in range(0, len(windows), batch_size):
         batch = windows[i: i + batch_size]
         inputs = tok(batch, return_tensors="pt", padding=True, truncation=True, max_length=512).to(device)
-        output = mod(**inputs)
-
-        if output.logits.shape[1] > 1:
-            batch_probs = F.softmax(output.logits, dim=-1)[:, 1].detach().cpu().numpy().tolist()
-        else:
-            batch_probs = torch.sigmoid(output.logits).detach().cpu().numpy().flatten().tolist()
+        logits = mod(input_ids=inputs["input_ids"], attention_mask=inputs.get("attention_mask"))
+        batch_probs = torch.sigmoid(logits).detach().cpu().numpy().flatten().tolist()
         probs.extend(batch_probs)
 
     lengths = [len(s.split()) for s in pure_sents]
     total_words = sum(lengths)
     weighted_avg = sum(p * l for p, l in zip(probs, lengths)) / total_words if total_words > 0 else 0
 
-    highlighted_html = "<div style='font-family:sans-serif; line-height:1.8;'>"
+    # HTML Heatmap
+    highlighted_html = "<div style='font-family: sans-serif; line-height: 1.8;'>"
     prob_map = {idx: probs[i] for i, idx in enumerate(pure_sents_indices)}
 
     for i, block in enumerate(blocks):
@@ -216,9 +243,9 @@ def analyze(text):
                 border = "1px solid transparent"
 
             highlighted_html += (
-                f"<span style='background:{bg}; padding:1px 2px; border-radius:3px; border-bottom:{border}; cursor:help;' "
+                f"<span style='background:{bg}; padding:1px 2px; border-radius:3px; border-bottom: {border}; cursor: help;' "
                 f"title='AI Confidence: {score:.2%}'>"
-                f"<span style='color:{color}; font-weight:bold; font-size:0.75em; vertical-align:super; margin-right:2px;'>{score:.0%}</span>"
+                f"<span style='color:{color}; font-weight: bold; font-size: 0.75em; vertical-align: super; margin-right: 2px;'>{score:.0%}</span>"
                 f"{block}</span>"
             )
         else:
@@ -234,24 +261,21 @@ def analyze(text):
 
 
 def nuke_and_reload():
-    """
-    UI button: wipe cache + force re-download + try to load.
-    Returns a status message.
-    """
     try:
-        get_model(force_redownload=True)
+        load_desklib_model(force_redownload=True)
         return (
             "✅ **Nuked cache and reloaded model successfully.**\n\n"
             "- Cache wiped\n"
             "- Fresh download forced\n"
+            "- Custom loader used (DesklibAIDetectionModel)\n"
             "- Model ready ✅"
         )
     except Exception as e:
         return (
             "❌ **Nuke attempted but model still failed to load.**\n\n"
             f"**Error:** `{str(e)}`\n\n"
-            "If this keeps happening, it usually means the downloaded weights are getting truncated "
-            "(network/storage) or the runtime stack (Python/Torch) is incompatible."
+            "If this error happens inside `load_file(model.safetensors)`, the file is truly corrupted/truncated.\n"
+            "If it happens after that, it’s likely key mismatches (shown in logs as missing/unexpected keys)."
         )