Update app.py

app.py

@@ -14,141 +14,178 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 dtype = torch.bfloat16 if (device.type=="cuda" and torch.cuda.is_bf16_supported()) else torch.float32
 model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME, dtype=dtype).to(device).eval()
 
-# -----------------------------
-# THRESHOLD
-# -----------------------------
 THRESHOLD = 0.80
 
 # -----------------------------
-# SENTENCE SPLITTING
+# ABBREVIATION PROTECTION
 # -----------------------------
-ABBR = ["e.g", "i.e", "mr", "mrs", "ms", "dr", "prof", "vs", "etc", "fig", "al",
-    "jr", "sr", "st", "no", "vol", "pp", "mt", "inc", "ltd", "co", "u.s", "u.k",
-    "a.m", "p.m"
+ABBR = [
+    "e.g", "i.e", "mr", "mrs", "ms", "dr", "prof", "vs", "etc", "fig", "al",
+    "jr", "sr", "st", "no", "vol", "pp", "mt", "inc", "ltd", "co",
+    "u.s", "u.k", "a.m", "p.m"
 ]
-ABBR_REGEX = re.compile(r"\b(" + "|".join(map(re.escape, ABBR)) + r")\.", flags=re.IGNORECASE)
+ABBR_REGEX = re.compile(r"\b(" + "|".join(map(re.escape, ABBR)) + r")\.", re.IGNORECASE)
 
 def _protect(text):
-    t = text.replace("...", "⟨ELLIPSIS⟩")
-    t = re.sub(r"(?<=\d)\.(?=\d)", "⟨DECIMAL⟩", t)
-    t = ABBR_REGEX.sub(r"\1⟨ABBRDOT⟩", t)
-    return t
+    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⟩", "...")
+    return (
+        text.replace("⟨ABBRDOT⟩", ".")
+            .replace("⟨DECIMAL⟩", ".")
+            .replace("⟨ELLIPSIS⟩", "...")
+    )
 
-def split_sentences_preserving(text):
-    protected = _protect(text)
-    parts = re.split(r"([.?!])(\s+)", protected)
+# -----------------------------
+# PERFECT PARAGRAPH-PRESERVING SPLITTER
+# -----------------------------
+def split_preserving_structure(text):
+    """
+    Splits text into:
+    - EXACT newline blocks (\n, \n\n, etc.)
+    - Sentences inside non-newline blocks
+    """
+    blocks = re.split(r"(\n+)", text)  # keep newline separators
+    final_blocks = []
 
-    sentences = []
-    current = ""
+    for block in blocks:
+        if block.startswith("\n"):
+            final_blocks.append(block)  # preserve EXACT paragraph spacing
+        else:
+            protected = _protect(block)
+            parts = re.split(r"([.?!])(\s+)", protected)
 
-    for i in range(0, len(parts), 3):
-        part = parts[i]
-        punct = parts[i+1] if i+1 < len(parts) else ""
-        space = parts[i+2] if i+2 < len(parts) else ""
+            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 ""
 
-        current = part + punct
-        sentences.append(_restore(current))
-        sentences.append(space)  # preserve exact spacing (spaces and newlines)
+                whole = sentence + punct
+                if whole.strip():
+                    final_blocks.append(_restore(whole))
 
-    return sentences  # alternating [sentence, whitespace, sentence, whitespace...]
+                if space:
+                    final_blocks.append(space)
 
-def extract_pure_sentences(sent_block):
-    return [s for s in sent_block if not s.isspace()]
+    return final_blocks
 
 
+def extract_sentences_only(blocks):
+    """Return only sentence blocks (no whitespace/newlines)."""
+    return [
+        b for b in blocks
+        if b.strip() != "" and not b.startswith("\n") and not b.isspace()
+    ]
+
 # -----------------------------
-# GROUP SENTENCES
+# GROUPING
 # -----------------------------
 def group_sentences(sents, size=3):
-    return [" ".join(sents[i:i+size]) for i in range(0, len(sents), size)]
-
+    return [" ".join(sents[i:i + size]) for i in range(0, len(sents), size)]
 
 # -----------------------------
-# MAIN ANALYSIS
+# ANALYSIS LOGIC
 # -----------------------------
 def analyze(text, max_len=512):
 
-    # 1. Split while preserving structure
-    blocks = split_sentences_preserving(text)
-    pure_sentences = extract_pure_sentences(blocks)
+    # Structured block split
+    blocks = split_preserving_structure(text)
+    pure_sentences = extract_sentences_only(blocks)
 
     if not pure_sentences:
         return "—", "—", "<em>Paste text to analyze.</em>", None
 
-    # 2. Group for model
+    # Group into 3-sentence windows (Turnitin style)
     grouped = group_sentences(pure_sentences, 3)
     clean_grouped = [re.sub(r"\s+", " ", g).strip() for g in grouped]
 
-    # 3. Run model
-    inputs = tokenizer(clean_grouped, return_tensors="pt", padding=True,
-                       truncation=True, max_length=max_len).to(device)
+    # Run model
+    inputs = tokenizer(clean_grouped, return_tensors="pt",
+                       padding=True, truncation=True,
+                       max_length=max_len).to(device)
 
     with torch.no_grad():
         logits = model(**inputs).logits
         chunk_probs = F.softmax(logits, dim=-1)[:, 1].cpu().tolist()
 
-    # 4. Expand chunk scores to per-sentence
-    sentence_ai = []
+    # Expand group scores back to individual sentences
+    ai_scores = []
     for idx, prob in enumerate(chunk_probs):
         start = idx * 3
         end = min(start + 3, len(pure_sentences))
         for _ in range(start, end):
-            sentence_ai.append(prob)
+            ai_scores.append(prob)
 
     # -----------------------------
-    # FINAL OUTPUT RECONSTRUCTION
+    # RECONSTRUCT ORIGINAL TEXT W/ HIGHLIGHTING
     # -----------------------------
     highlighted = ""
-    sent_index = 0
+    current_sentence = 0
 
     for block in blocks:
+
+        # newline block → keep EXACT
+        if block.startswith("\n"):
+            highlighted += block
+            continue
+
+        # whitespace block → keep
         if block.isspace():
-            highlighted += block  # preserve exact spacing
-        else:
-            # this block is a real sentence
-            ai_p = sentence_ai[sent_index]
-            sent_index += 1
+            highlighted += block
+            continue
 
-            pct = f"{ai_p*100:.1f}%"
+        # real sentence → highlight
+        ai_p = ai_scores[current_sentence]
+        current_sentence += 1
 
-            if ai_p < 0.30:
-                color = "#11823b"
-            elif ai_p < 0.70:
-                color = "#b8860b"
-            else:
-                color = "#b80d0d"
+        pct = f"{ai_p * 100:.1f}%"
+
+        if ai_p < 0.30:
+            color = "#11823b"
+        elif ai_p < 0.70:
+            color = "#b8860b"
+        else:
+            color = "#b80d0d"
 
-            highlighted += f"<span style='background-color:rgba(0,0,0,0.03); padding:3px 4px; border-radius:4px;'><strong style='color:{color}'>[{pct}]</strong> {block.strip()}</span> "
+        highlighted += (
+            f"<span style='background:rgba(0,0,0,0.03); padding:3px 4px; "
+            f"border-radius:4px;'><strong style='color:{color}'>[{pct}]</strong> "
+            f"{block.strip()}</span>"
+        )
 
-    # Overall score
-    overall = sum(sentence_ai) / len(sentence_ai)
-    overall_pct = f"{overall*100:.1f}%"
+        # maintain spacing after sentence
+        highlighted += " "
+
+    # -----------------------------
+    # OVERALL SCORE
+    # -----------------------------
+    overall = sum(ai_scores) / len(ai_scores)
+    overall_pct = f"{overall * 100:.1f}%"
     overall_label = "🤖 Likely AI Written" if overall >= THRESHOLD else "🧒 Likely Human Written"
 
     # Table output
     df = pd.DataFrame(
-        [[i+1, s, sentence_ai[i]] for i, s in enumerate(pure_sentences)],
+        [[i + 1, s, ai_scores[i]] for i, s in enumerate(pure_sentences)],
         columns=["#", "Sentence", "AI_Prob"]
     )
 
     return overall_label, overall_pct, highlighted, df
 
-
 # -----------------------------
 # UI
 # -----------------------------
 with gr.Blocks() as demo:
-    gr.Markdown("### 🕵️ AI Sentence-Level Detector — Original Format Highlighting")
+    gr.Markdown("### 🕵️ AI Sentence-Level Detector — Exact Structure Highlighting")
 
-    text_input = gr.Textbox(label="Paste text", lines=14)
+    text_input = gr.Textbox(label="Paste text", lines=14, placeholder="Your text…")
     btn = gr.Button("Analyze")
 
-    verdict = gr.Label(label="Overall Verdict")
-    score = gr.Label(label="Overall AI Score")
-    highlights = gr.HTML(label="Highlighted Text (Original Format)")
+    verdict = gr.Label(label="Verdict (Overall)")
+    score = gr.Label(label="AI Score")
+    highlights = gr.HTML(label="Highlighted Text (Exact Structure)")
     table = gr.Dataframe(headers=["#", "Sentence", "AI_Prob"], wrap=True)
 
     btn.click(analyze, inputs=[text_input], outputs=[verdict, score, highlights, table])