Update app.py

app.py

@@ -2,41 +2,35 @@ import gradio as gr
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 import torch
 import re
-import matplotlib
-matplotlib.use("Agg")
 import matplotlib.pyplot as plt
 from tokenizers.normalizers import Sequence, Replace, Strip
 from tokenizers import Regex
 
-# -------------------------
-# Device setup
-# -------------------------
+# ---- Device setup ----
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
-# -------------------------
-# Model and Tokenizer Setup
-# -------------------------
+# ---- Model and Tokenizer Setup ----
 model1_path = "modernbert.bin"
 model2_path = "https://huggingface.co/mihalykiss/modernbert_2/resolve/main/Model_groups_3class_seed12"
 model3_path = "https://huggingface.co/mihalykiss/modernbert_2/resolve/main/Model_groups_3class_seed22"
 
 tokenizer = AutoTokenizer.from_pretrained("answerdotai/ModernBERT-base")
 
-def safe_load_model(base_name, weights_path):
-    model = AutoModelForSequenceClassification.from_pretrained(base_name, num_labels=41)
-    state_dict = torch.hub.load_state_dict_from_url(weights_path, map_location=device) if weights_path.startswith("http") else torch.load(weights_path, map_location=device)
-    model.load_state_dict(state_dict)
-    model.to(device).eval()
-    return model
-
-print("Loading models...")
-model_1 = safe_load_model("answerdotai/ModernBERT-base", model1_path)
-model_2 = safe_load_model("answerdotai/ModernBERT-base", model2_path)
-model_3 = safe_load_model("answerdotai/ModernBERT-base", model3_path)
-
-# -------------------------
-# Label Mapping
-# -------------------------
+# Load models
+model_1 = AutoModelForSequenceClassification.from_pretrained("answerdotai/ModernBERT-base", num_labels=41)
+model_1.load_state_dict(torch.load(model1_path, map_location=device))
+model_1.to(device).eval()
+
+model_2 = AutoModelForSequenceClassification.from_pretrained("answerdotai/ModernBERT-base", num_labels=41)
+model_2.load_state_dict(torch.hub.load_state_dict_from_url(model2_path, map_location=device))
+model_2.to(device).eval()
+
+model_3 = AutoModelForSequenceClassification.from_pretrained("answerdotai/ModernBERT-base", num_labels=41)
+model_3.load_state_dict(torch.hub.load_state_dict_from_url(model3_path, map_location=device))
+model_3.to(device).eval()
+
+
+# ---- Label Mapping ----
 label_mapping = {
     0: '13B', 1: '30B', 2: '65B', 3: '7B', 4: 'GLM130B', 5: 'bloom_7b',
     6: 'bloomz', 7: 'cohere', 8: 'davinci', 9: 'dolly', 10: 'dolly-v2-12b',
@@ -50,9 +44,7 @@ label_mapping = {
     39: 'text-davinci-002', 40: 'text-davinci-003'
 }
 
-# -------------------------
-# Text Cleaning
-# -------------------------
+# ---- Text Cleaning ----
 def clean_text(text: str) -> str:
     text = re.sub(r'\s{2,}', ' ', text)
     text = re.sub(r'\s+([,.;:?!])', r'\1', text)
@@ -65,21 +57,20 @@ tokenizer.backend_tokenizer.normalizer = Sequence([
     Strip()
 ])
 
-# -------------------------
-# Classification Function
-# -------------------------
+
+# ---- Classification Function ----
 def classify_text(text):
     cleaned_text = clean_text(text)
     if not cleaned_text.strip():
-        return "<b style='color:red;'>Please enter some text to analyze.</b>", None
+        return "**Error:** Please enter some text to analyze.", None
 
+    # Split into paragraphs
     paragraphs = [p.strip() for p in re.split(r'\n{2,}', cleaned_text) if p.strip()]
     chunk_scores = []
     all_probabilities = []
 
     for i, paragraph in enumerate(paragraphs):
         inputs = tokenizer(paragraph, return_tensors="pt", truncation=True, padding=True).to(device)
-
         with torch.no_grad():
             logits_1 = model_1(**inputs).logits
             logits_2 = model_2(**inputs).logits
@@ -88,51 +79,50 @@ def classify_text(text):
             softmax_1 = torch.softmax(logits_1, dim=1)
             softmax_2 = torch.softmax(logits_2, dim=1)
             softmax_3 = torch.softmax(logits_3, dim=1)
-            averaged_probabilities = (softmax_1 + softmax_2 + softmax_3) / 3
-            probabilities = averaged_probabilities[0]
-            all_probabilities.append(probabilities.cpu())
 
-        human_prob = probabilities[24].item()
-        ai_probs_clone = probabilities.clone()
-        ai_probs_clone[24] = 0
-        ai_total_prob = ai_probs_clone.sum().item()
+            avg_probs = (softmax_1 + softmax_2 + softmax_3) / 3
+            probs = avg_probs[0]
+            all_probabilities.append(probs.cpu())
 
-        total = human_prob + ai_total_prob
+        human_prob = probs[24].item()
+        ai_probs_clone = probs.clone()
+        ai_probs_clone[24] = 0
+        ai_total = ai_probs_clone.sum().item()
+        total = human_prob + ai_total
         human_pct = (human_prob / total) * 100
-        ai_pct = (ai_total_prob / total) * 100
+        ai_pct = (ai_total / total) * 100
         ai_model = label_mapping[torch.argmax(ai_probs_clone).item()]
 
         chunk_scores.append({
-            "paragraph": paragraph[:150] + ("..." if len(paragraph) > 150 else ""),
             "human": human_pct,
             "ai": ai_pct,
-            "model": ai_model
+            "model": ai_model,
+            "text": paragraph[:200].replace('\n', ' ') + ("..." if len(paragraph) > 200 else "")
         })
 
-    # --- Overall ---
+    # ---- Overall Averages ----
     avg_human = sum(c["human"] for c in chunk_scores) / len(chunk_scores)
     avg_ai = sum(c["ai"] for c in chunk_scores) / len(chunk_scores)
+
     if avg_human > avg_ai:
         result_message = f"**Overall Result:** <span class='highlight-human'>{avg_human:.2f}% Human-written</span>"
     else:
         top_model = max(chunk_scores, key=lambda c: c['ai'])['model']
         result_message = f"**Overall Result:** <span class='highlight-ai'>{avg_ai:.2f}% AI-generated (likely {top_model})</span>"
 
-    # --- Paragraph Breakdown ---
-    paragraph_html = "<h3>Paragraph Analysis:</h3>"
-    for idx, c in enumerate(chunk_scores, 1):
-        color = "#4CAF50" if c['human'] > c['ai'] else "#FF5733"
-        paragraph_html += f"""
-        <div style='margin-bottom:10px; border-left:4px solid {color}; padding-left:10px;'>
-        <b>Paragraph {idx}</b>: {c['human']:.2f}% Human | {c['ai']:.2f}% AI → <i>{c['model']}</i><br>
-        <small>{c['paragraph']}</small></div>
-        """
-
-    # --- Plot ---
+    # ---- Paragraph Analysis (Markdown Clean) ----
+    paragraph_text = "\n\n**Paragraph Analysis:**\n"
+    for i, c in enumerate(chunk_scores, 1):
+        paragraph_text += (
+            f"**Paragraph {i}:** {c['human']:.2f}% Human | {c['ai']:.2f}% AI → *{c['model']}*\n"
+            f"{c['text']}\n\n"
+        )
+
+    # ---- Top 5 Models Plot ----
     mean_probs = torch.mean(torch.stack(all_probabilities), dim=0)
-    top_5_probs, top_5_indices = torch.topk(mean_probs, 5)
+    top_5_probs, top_5_idx = torch.topk(mean_probs, 5)
     top_5_probs = top_5_probs.cpu().numpy()
-    top_5_labels = [label_mapping[i.item()] for i in top_5_indices]
+    top_5_labels = [label_mapping[i.item()] for i in top_5_idx]
 
     fig, ax = plt.subplots(figsize=(10, 5))
     bars = ax.barh(top_5_labels, top_5_probs, color='#4CAF50')
@@ -144,30 +134,31 @@ def classify_text(text):
         ax.text(width + 0.005, bar.get_y() + bar.get_height() / 2, f'{width:.2%}', va='center')
     plt.tight_layout()
 
-    return result_message + "<br><br>" + paragraph_html, fig
+    return result_message + "\n\n" + paragraph_text, fig
 
 
-# -------------------------
-# UI Setup
-# -------------------------
+# ---- UI Setup ----
 title = "AI Text Detector"
 description = """
-This tool uses <b>ModernBERT</b> to detect AI-generated text.<br>
+This tool uses **ModernBERT** to detect AI-generated text.  
 Each paragraph is analyzed separately to show which parts are likely AI-generated.
 """
 bottom_text = "**Developed by SzegedAI – Extended by Saber**"
 
 AI_texts = [
-"Artificial intelligence (AI) is reshaping industries by automating tasks, enhancing decision-making, and driving innovation. From predictive analytics in finance to autonomous vehicles in transportation, AI technologies are becoming integral to daily operations."
+"Artificial intelligence (AI) is reshaping industries by automating tasks, enhancing decision-making, and driving innovation. From predictive analytics in finance to autonomous vehicles in transportation, AI technologies are becoming integral to daily operations. The future of AI lies not only in technological advancement but also in ensuring ethical use, transparency, and accountability."
 ]
 
 Human_texts = [
-"Mathematics has always been a cornerstone of scientific discovery. It provides a precise language for describing natural phenomena, from the orbit of planets to the behavior of subatomic particles."
+"Mathematics has always been a cornerstone of scientific discovery. It provides a precise language for describing natural phenomena, from the orbit of planets to the behavior of subatomic particles. The beauty of mathematics lies in its universality—its principles hold true regardless of context or culture."
 ]
 
 iface = gr.Blocks(css="""
     @import url('https://fonts.googleapis.com/css2?family=Roboto+Mono:wght@400;700&display=swap');
-    body { font-family: 'Roboto Mono', sans-serif !important; }
+    #text_input_box { border-radius: 10px; border: 2px solid #4CAF50; font-size: 18px; padding: 15px; margin-bottom: 20px; width: 60%; margin: auto; }
+    #result_output_box { border-radius: 10px; border: 2px solid #4CAF50; font-size: 16px; padding: 15px; margin-top: 20px; width: 80%; margin: auto; }
+    body { font-family: 'Roboto Mono', sans-serif !important; padding: 20px; }
+    .gradio-container { border: 1px solid #4CAF50; border-radius: 15px; padding: 30px; box-shadow: 0px 0px 10px rgba(0,255,0,0.4); max-width: 900px; margin: auto; }
     .highlight-human { color: #4CAF50; font-weight: bold; }
     .highlight-ai { color: #FF5733; font-weight: bold; }
 """)
@@ -175,18 +166,14 @@ iface = gr.Blocks(css="""
 with iface:
     gr.Markdown(f"# {title}")
     gr.Markdown(description)
-    text_input = gr.Textbox(label="", placeholder="Paste your article here...", lines=10)
-    analyze_btn = gr.Button("🔍 Analyze", variant="primary")
-    result_output = gr.HTML(label="Result")
+    text_input = gr.Textbox(label="", placeholder="Paste your article here...", elem_id="text_input_box", lines=10)
+    result_output = gr.HTML("", elem_id="result_output_box")
     plot_output = gr.Plot(label="Model Probability Distribution")
-
-    analyze_btn.click(classify_text, inputs=text_input, outputs=[result_output, plot_output])
-
+    text_input.change(classify_text, inputs=text_input, outputs=[result_output, plot_output])
     with gr.Tab("AI Examples"):
         gr.Examples(AI_texts, inputs=text_input)
     with gr.Tab("Human Examples"):
         gr.Examples(Human_texts, inputs=text_input)
-
     gr.Markdown(bottom_text)
 
 iface.launch(share=True)