Update app.py

app.py

@@ -1,10 +1,12 @@
+import os
+import re
+import shutil
 import torch
 import torch.nn.functional as F
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
-import re
 import pandas as pd
 import gradio as gr
-import os
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+
 
 # -----------------------------
 # MODEL INITIALIZATION
@@ -14,51 +16,109 @@ tokenizer = None
 model = None
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
+
+def purge_model_cache(model_id: str) -> None:
+    """
+    Remove cached weights/tokenizer for this model from common HF cache locations.
+    This fixes the 'state dictionary ... corrupted' error caused by partial downloads.
+    """
+    safe = model_id.replace("/", "--")
+
+    candidates = [
+        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}"),
+    ]
+
+    for path in candidates:
+        if os.path.exists(path):
+            try:
+                shutil.rmtree(path, ignore_errors=True)
+                print(f"🧹 Removed cache: {path}")
+            except Exception as e:
+                print(f"⚠️ Failed to remove cache at {path}: {e}")
+
+
 def get_model():
+    """
+    Loads tokenizer + model with safetensors preferred.
+    If load fails (often due to corrupted HF cache), purge cache + force download.
+    """
     global tokenizer, model
-    if model is None:
-        try:
-            print(f"🚀 Loading Model: {MODEL_NAME} on {device}")
-            tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-            # added use_safetensors=True and low_cpu_mem_usage for stability
-            model = AutoModelForSequenceClassification.from_pretrained(
-                MODEL_NAME, 
-                ignore_mismatched_sizes=True,
-                low_cpu_mem_usage=True
-            ).to(device).eval()
-        except Exception as e:
-            print(f"⚠️ Initial load failed: {e}. Attempting recovery...")
-            # If it fails, it might be a corrupted cache. 
-            # This forces a clean redownload of the weights.
-            model = AutoModelForSequenceClassification.from_pretrained(
-                MODEL_NAME,
-                ignore_mismatched_sizes=True,
-                force_download=True 
-            ).to(device).eval()
-    return tokenizer, model
-
-THRESHOLD = 0.59 
+
+    if model is not None and tokenizer is not None:
+        return tokenizer, model
+
+    print(f"🚀 Loading Model: {MODEL_NAME} on {device}")
+
+    try:
+        tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+
+        model = AutoModelForSequenceClassification.from_pretrained(
+            MODEL_NAME,
+            use_safetensors=True,          # ✅ prefer safetensors
+            ignore_mismatched_sizes=True,
+            low_cpu_mem_usage=True,
+        ).to(device).eval()
+
+        return tokenizer, model
+
+    except Exception as e:
+        print(f"⚠️ Initial load failed: {e}")
+        print("🔁 Attempting recovery: purge cache + force re-download...")
+
+        purge_model_cache(MODEL_NAME)
+
+        # Redownload everything cleanly
+        tokenizer = AutoTokenizer.from_pretrained(
+            MODEL_NAME,
+            force_download=True,
+        )
+
+        model = AutoModelForSequenceClassification.from_pretrained(
+            MODEL_NAME,
+            use_safetensors=True,          # ✅ keep safetensors on recovery too
+            ignore_mismatched_sizes=True,
+            force_download=True,
+        ).to(device).eval()
+
+        return tokenizer, model
+
+
+THRESHOLD = 0.59
+
 
 # -----------------------------
 # 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):
+
+def _protect(text: str) -> str:
     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):
+def _restore(text: str) -> str:
+    return (
+        text.replace("⟨ABBRDOT⟩", ".")
+            .replace("⟨DECIMAL⟩", ".")
+            .replace("⟨ELLIPSIS⟩", "...")
+    )
+
+
+def split_preserving_structure(text: str):
     blocks = re.split(r"(\n+)", text)
     final_blocks = []
     for block in blocks:
-        if not block: continue
+        if not block:
+            continue
         if block.startswith("\n"):
             final_blocks.append(block)
         else:
@@ -66,27 +126,36 @@ def split_preserving_structure(text):
             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 ""
+                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.strip()
+    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", f"<div style='color: #b80d0d; padding: 20px; border: 1px solid #b80d0d; border-radius: 8px;'>{warning_msg}</div>", None
+        warning_msg = (
+            f"⚠️ <b>Insufficient Text:</b> Your input has {word_count} words. "
+            f"Please enter at least 250 words for accurate results."
+        )
+        return (
+            "Too Short",
+            "N/A",
+            f"<div style='color:#b80d0d; padding:20px; border:1px solid #b80d0d; border-radius:8px;'>{warning_msg}</div>",
+            None,
+        )
 
     try:
         tok, mod = get_model()
@@ -96,7 +165,7 @@ def analyze(text):
     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
 
@@ -108,30 +177,43 @@ def analyze(text):
 
     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)
+        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].cpu().numpy().tolist()
+            batch_probs = F.softmax(output.logits, dim=-1)[:, 1].detach().cpu().numpy().tolist()
         else:
-            batch_probs = torch.sigmoid(output.logits).cpu().numpy().flatten().tolist()
+            batch_probs = torch.sigmoid(output.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
+    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;'>"
+    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:
@@ -140,49 +222,59 @@ def analyze(text):
             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"<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:
             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
 
+
 # -----------------------------
 # 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...")
+            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")
-        
+
         with gr.Column(scale=1):
             verdict_out = gr.Label(label="Global Verdict")
             score_out = gr.Label(label="Weighted Probability")
-    
+
     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])
-    clear_btn.click(lambda: ["", "", "", "", None], outputs=[text_input, verdict_out, score_out, html_out, table_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])
 
 if __name__ == "__main__":
     demo.launch()