"""
AI Writing Analyzer — sentence-level heat map for human vs. AI-generated text.
Built for classroom use. Loads a RoBERTa-based ChatGPT detector from
Hugging Face and runs it on each sentence independently, then renders the
input text with per-sentence color coding indicating the probability that
the sentence was AI-generated.
Runs comfortably on the free CPU tier.
"""
import re
import html
import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
# ---------------------------------------------------------------------------
# Model
# ---------------------------------------------------------------------------
# Hello-SimpleAI's RoBERTa detector — small, CPU-friendly, widely used.
MODEL_NAME = "Hello-SimpleAI/chatgpt-detector-roberta"
print(f"Loading model: {MODEL_NAME}")
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME)
model.eval()
# The model's label order: index 0 = Human, index 1 = ChatGPT/AI.
# (Confirmed from the model card's id2label.)
AI_INDEX = 1
# ---------------------------------------------------------------------------
# Sentence splitting
# ---------------------------------------------------------------------------
_SENT_SPLIT_RE = re.compile(r"(?<=[.!?])\s+(?=[A-Z\"'\(\[])")
def split_sentences(text: str):
"""Lightweight sentence splitter — no NLTK download needed on free CPU."""
text = text.strip()
if not text:
return []
# First split on paragraph breaks to preserve structure, then sentences.
chunks = []
for para in re.split(r"\n\s*\n", text):
para = para.strip()
if not para:
continue
parts = _SENT_SPLIT_RE.split(para)
parts = [p.strip() for p in parts if p.strip()]
chunks.extend(parts)
return chunks
# ---------------------------------------------------------------------------
# Scoring
# ---------------------------------------------------------------------------
@torch.no_grad()
def score_sentence(sentence: str) -> float:
"""Return probability that `sentence` is AI-generated (0.0 – 1.0)."""
inputs = tokenizer(
sentence,
return_tensors="pt",
truncation=True,
max_length=512,
)
logits = model(**inputs).logits
probs = torch.softmax(logits, dim=-1)[0]
return float(probs[AI_INDEX].item())
# ---------------------------------------------------------------------------
# Rendering
# ---------------------------------------------------------------------------
def prob_to_color(p: float) -> str:
"""
Map probability 0..1 to a background color.
Low (human) -> cool teal
Mid -> amber
High (AI) -> warm red
"""
# Interpolate between three stops in RGB.
if p < 0.5:
t = p / 0.5
r = int(56 + (245 - 56) * t)
g = int(189 + (191 - 189) * t)
b = int(248 + (66 - 248) * t)
else:
t = (p - 0.5) / 0.5
r = int(245 + (248 - 245) * t)
g = int(191 + (80 - 191) * t)
b = int(66 + (80 - 66) * t)
# Higher opacity for a vivid highlight; text is forced light on top.
return f"rgba({r}, {g}, {b}, 0.42)"
def border_color(p: float) -> str:
if p < 0.5:
t = p / 0.5
r = int(56 + (245 - 56) * t)
g = int(189 + (191 - 189) * t)
b = int(248 + (66 - 248) * t)
else:
t = (p - 0.5) / 0.5
r = int(245 + (248 - 245) * t)
g = int(191 + (80 - 191) * t)
b = int(66 + (80 - 66) * t)
return f"rgba({r}, {g}, {b}, 0.95)"
def render_heatmap(sentences, scores) -> str:
if not sentences:
return (
""
"Paste some writing above and click Analyze to see a "
"sentence-by-sentence breakdown.
"
)
pieces = []
for sent, p in zip(sentences, scores):
bg = prob_to_color(p)
bd = border_color(p)
pct = int(round(p * 100))
safe = html.escape(sent)
pieces.append(
f"{safe} "
f"{pct}%"
)
body = " ".join(pieces)
avg = sum(scores) / len(scores)
verdict, vcolor = classify_overall(avg)
summary = (
f""
f"
Overall assessment
"
f"
"
f"{verdict}
"
f"
"
f"Avg. AI likelihood: {int(round(avg*100))}% "
f" · Sentences: {len(sentences)}
"
f"
"
"Legend:"
"Likely human"
"Uncertain"
"Likely AI"
"
"
)
return (
f""
f"{summary}"
f"
{body}
"
f"{legend}"
f"