Update app.py

app.py

@@ -0,0 +1,318 @@
+import os
+import re
+import shutil
+
+# ============================================================
+# ENV (set BEFORE transformers/hub usage)
+# ============================================================
+os.environ.setdefault("HF_HOME", "/tmp/hf")
+os.environ.setdefault("HUGGINGFACE_HUB_CACHE", "/tmp/hf/hub")
+os.environ.setdefault("TRANSFORMERS_CACHE", "/tmp/hf/transformers")
+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
+# -----------------------------
+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.
+    Returns number of cache directories removed.
+    """
+    safe = model_id.replace("/", "--")
+    candidates = [
+        # our /tmp cache (recommended)
+        f"/tmp/hf/hub/models--{safe}",
+        f"/tmp/hf/transformers/models--{safe}",
+        # default home cache (in case something wrote there)
+        os.path.expanduser(f"~/.cache/huggingface/hub/models--{safe}"),
+        os.path.expanduser(f"~/.cache/huggingface/transformers/models--{safe}"),
+        os.path.expanduser(f"~/.cache/huggingface/modules/models--{safe}"),
+    ]
+
+    removed = 0
+    for path in candidates:
+        if os.path.exists(path):
+            shutil.rmtree(path, ignore_errors=True)
+            removed += 1
+    return removed
+
+
+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 load_desklib_model(force_redownload: bool = False):
+    """
+    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 tokenizer is not None and model is not None:
+        return tokenizer, model
+
+    if force_redownload:
+        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 tokenizer/config: {MODEL_NAME}")
+    config = AutoConfig.from_pretrained(MODEL_NAME, force_download=force_redownload)
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, force_download=force_redownload)
+
+    print("⬇️ Downloading model.safetensors explicitly...")
+    weights_path = hf_hub_download(
+        repo_id=MODEL_NAME,
+        filename="model.safetensors",
+        force_download=force_redownload,
+    )
+
+    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_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 = []
+    for block in blocks:
+        if not block:
+            continue
+        if block.startswith("\n"):
+            final_blocks.append(block)
+        else:
+            protected = _protect(block)
+            parts = re.split(r"([.?!])(\s+)", protected)
+            for i in range(0, len(parts), 3):
+                sentence = parts[i]
+                punct = parts[i + 1] if i + 1 < len(parts) else ""
+                space = parts[i + 2] if i + 2 < len(parts) else ""
+                if sentence.strip():
+                    final_blocks.append(_restore(sentence + punct))
+                if space:
+                    final_blocks.append(space)
+    return final_blocks
+
+
+# -----------------------------
+# ANALYSIS
+# -----------------------------
+@torch.inference_mode()
+def analyze(text):
+    text = (text or "").strip()
+    if not text:
+        return "—", "—", "<em>Please enter text...</em>", None, ""
+
+    word_count = len(text.split())
+    if word_count < 250:
+        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 = 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, ""
+
+    blocks = split_preserving_structure(text)
+    pure_sents_indices = [i for i, b in enumerate(blocks) if b.strip() and not b.startswith("\n")]
+    pure_sents = [blocks[i] for i in pure_sents_indices]
+
+    if not pure_sents:
+        return "—", "—", "<em>No sentences detected.</em>", None, ""
+
+    windows = []
+    for i in range(len(pure_sents)):
+        start = max(0, i - 1)
+        end = min(len(pure_sents), i + 2)
+        windows.append(" ".join(pure_sents[start:end]))
+
+    batch_size = 8
+    probs = []
+    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)
+        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
+
+    # 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):
+        if block.startswith("\n") or block.isspace():
+            highlighted_html += block.replace("\n", "<br>")
+            continue
+
+        if i in prob_map:
+            score = prob_map[i]
+            if score >= THRESHOLD:
+                color, bg = "#d32f2f", "rgba(211, 47, 47, 0.12)"
+                border = "2px solid #d32f2f"
+            else:
+                color, bg = "#2e7d32", "rgba(46, 125, 50, 0.08)"
+                border = "1px solid transparent"
+
+            highlighted_html += (
+                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"{block}</span>"
+            )
+        else:
+            highlighted_html += block
+
+    highlighted_html += "</div>"
+
+    label = f"{weighted_avg:.1%} AI Written"
+    display_score = f"{weighted_avg:.2%}"
+    df = pd.DataFrame({"Sentence": pure_sents, "AI Confidence": [f"{p:.2%}" for p in probs]})
+
+    return label, display_score, highlighted_html, df, ""
+
+
+def nuke_and_reload():
+    try:
+        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 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)."
+        )
+
+
+# -----------------------------
+# INTERFACE
+# -----------------------------
+with gr.Blocks(theme=gr.themes.Soft(), title="AI Detector Pro") as demo:
+    gr.Markdown("# 🕵️ AI Detector Pro")
+    gr.Markdown(f"Model: **{MODEL_NAME}** | Highlight Threshold: **{THRESHOLD*100:.0f}%**")
+
+    with gr.Row():
+        with gr.Column(scale=3):
+            text_input = gr.Textbox(label="Input Text", lines=15, placeholder="Enter at least 250 words...")
+            with gr.Row():
+                clear_btn = gr.Button("Clear")
+                run_btn = gr.Button("Analyze Text", variant="primary")
+                nuke_btn = gr.Button("💣 Nuke Model Cache", variant="stop")
+
+        with gr.Column(scale=1):
+            verdict_out = gr.Label(label="Global Verdict")
+            score_out = gr.Label(label="Weighted Probability")
+
+    status_out = gr.Markdown()
+
+    with gr.Tabs():
+        with gr.TabItem("Visual Heatmap"):
+            html_out = gr.HTML()
+        with gr.TabItem("Data Breakdown"):
+            table_out = gr.Dataframe(headers=["Sentence", "AI Confidence"], wrap=True)
+
+    run_btn.click(analyze, inputs=text_input, outputs=[verdict_out, score_out, html_out, table_out, status_out])
+
+    def _clear():
+        return "", "—", "—", "<em>Please enter text...</em>", None, ""
+
+    clear_btn.click(_clear, outputs=[text_input, verdict_out, score_out, html_out, table_out, status_out])
+    nuke_btn.click(nuke_and_reload, outputs=status_out)
+
+if __name__ == "__main__":
+    demo.launch()