Update app.py

app.py

@@ -13,9 +13,7 @@ MODEL_NAME = "distilgpt2"
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_NAME
-).to(device).eval()
+model = AutoModelForCausalLM.from_pretrained(MODEL_NAME).to(device).eval()
 
 
 # ----------------------------------------------------
@@ -28,7 +26,7 @@ def sentence_split(text):
 
 
 # ----------------------------------------------------
-# PERPLEXITY FUNCTION (very fast on distilgpt2)
+# PERPLEXITY FUNCTION
 # ----------------------------------------------------
 def perplexity(sentence):
     inputs = tokenizer(sentence, return_tensors="pt").to(device)
@@ -50,27 +48,31 @@ def perturb(text):
 
 
 # ----------------------------------------------------
-# DETECTGPT SCORE
+# BASE PERPLEXITY + DETECTGPT SCORE
 # ----------------------------------------------------
-def detectgpt_score(sentence, perturbations=5):
+def detectgpt_base_and_score(sentence, perturbations=3):
+    """
+    Returns:
+      base_perplexity, detectgpt_score
+    """
     try:
         base = perplexity(sentence)
-    except:
-        return 0
+    except Exception:
+        return None, 0.0
 
     pert_scores = []
     for _ in range(perturbations):
         p = perturb(sentence)
         try:
             pert_scores.append(perplexity(p))
-        except:
+        except Exception:
             continue
 
     if not pert_scores:
-        return 0
+        return base, 0.0
 
-    # DetectGPT signal
-    return float(np.mean(pert_scores) - base)
+    score = float(np.mean(pert_scores) - base)
+    return base, score
 
 
 # ----------------------------------------------------
@@ -81,19 +83,55 @@ def classify_text(text):
         return "⚠️ Please enter some text.", None, None
 
     sentences = sentence_split(text)
-    results = []
-    highlighted = []
+    if not sentences:
+        return "⚠️ No valid sentences found.", None, None
+
+    perps = []
     scores = []
+    tmp_results = []
 
+    # 1. Compute base perplexity & DetectGPT score per sentence
     for s in sentences:
-        score = detectgpt_score(s)
+        base_perp, score = detectgpt_base_and_score(s)
+        if base_perp is None:
+            base_perp = float("nan")
+
+        perps.append(base_perp)
         scores.append(score)
+        tmp_results.append({"sentence": s, "perp": base_perp, "score": score})
+
+    # Handle NaNs if any
+    perps_clean = [p for p in perps if not np.isnan(p)]
+    if perps_clean:
+        median_perp = float(np.median(perps_clean))
+    else:
+        median_perp = np.nan
+
+    # 2. Classify using calibrated rule
+    results = []
+    highlighted = []
+    ai_count = 0
+    total = len(tmp_results)
+
+    SCORE_THRESHOLD = 0.05  # Require meaningful positive signal
 
-        label = "AI" if score > 0 else "Human"
+    for item in tmp_results:
+        s = item["sentence"]
+        perp = item["perp"]
+        score = item["score"]
 
-        results.append([s, label, f"{abs(score):.4f}"])
+        # Default label is Human
+        label = "Human"
+
+        # Conditions for AI-like:
+        # - score significantly positive
+        # - perplexity lower than median (more predictable)
+        if not np.isnan(perp) and not np.isnan(median_perp):
+            if (score > SCORE_THRESHOLD) and (perp < median_perp):
+                label = "AI"
 
         if label == "AI":
+            ai_count += 1
             highlighted.append(
                 f"<p style='color:red;font-weight:bold'>{s}</p>"
             )
@@ -102,21 +140,33 @@ def classify_text(text):
                 f"<p style='color:green;font-weight:bold'>{s}</p>"
             )
 
-    # Document-level score (Turnitin-like scaling)
-    avg = np.mean(scores)
-    doc_ai = max(0, min(100, (avg + 1) * 50))
-
-    df = pd.DataFrame(results, columns=["Sentence", "Label", "Score"])
+        results.append([
+            s,
+            label,
+            f"{perp:.2f}" if not np.isnan(perp) else "NaN",
+            f"{score:.4f}"
+        ])
+
+    # 3. Document-level AI percentage = fraction of AI sentences
+    if total > 0:
+        doc_ai_percent = (ai_count / total) * 100.0
+    else:
+        doc_ai_percent = 0.0
+
+    df = pd.DataFrame(
+        results,
+        columns=["Sentence", "Label", "Perplexity", "DetectGPT Score"]
+    )
     html = "\n".join(highlighted)
 
-    return f"⚖️ Document AI Likelihood: {doc_ai:.1f}%", html, df
+    return f"⚖️ Document AI Likelihood (approx): {doc_ai_percent:.1f}%", html, df
 
 
 # ----------------------------------------------------
 # GRADIO UI
 # ----------------------------------------------------
 with gr.Blocks() as demo:
-    gr.Markdown("## 🧠 Writenix DetectGPT (distilgpt2 — Ultra Fast Version)")
+    gr.Markdown("## 🧠 Writenix DetectGPT (Calibrated, distilgpt2)")
 
     text_input = gr.Textbox(
         label="Enter text",
@@ -128,7 +178,7 @@ with gr.Blocks() as demo:
 
     ai_score = gr.Label(label="Overall AI Likelihood")
     highlighted = gr.HTML()
-    table = gr.Dataframe(headers=["Sentence", "Label", "Score"], wrap=True)
+    table = gr.Dataframe(headers=["Sentence", "Label", "Perplexity", "DetectGPT Score"], wrap=True)
 
     classify_btn.click(classify_text, text_input, [ai_score, highlighted, table])