Update app.py

app.py

@@ -1,318 +0,0 @@
-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 = "distilbert-base-uncased-finetuned-sst-2-english"
-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()