Update app.py

app.py

@@ -15,12 +15,12 @@ dtype = torch.bfloat16 if (device.type=="cuda" and torch.cuda.is_bf16_supported(
 model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME, dtype=dtype).to(device).eval()
 
 # -----------------------------
-# INTERNAL THRESHOLD for sentence labels/colors
+# THRESHOLD FOR LABEL COLOR
 # -----------------------------
-THRESHOLD = 0.70  # used only for per-sentence "AI/Human" tags & color
+THRESHOLD = 0.70
 
 # -----------------------------
-# SENTENCE SPLITTER (protect → split → restore; no lookbehinds)
+# SENTENCE SPLITTING UTILITIES
 # -----------------------------
 ABBR = [
     "e.g", "i.e", "mr", "mrs", "ms", "dr", "prof", "vs", "etc", "fig", "al",
@@ -33,10 +33,10 @@ def _protect(text: str) -> str:
     t = text.strip()
     if not t:
         return ""
-    t = re.sub(r"\s*\n+\s*", " ", t)            # normalize newlines
-    t = t.replace("...", "⟨ELLIPSIS⟩")          # ellipses
-    t = re.sub(r"(?<=\d)\.(?=\d)", "⟨DECIMAL⟩", t)  # decimals like 3.14
-    t = ABBR_REGEX.sub(r"\1⟨ABBRDOT⟩", t)       # abbreviations' dot
+    t = re.sub(r"\s*\n+\s*", " ", t)
+    t = t.replace("...", "⟨ELLIPSIS⟩")
+    t = re.sub(r"(?<=\d)\.(?=\d)", "⟨DECIMAL⟩", t)
+    t = ABBR_REGEX.sub(r"\1⟨ABBRDOT⟩", t)
     return t
 
 def _restore(text: str) -> str:
@@ -49,7 +49,6 @@ def sentence_split(text: str):
     t = _protect(text)
     if not t:
         return []
-    # split on [.?!] followed by whitespace and likely sentence start or end
     parts = re.split(r"([.?!])\s+(?=(?:[\"“”‘’']?\s*[A-Z(])|$)", t)
     sentences, buf = [], ""
     for i, chunk in enumerate(parts):
@@ -57,63 +56,95 @@ def sentence_split(text: str):
             buf += chunk
         else:
             buf += chunk
-            sentences.append(buf.strip()); buf = ""
+            sentences.append(buf.strip())
+            buf = ""
     if buf.strip():
         sentences.append(buf.strip())
     return [_restore(s).strip() for s in sentences if s.strip()]
 
 # -----------------------------
-# CORE: overall AI score + highlights
+# GROUP SENTENCES (TURNITIN STYLE)
+# -----------------------------
+def group_sentences(sents, size=3):
+    grouped = []
+    for i in range(0, len(sents), size):
+        chunk = " ".join(sents[i:i+size])
+        grouped.append(chunk)
+    return grouped
+
+# -----------------------------
+# CORE ANALYSIS
 # -----------------------------
 def analyze(text, max_len=512):
     sents = sentence_split(text)
     if not sents:
         return "—", "—", "<em>Paste some text to analyze.</em>", None
 
-    # light clean (per model card vibe)
-    clean_sents = [re.sub(r"\s+", " ", s).strip() for s in sents]
+    # GROUP sentences into 3-sentence chunks
+    grouped = group_sentences(sents, size=3)
+    clean_grouped = [re.sub(r"\s+", " ", g).strip() for g in grouped]
 
+    # tokenize grouped
     inputs = tokenizer(
-        clean_sents, return_tensors="pt", padding=True, truncation=True, max_length=max_len
+        clean_grouped, return_tensors="pt", padding=True,
+        truncation=True, max_length=max_len
     ).to(device)
 
+    # model inference
     with torch.no_grad():
         logits = model(**inputs).logits
-        probs = F.softmax(logits, dim=-1)  # [:,0]=Human, [:,1]=AI
+        chunk_probs = F.softmax(logits, dim=-1)[:, 1].detach().cpu().tolist()
 
-    ai_probs = probs[:, 1].detach().cpu().tolist()
+    # EXPAND chunk-level probabilities back to per-sentence
+    ai_probs = []
+    for idx, prob in enumerate(chunk_probs):
+        start = idx * 3
+        end = min(start + 3, len(sents))
+        for _ in range(start, end):
+            ai_probs.append(prob)
+
+    # overall AI score
     overall_ai = sum(ai_probs) / len(ai_probs)
     overall_pct = f"{overall_ai * 100:.1f}%"
-    overall_label = "🤖 Likely AI Written" if overall_ai >= THRESHOLD else "🧒 Likely Human Written"
+    overall_label = (
+        "🤖 Likely AI Written" if overall_ai >= THRESHOLD else "🧒 Likely Human Written"
+    )
 
-    # Per-sentence highlights (use THRESHOLD only for the tag/color)
+    # HIGHLIGHTS + TABLE
     rows, highlights = [], []
     for i, orig in enumerate(sents, start=1):
         ai_p = float(ai_probs[i-1])
         label = "AI" if ai_p >= THRESHOLD else "Human"
         pct = f"{ai_p*100:.1f}%"
+
+        # color logic
         if ai_p < 0.30:
-            color = "#11823b"      # green
+            color = "#11823b"
         elif ai_p < 0.70:
-            color = "#b8860b"      # amber
+            color = "#b8860b"
         else:
-            color = "#b80d0d"      # red
+            color = "#b80d0d"
+
         normalized = re.sub(r"\s+", " ", orig)
         highlights.append(
-            "<div style='margin:6px 0; padding:6px 8px; border-radius:6px; background:rgba(0,0,0,0.03)'>"
-            f"<strong style='color:{color}'>[{pct} {label}]</strong> {normalized}</div>"
+            "<div style='margin:6px 0; padding:6px 8px; border-radius:6px; "
+            "background:rgba(0,0,0,0.03)'>"
+            f"<strong style='color:{color}'>[{pct} {label}]</strong> "
+            f"{normalized}</div>"
         )
+
         rows.append([i, orig, round(ai_p, 4), label])
 
-    html = "\n".join(highlights)
     df = pd.DataFrame(rows, columns=["#", "Sentence", "AI_Prob", "Label"])
+    html = "\n".join(highlights)
+
     return overall_label, overall_pct, html, df
 
 # -----------------------------
-# GRADIO UI (verdict + score, plus highlights)
+# GRADIO UI
 # -----------------------------
 with gr.Blocks() as demo:
-    gr.Markdown("### 🕵️ AI Written Text Detector — Fakespot Model")
+    gr.Markdown("### 🕵️ AI Written Text Detector — Fakespot Model (Turnitin-Style)")
 
     text_input = gr.Textbox(label="Paste text", lines=14, placeholder="Your content…")
     btn = gr.Button("Analyze")