Update app.py

app.py

@@ -4,10 +4,6 @@ from transformers import AutoTokenizer, AutoModelForSequenceClassification
 import re
 import pandas as pd
 import gradio as gr
-import os
-
-# 1. OPTIMIZE CPU PERFORMANCE
-torch.set_num_threads(os.cpu_count() or 4) # Use all cores
 
 # -----------------------------
 # MODEL INITIALIZATION
@@ -15,20 +11,17 @@ torch.set_num_threads(os.cpu_count() or 4) # Use all cores
 MODEL_NAME = "desklib/ai-text-detector-v1.01"
 tokenizer = None
 model = None
-device = torch.device("cpu") # Hardcoding CPU as requested
-dtype = torch.float32
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 def get_model():
     global tokenizer, model
     if model is None:
-        print(f"Loading Model: {MODEL_NAME} on CPU...")
-        tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, use_fast=False)
+        print(f"Loading Specialized Model: {MODEL_NAME} on {device}")
+        tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+        # Load with default labels; if the model has 2 (Human/AI), we handle it in analyze()
         model = AutoModelForSequenceClassification.from_pretrained(
             MODEL_NAME, 
-            num_labels=1, 
-            ignore_mismatched_sizes=True,
-            low_cpu_mem_usage=True,
-            torch_dtype=dtype
+            ignore_mismatched_sizes=True
         ).to(device).eval()
     return tokenizer, model
 
@@ -69,101 +62,125 @@ def split_preserving_structure(text):
     return final_blocks
 
 # -----------------------------
-# ANALYSIS (With Batching & Progress)
+# ANALYSIS
 # -----------------------------
 @torch.inference_mode()
-def analyze(text, progress=gr.Progress(track_tqdm=True)):
+def analyze(text):
     text = text.strip()
     if not text:
         return "—", "—", "<em>Please enter text...</em>", None
-            
+        
     word_count = len(text.split())
-    if word_count < 300:
-        warning_msg = f"⚠️ <b>Insufficient Text:</b> {word_count} words. Min 300 required."
-        return "Too Short", "N/A", f"<div style='color: #b80d0d; border: 1px solid #b80d0d; padding: 10px;'>{warning_msg}</div>", None
+    if word_count < 300: # Slightly lowered for testing flexibility
+        warning_msg = f"⚠️ <b>Insufficient Text:</b> Your input has {word_count} words. Please enter at least 250-300 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()
     except Exception as e:
-        return "ERROR", "0%", f"Load Error: {str(e)}", None
+        return "ERROR", "0%", f"Failed to load model: {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 "—", "—", "No sentences found.", None
+        return "—", "—", "<em>No sentences detected.</em>", None
 
-    # Windows for context
+    # Sliding Window Generation (Context of 3 sentences)
     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]))
 
-    # 2. BATCHED INFERENCE (Crucial for CPU)
+    # --- BATCHED INFERENCE (Prevents OOM) ---
+    batch_size = 8
     probs = []
-    batch_size = 4 # Small batches so CPU doesn't hang
-    
-    progress(0, desc="Starting Analysis...")
     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)
-        batch_probs = torch.sigmoid(output.logits).cpu().numpy().flatten().tolist()
+        
+        # Check if model is binary classification (2 labels) or regression (1 label)
+        if output.logits.shape[1] > 1:
+            # Assumes Label 1 is 'AI'
+            batch_probs = F.softmax(output.logits, dim=-1)[:, 1].cpu().numpy().tolist()
+        else:
+            batch_probs = torch.sigmoid(output.logits).cpu().numpy().flatten().tolist()
         probs.extend(batch_probs)
-        progress((i + batch_size) / len(windows), desc=f"Analyzing sentences {i+1}-{min(i+batch_size, len(windows))}...")
 
-    # Statistics
+    # Calculation for Final Score
     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
 
-    # 3. HTML GENERATION
-    highlighted_html = "<div style='font-family: sans-serif; line-height: 1.8;'>"
+    # -----------------------------
+    # HTML RECONSTRUCTION
+    # -----------------------------
+    highlighted_html = "<div style='font-family: -apple-system, BlinkMacSystemFont, \"Segoe UI\", Roboto, 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]
-            color, bg = ("#b80d0d", "rgba(184, 13, 13, 0.15)") if score >= THRESHOLD else ("#11823b", "rgba(17, 130, 59, 0.15)")
+            # Color logic based on Threshold
+            if score >= THRESHOLD:
+                color, bg = "#d32f2f", "rgba(211, 47, 47, 0.12)" # Soft Red
+                border = "2px solid #d32f2f"
+            else:
+                color, bg = "#2e7d32", "rgba(46, 125, 50, 0.08)" # Soft Green
+                border = "1px solid transparent"
+            
             highlighted_html += (
-                f"<span style='background:{bg}; padding:2px 4px; border-radius:4px; border-bottom: 2px solid {color};'>"
-                f"<b style='color:{color}; font-size: 0.8em;'>[{score:.1%}]</b> {block}</span>"
+                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 f"{weighted_avg:.1%} AI Probability", f"{weighted_avg:.2%}", highlighted_html, df
+    
+    return label, display_score, highlighted_html, df
 
 # -----------------------------
 # GRADIO INTERFACE
 # -----------------------------
-with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("## 🕵️ AI Detector Pro (CPU Optimized)")
+with gr.Blocks(theme=gr.themes.Soft(), title="AI Detector Pro") as demo:
+    gr.Markdown("# 🕵️ AI Detector Pro")
+    gr.Markdown(f"Utilizing **{MODEL_NAME}**. Values above **{THRESHOLD*100:.0f}%** are flagged as highly likely AI.")
     
     with gr.Row():
         with gr.Column(scale=3):
-            text_input = gr.Textbox(label="Paste Text", lines=12, placeholder="Min 300 words...")
-            run_btn = gr.Button("Run Analysis", variant="primary")
+            text_input = gr.Textbox(label="Input Text", lines=15, placeholder="Paste your essay here (minimum 250 words for accuracy)...")
+            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="Weighted Verdict")
-            score_out = gr.Label(label="Exact Score")
-            
+            verdict_out = gr.Label(label="Global Verdict")
+            score_out = gr.Label(label="Weighted Probability")
+            gr.Markdown("---")
+            gr.Markdown("### How to read:")
+            gr.Markdown("- **Red Highlight:** High AI probability\n- **Green Highlight:** Likely Human\n- **Super-script:** Exact sentence-level AI score")
+    
     with gr.Tabs():
         with gr.TabItem("Visual Heatmap"):
-            html_out = gr.HTML()
-        with gr.TabItem("Detailed Breakdown"):
+            html_out = gr.HTML(label="Heatmap")
+        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])
 
 if __name__ == "__main__":
     demo.launch()