Update app.py

app.py

@@ -9,16 +9,23 @@ tokenizer = AutoTokenizer.from_pretrained(MODEL)
 model = AutoModelForSequenceClassification.from_pretrained(MODEL)
 
 def get_color(ai_score):
-    """Convert AI score (0-1) into a smooth green-yellow-red gradient."""
-    red = int(ai_score * 255)
-    green = int((1 - ai_score) * 255)
+    """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 by paragraphs instead of sentences
+    # Split into paragraphs (by double newlines), fallback to sentences if no paragraphs
     paragraphs = re.split(r"\n\s*\n", text.strip())
-    results = []
+    if len(paragraphs) == 1:  # no clear paragraphs
+        paragraphs = re.split(r'(?<=[.!?]) +', text)
 
+    results = []
     for para in paragraphs:
         if not para.strip():
             continue
@@ -27,19 +34,31 @@ def detect_ai(text):
             outputs = model(**inputs)
             probs = torch.softmax(outputs.logits, dim=1)
         ai_score = float(probs[0][1])
-        results.append({"paragraph": para, "ai_score": ai_score})
+
+        # Add label for clarity
+        if ai_score < 0.3:
+            label = "🟢 Likely Human"
+        elif ai_score < 0.7:
+            label = "🟡 Mixed"
+        else:
+            label = "🔴 Likely AI"
+
+        results.append({"text": para, "ai_score": ai_score, "label": label})
 
     # 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:4px'>{r['paragraph']}</div>"
+        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>"
+        )
 
-    # Compute total AI percentage (average across paragraphs)
+    # Compute total AI probability (average)
     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>Total AI Probability: {total_percent}%</b></p>"
+        highlighted += f"<p><b>⚖️ Overall AI Probability: {total_percent}%</b></p>"
     else:
         total_percent = 0.0
 
@@ -47,8 +66,8 @@ def detect_ai(text):
 
 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=8, placeholder="Enter text here with paragraphs separated by blank lines...")
+    gr.Markdown("Paste your text below. Green = Human-like, Yellow = Mixed, Red = AI-like.")
+    input_text = gr.Textbox(lines=10, placeholder="Enter text here...")
     output_html = gr.HTML()
     output_json = gr.JSON()
     run_btn = gr.Button("Detect AI")