Update app.py

app.py

@@ -8,22 +8,9 @@ MODEL = "roberta-base-openai-detector"
 tokenizer = AutoTokenizer.from_pretrained(MODEL)
 model = AutoModelForSequenceClassification.from_pretrained(MODEL)
 
-def get_color(ai_score):
-    """Convert AI score (0-1) into green-yellow-red gradient."""
-    # Green (0,255,0) -> Yellow (255,255,0) -> Red (255,0,0)
-    if ai_score < 0.5:  # Green to Yellow
-        red = int(ai_score * 510)
-        green = 255
-    else:  # Yellow to Red
-        red = 255
-        green = int(255 - (ai_score - 0.5) * 510)
-    return f"rgb({red},{green},0)"
-
 def detect_ai(text):
-    # Split into paragraphs (by double newlines), fallback to sentences if no paragraphs
+    # Split into paragraphs (double newlines)
     paragraphs = re.split(r"\n\s*\n", text.strip())
-    if len(paragraphs) == 1:  # no clear paragraphs
-        paragraphs = re.split(r'(?<=[.!?]) +', text)
 
     results = []
     for para in paragraphs:
@@ -33,41 +20,43 @@ def detect_ai(text):
         with torch.no_grad():
             outputs = model(**inputs)
             probs = torch.softmax(outputs.logits, dim=1)
-        ai_score = float(probs[0][1])
 
-        # Add label for clarity
-        if ai_score < 0.3:
-            label = "🟢 Likely Human"
-        elif ai_score < 0.7:
-            label = "🟡 Mixed"
+        ai_score = float(probs[0][1])  # AI likelihood
+        human_score = 1 - ai_score     # Human likelihood
+
+        # Decide label: if <90% human → AI
+        if human_score < 0.9:
+            label = "🔴 AI"
+            color = "rgb(255,120,120)"  # red
         else:
-            label = "🔴 Likely AI"
+            label = "🟢 Human"
+            color = "rgb(120,255,120)"  # green
 
-        results.append({"text": para, "ai_score": ai_score, "label": label})
+        results.append({"paragraph": para, "ai_score": ai_score, "human_score": human_score, "label": label, "color": color})
 
     # Build highlighted HTML
     highlighted = ""
     for r in results:
-        color = get_color(r['ai_score'])
         highlighted += (
-            f"<div style='background-color:{color}; padding:6px; margin-bottom:6px; border-radius:6px'>"
-            f"<b>{r['label']} ({round(r['ai_score']*100,1)}%)</b><br>{r['text']}</div>"
+            f"<div style='background-color:{r['color']}; padding:8px; margin-bottom:6px; border-radius:6px'>"
+            f"<b>{r['label']} — Human {round(r['human_score']*100,1)}% | AI {round(r['ai_score']*100,1)}%</b><br>"
+            f"{r['paragraph']}</div>"
         )
 
-    # Compute total AI probability (average)
+    # Compute overall human %
     if results:
-        avg_ai = sum(r['ai_score'] for r in results) / len(results)
-        total_percent = round(avg_ai * 100, 2)
-        highlighted += f"<p><b>⚖️ Overall AI Probability: {total_percent}%</b></p>"
+        avg_human = sum(r['human_score'] for r in results) / len(results)
+        total_human = round(avg_human * 100, 2)
+        highlighted += f"<p><b>⚖️ Overall Human Probability: {total_human}%</b></p>"
     else:
-        total_percent = 0.0
+        total_human = 0.0
 
-    return highlighted, {"paragraphs": results, "total_ai_percent": total_percent}
+    return highlighted, {"paragraphs": results, "overall_human_percent": total_human}
 
 with gr.Blocks() as demo:
     gr.Markdown("## 🤖 AI Detector (Paragraph-level)")
-    gr.Markdown("Paste your text below. Green = Human-like, Yellow = Mixed, Red = AI-like.")
-    input_text = gr.Textbox(lines=10, placeholder="Enter text here...")
+    gr.Markdown("Paragraphs with <90% human probability are flagged as **AI**.")
+    input_text = gr.Textbox(lines=12, placeholder="Paste your essay or report here...")
     output_html = gr.HTML()
     output_json = gr.JSON()
     run_btn = gr.Button("Detect AI")