Update app.py

app.py

@@ -1,150 +1,222 @@
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM
-import numpy as np
-import pandas as pd
+import torch.nn.functional as F
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
 import re
+import math
+import pandas as pd
 import gradio as gr
 
-# ----------------------------------------------
-# LOAD FAST MODEL (DistilGPT2)
-# ----------------------------------------------
-MODEL_NAME = "distilgpt2"
-
+# -----------------------------
+# MODEL
+# -----------------------------
+MODEL_NAME = "openai-community/roberta-base-openai-detector"
 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()
-
-
-# ----------------------------------------------
-# SENTENCE SPLITTER
-# ----------------------------------------------
-def sentence_split(text):
-    text = text.replace("\n", ". ")
-    s = re.split(r'(?<=[.!?])\s+', text)
-    return [x.strip() for x in s if x.strip()]
-
-
-# ----------------------------------------------
-# PERPLEXITY
-# ----------------------------------------------
-def perplexity(sentence):
-    enc = tokenizer(sentence, return_tensors="pt").to(device)
-    with torch.no_grad():
-        out = model(**enc, labels=enc["input_ids"])
-    return float(torch.exp(out.loss))
-
-
-# ----------------------------------------------
-# TOKEN-LEVEL ENTROPY
-# ----------------------------------------------
-def token_entropy(sentence):
-    enc = tokenizer(sentence, return_tensors="pt").to(device)
-    input_ids = enc["input_ids"][0]
-
-    with torch.no_grad():
-        outputs = model(enc["input_ids"], labels=enc["input_ids"])
-        logits = outputs.logits[0]
-
-    entropies = []
-    for i in range(1, len(input_ids)):
-        probs = torch.softmax(logits[i-1], dim=-1)
-        entropy = -torch.sum(probs * torch.log(probs + 1e-10))
-        entropies.append(float(entropy))
-
-    return np.mean(entropies), np.std(entropies)
-
-
-# ----------------------------------------------
-# TURNITIN-STYLE SCORING PIPELINE
-# ----------------------------------------------
-def analyze_sentence(sentence):
-    perp = perplexity(sentence)
-    mean_ent, std_ent = token_entropy(sentence)
-    length = len(sentence.split())
-    punct = sum([sentence.count(p) for p in ".,;:!?"])
-
-    return {
-        "sentence": sentence,
-        "perplexity": perp,
-        "entropy_mean": mean_ent,
-        "entropy_std": std_ent,
-        "length": length,
-        "punctuation": punct
-    }
-
-
-# ----------------------------------------------
-# MAIN TURNITIN STYLE DETECTOR
-# ----------------------------------------------
-def classify_text(text):
-
-    sentences = sentence_split(text)
-    stats = [analyze_sentence(s) for s in sentences]
-
-    df = pd.DataFrame(stats)
-
-    # ---------- TURNITIN STYLE METRICS ----------
-    perplexity_mean = df["perplexity"].mean()
-    perplexity_std  = df["perplexity"].std()
-
-    entropy_mean = df["entropy_mean"].mean()
-    entropy_std  = df["entropy_std"].mean()
-
-    length_std = df["length"].std()
-    punct_std  = df["punctuation"].std()
-
-    # ---------- NORMALIZED SCORES ----------
-    # Low variance = AI-like
-    burstiness_score = np.exp(-perplexity_std)
-
-    entropy_smoothness = np.exp(-entropy_std)
-
-    length_uniformity = np.exp(-length_std / (df["length"].mean() + 1e-5))
-    punct_uniformity  = np.exp(-punct_std / (df["punctuation"].mean() + 1e-5))
-
-    # ---------- ENSEMBLE SCORE (Turnitin-like) ----------
-    ai_score = (
-        0.35 * burstiness_score +
-        0.25 * entropy_smoothness +
-        0.20 * length_uniformity +
-        0.20 * punct_uniformity
-    )
-
-    ai_percent = float(ai_score * 100)
-
-    # ---------- PER-SENTENCE LABELS ----------
-    highlighted = []
-    for i, row in df.iterrows():
-        is_ai = row["perplexity"] < perplexity_mean * 0.75 and row["entropy_std"] < entropy_std * 0.8
-        if is_ai:
-            highlighted.append(f"<p style='color:red;font-weight:bold'>{row['sentence']}</p>")
+dtype = torch.bfloat16 if (device.type=="cuda" and torch.cuda.is_bf16_supported()) else torch.float32
+model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME, torch_dtype=dtype).to(device).eval()
+
+# -----------------------------
+# SENTENCE SPLITTER (robust, no externals)
+# -----------------------------
+_ABBR = r"(?: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)\."
+_QUOTE = r"[\"“”‘’']?"
+# Split on ., ?, ! when followed by space/newline + a capital/quote or end of text,
+# while avoiding common abbreviations and decimals.
+_SENT_PAT = re.compile(
+    rf"""
+    (?<!\b{_ABBR})           # not common abbreviation
+    (?<!\d)\.|\?|!           # ., ?, !
+    (?=\s+{_QUOTE}[A-Z(]|$)  # lookahead for next sentence start or end
+    """,
+    re.VERBOSE
+)
+
+def sentence_split(text: str):
+    # Normalize hard breaks to spaces (Turnitin-like continuous flow)
+    t = re.sub(r"\s*\n+\s*", " ", text.strip())
+    if not t:
+        return []
+    # Temporarily protect ellipses to avoid over-splitting
+    t = t.replace("...", "…")
+    pieces = []
+    start = 0
+    for m in _SENT_PAT.finditer(t):
+        end = m.end()
+        chunk = t[start:end].strip()
+        if chunk:
+            pieces.append(chunk)
+        start = end
+    # tail
+    tail = t[start:].strip()
+    if tail:
+        pieces.append(tail)
+    # Restore ellipses
+    return [s.replace("…", "...") for s in pieces]
+
+# -----------------------------
+# UTILITIES
+# -----------------------------
+def batched(iterable, n=64):
+    for i in range(0, len(iterable), n):
+        yield iterable[i:i+n], i
+
+def contig_spans(labels):
+    """Return (num_spans, longest_span_len) for consecutive 'AI' labels."""
+    longest = 0
+    count = 0
+    run = 0
+    for lab in labels:
+        if lab == "AI":
+            run += 1
+            longest = max(longest, run)
         else:
-            highlighted.append(f"<p style='color:green;font-weight:bold'>{row['sentence']}</p>")
-
-    html = "\n".join(highlighted)
-
-    # Display readable columns
-    df_display = df[["sentence", "perplexity", "entropy_mean", "entropy_std", "length", "punctuation"]]
-
-    return f"⚖️ Estimated AI Probability (Turnitin-style): {ai_percent:.1f}%", html, df_display
+            if run > 0:
+                count += 1
+            run = 0
+    if run > 0:
+        count += 1
+    return count, longest
+
+def verdict_from_stats(flag_pct, longest_span, avg_ai_prob):
+    """
+    Turnitin-ish qualitative summary.
+    - Emphasize consecutive AI-like sentences (spans) and overall prevalence.
+    """
+    if flag_pct >= 85 and longest_span >= 6 and avg_ai_prob >= 0.80:
+        return "⚠️ Highly likely AI-generated (long consecutive spans and high prevalence)."
+    if flag_pct >= 60 and longest_span >= 4:
+        return "⚠️ Strong AI signals (multiple/long spans)."
+    if flag_pct >= 30 or longest_span >= 3:
+        return "△ Some AI indicators (partial/short spans)."
+    return "✓ No clear AI indication (by this detector)."
+
+# -----------------------------
+# CORE CLASSIFIER
+# -----------------------------
+def classify_sentences(text, ai_threshold=0.70, batch_size=64, max_len=512):
+    sents = sentence_split(text)
+    if not sents:
+        return [], [], 0.0, 0.0, (0, 0)
+
+    all_probs = []
+    all_labels = []
+
+    for chunk, base in batched(sents, n=batch_size):
+        inputs = tokenizer(
+            chunk,
+            return_tensors="pt",
+            padding=True,
+            truncation=True,
+            max_length=max_len
+        ).to(device)
+        with torch.no_grad():
+            logits = model(**inputs).logits
+            probs = F.softmax(logits, dim=-1)  # [:, 0]=Human, [:, 1]=AI
+
+        ai_probs = probs[:, 1].detach().cpu().tolist()
+        all_probs.extend(ai_probs)
+
+    for p in all_probs:
+        all_labels.append("AI" if p >= ai_threshold else "Human")
+
+    avg_ai_prob = float(sum(all_probs) / len(all_probs))
+    flagged_pct = 100.0 * sum(1 for l in all_labels if l == "AI") / len(all_labels)
+    spans = contig_spans(all_labels)
+
+    rows = []
+    for i, (s, p, lab) in enumerate(zip(sents, all_probs, all_labels), start=1):
+        rows.append({
+            "Sentence #": i,
+            "Sentence": s,
+            "AI Probability": round(p, 4),
+            "Label": lab
+        })
+
+    return sents, rows, avg_ai_prob, flagged_pct, spans
+
+# -----------------------------
+# HTML HIGHLIGHT (Turnitin-ish)
+# -----------------------------
+def color_for_prob(p):
+    # 0-0.3 green, 0.3-0.7 yellow, 0.7-1.0 red
+    if p < 0.30:
+        return "#11823b"
+    if p < 0.70:
+        return "#b8860b"
+    return "#b80d0d"
+
+def build_highlight_html(rows):
+    blocks = []
+    for r in rows:
+        p = r["AI Probability"]
+        col = color_for_prob(p)
+        pct = f"{p*100:.1f}%"
+        text = re.sub(r"\s+", " ", r["Sentence"]).strip()
+        blocks.append(
+            f"<span style='background:rgba(0,0,0,0.02); "
+            f"padding:4px 6px; border-radius:6px; display:block; margin:6px 0;'>"
+            f"<strong style='color:{col}'>[{pct} {r['Label']}]</strong> {text}</span>"
+        )
+    return "\n".join(blocks)
+
+# -----------------------------
+# PUBLIC API FOR GRADIO
+# -----------------------------
+def generate_report(text, threshold):
+    if not text or not text.strip():
+        return "⚠️ Please enter some text.", None, None, None
+
+    sents, rows, avg_ai_prob, flagged_pct, (span_count, longest_span) = classify_sentences(
+        text, ai_threshold=threshold
+    )
 
+    verdict = verdict_from_stats(flagged_pct, longest_span, avg_ai_prob)
+    overall = (
+        f"⚖️ Turnitin-style Summary\n"
+        f"- Overall AI probability (avg per sentence): {avg_ai_prob*100:.1f}%\n"
+        f"- Sentences flagged as AI ≥ {int(threshold*100)}%: {flagged_pct:.1f}%\n"
+        f"- Consecutive AI spans: {span_count} (longest: {longest_span})\n"
+        f"- Verdict: {verdict}\n"
+        f"\nⓘ This is an approximation using an open detector; "
+        f"actual Turnitin results may differ."
+    )
 
+    html = build_highlight_html(rows)
+    df = pd.DataFrame(rows, columns=["Sentence #", "Sentence", "AI Probability", "Label"])
+    return overall, html, df, f"{flagged_pct:.1f}%"
 
-# ----------------------------------------------
+# -----------------------------
 # GRADIO UI
-# ----------------------------------------------
+# -----------------------------
 with gr.Blocks() as demo:
-    gr.Markdown("## 🧠 Writenix — Turnitin-Style AI Detector")
-
-    text_input = gr.Textbox(label="Enter text", lines=10, placeholder="Paste your essay...")
-
-    classify_btn = gr.Button("🚀 Analyze")
-
-    ai_score = gr.Label(label="Turnitin-Style AI Likelihood")
-    highlighted = gr.HTML()
-    table = gr.Dataframe(headers=["Sentence", "Perplexity", "Entropy Mean", "Entropy Std", "Length", "Punctuation"], wrap=True)
-
-    classify_btn.click(classify_text, text_input, [ai_score, highlighted, table])
+    gr.Markdown("## 🧭 Writenix AI Detector — Turnitin-style (Sentence-Level)")
+
+    with gr.Row():
+        text_input = gr.Textbox(
+            label="Paste your content",
+            lines=16,
+            placeholder="Drop your essay/article here…"
+        )
+
+    with gr.Row():
+        threshold = gr.Slider(
+            0.50, 0.95, value=0.70, step=0.01,
+            label="AI Flag Threshold (probability ≥ threshold ⇒ AI)"
+        )
+        detect_btn = gr.Button("🔎 Analyze")
+
+    with gr.Row():
+        ai_summary = gr.Textbox(label="Report Summary", lines=8)
+        flagged_pct = gr.Label(label="% Sentences Flagged (AI)")
+
+    highlighted = gr.HTML(label="Per-Sentence Highlights")
+    table = gr.Dataframe(headers=["Sentence #", "Sentence", "AI Probability", "Label"], wrap=True)
+
+    detect_btn.click(
+        fn=generate_report,
+        inputs=[text_input, threshold],
+        outputs=[ai_summary, highlighted, table, flagged_pct]
+    )
 
 if __name__ == "__main__":
     demo.launch()