Update app.py

app.py

@@ -1,153 +1,192 @@
+import gradio as gr
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
 import torch
 import re
-import io
-import base64
 import matplotlib
 matplotlib.use("Agg")
 import matplotlib.pyplot as plt
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
-import gradio as gr
-
-# ===============================
-# Safe model loading
-# ===============================
-device = "cuda" if torch.cuda.is_available() else "cpu"
-
-def safe_load_model(name):
-    try:
-        return AutoModelForSequenceClassification.from_pretrained(name).to(device)
-    except Exception as e:
-        print(f"[WARN] Failed to load {name}: {e}")
-        return None
+from tokenizers.normalizers import Sequence, Replace, Strip
+from tokenizers import Regex
+
+# -------------------------
+# Device setup
+# -------------------------
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+# -------------------------
+# 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("roberta-base-openai-detector")
-model_2 = safe_load_model("roberta-large-openai-detector")
-model_3 = safe_load_model("Hello-SimpleAI/chatgpt-detector-roberta")
-tokenizer = AutoTokenizer.from_pretrained("roberta-base-openai-detector")
-
-label_mapping = {i: f"model_{i}" for i in range(25)}
-label_mapping[24] = "Human"
-
-# ===============================
-# Helper functions
-# ===============================
-def clean_text(text):
-    return re.sub(r'\s+', ' ', text).strip()
-
-def plot_to_base64(fig):
-    buf = io.BytesIO()
-    fig.savefig(buf, format="png", bbox_inches="tight")
-    buf.seek(0)
-    img_base64 = base64.b64encode(buf.getvalue()).decode("utf-8")
-    plt.close(fig)
-    return f"<img src='data:image/png;base64,{img_base64}' style='max-width:100%; border-radius:8px;'>"
-
-# ===============================
-# Main classification logic
-# ===============================
+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
+# -------------------------
+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',
+    11: 'flan_t5_base', 12: 'flan_t5_large', 13: 'flan_t5_small',
+    14: 'flan_t5_xl', 15: 'flan_t5_xxl', 16: 'gemma-7b-it', 17: 'gemma2-9b-it',
+    18: 'gpt-3.5-turbo', 19: 'gpt-35', 20: 'gpt4', 21: 'gpt4o',
+    22: 'gpt_j', 23: 'gpt_neox', 24: 'human', 25: 'llama3-70b', 26: 'llama3-8b',
+    27: 'mixtral-8x7b', 28: 'opt_1.3b', 29: 'opt_125m', 30: 'opt_13b',
+    31: 'opt_2.7b', 32: 'opt_30b', 33: 'opt_350m', 34: 'opt_6.7b',
+    35: 'opt_iml_30b', 36: 'opt_iml_max_1.3b', 37: 't0_11b', 38: 't0_3b',
+    39: 'text-davinci-002', 40: 'text-davinci-003'
+}
+
+# -------------------------
+# Text Cleaning
+# -------------------------
+def clean_text(text: str) -> str:
+    text = re.sub(r'\s{2,}', ' ', text)
+    text = re.sub(r'\s+([,.;:?!])', r'\1', text)
+    return text
+
+newline_to_space = Replace(Regex(r'\s*\n\s*'), " ")
+tokenizer.backend_tokenizer.normalizer = Sequence([
+    tokenizer.backend_tokenizer.normalizer,
+    newline_to_space,
+    Strip()
+])
+
+# -------------------------
+# Classification Function
+# -------------------------
 def classify_text(text):
     cleaned_text = clean_text(text)
-    if not cleaned_text:
-        return "<b style='color:red'>Please enter some text.</b>"
-
-    paragraphs = re.split(r'\n{2,}', cleaned_text)
-    if len(paragraphs) == 1 and len(cleaned_text.split()) > 300:
-        words = cleaned_text.split()
-        paragraphs = [' '.join(words[i:i + 300]) for i in range(0, len(words), 300)]
+    if not cleaned_text.strip():
+        return "<b style='color:red;'>Please enter some text to analyze.</b>", None
 
-    paragraph_scores = []
+    paragraphs = [p.strip() for p in re.split(r'\n{2,}', cleaned_text) if p.strip()]
+    chunk_scores = []
     all_probabilities = []
 
-    for i, para in enumerate(paragraphs):
-        inputs = tokenizer(para, return_tensors="pt", truncation=True, padding=True).to(device)
-        softmax_outputs = []
+    for i, paragraph in enumerate(paragraphs):
+        inputs = tokenizer(paragraph, return_tensors="pt", truncation=True, padding=True).to(device)
 
-        for m in [model_1, model_2, model_3]:
-            if m is None:
-                continue
-            with torch.no_grad():
-                logits = m(**inputs).logits
-                softmax_outputs.append(torch.softmax(logits, dim=1))
+        with torch.no_grad():
+            logits_1 = model_1(**inputs).logits
+            logits_2 = model_2(**inputs).logits
+            logits_3 = model_3(**inputs).logits
 
-        if not softmax_outputs:
-            return "<b style='color:red'>Error: No models loaded successfully.</b>"
-
-        avg_probs = sum(softmax_outputs) / len(softmax_outputs)
-        probabilities = avg_probs[0]
-        all_probabilities.append(probabilities.cpu())
+            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 = probabilities.clone()
-        ai_probs[24] = 0
-        ai_prob = ai_probs.sum().item()
+        ai_probs_clone = probabilities.clone()
+        ai_probs_clone[24] = 0
+        ai_total_prob = ai_probs_clone.sum().item()
 
-        total = human_prob + ai_prob
+        total = human_prob + ai_total_prob
         human_pct = (human_prob / total) * 100
-        ai_pct = (ai_prob / total) * 100
-        ai_model = label_mapping[torch.argmax(ai_probs).item()]
+        ai_pct = (ai_total_prob / total) * 100
+        ai_model = label_mapping[torch.argmax(ai_probs_clone).item()]
 
-        preview = para[:180].strip() + ("..." if len(para) > 180 else "")
-        paragraph_scores.append({
-            "id": i + 1,
+        chunk_scores.append({
+            "paragraph": paragraph[:150] + ("..." if len(paragraph) > 150 else ""),
             "human": human_pct,
             "ai": ai_pct,
-            "model": ai_model,
-            "preview": preview
+            "model": ai_model
         })
 
-    avg_human = sum(p["human"] for p in paragraph_scores) / len(paragraph_scores)
-    avg_ai = sum(p["ai"] for p in paragraph_scores) / len(paragraph_scores)
-
+    # --- Overall ---
+    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:
-        overall = f"<b>Overall Result:</b> <span class='highlight-human'>{avg_human:.2f}% Human-written</span>"
+        result_message = f"**Overall Result:** <span class='highlight-human'>{avg_human:.2f}% Human-written</span>"
     else:
-        top_model = max(paragraph_scores, key=lambda p: p["ai"])["model"]
-        overall = f"<b>Overall Result:</b> <span class='highlight-ai'>{avg_ai:.2f}% AI-generated (likely {top_model})</span>"
+        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>
+        """
 
-    # --- Top 5 chart ---
+    # --- Plot ---
     mean_probs = torch.mean(torch.stack(all_probabilities), dim=0)
     top_5_probs, top_5_indices = torch.topk(mean_probs, 5)
-    labels = [label_mapping[i.item()] for i in top_5_indices]
-    values = top_5_probs.cpu().numpy()
+    top_5_probs = top_5_probs.cpu().numpy()
+    top_5_labels = [label_mapping[i.item()] for i in top_5_indices]
 
-    fig, ax = plt.subplots(figsize=(8, 4))
-    ax.barh(labels, values, color='#4CAF50')
-    ax.set_xlabel("Probability")
-    ax.set_title("Top 5 Model Predictions")
+    fig, ax = plt.subplots(figsize=(10, 5))
+    bars = ax.barh(top_5_labels, top_5_probs, color='#4CAF50')
+    ax.set_xlabel('Probability')
+    ax.set_title('Top 5 Model Predictions')
     ax.invert_yaxis()
-    chart_html = plot_to_base64(fig)
-
-    # --- Paragraph breakdown ---
-    html = f"<div style='font-family:Arial, sans-serif; line-height:1.6'>{overall}<br><br>"
-    html += "<h3>Paragraph Analysis:</h3>"
-
-    for p in paragraph_scores:
-        color = "#28a745" if p["human"] > p["ai"] else "#FF5733"
-        html += f"""
-        <div style='margin-bottom:10px; border-left:5px solid {color}; padding-left:10px; background:#f9f9f9; border-radius:6px;'>
-            <b>Paragraph {p["id"]}</b>: {p["human"]:.2f}% Human | {p["ai"]:.2f}% AI → <i>{p["model"]}</i><br>
-            <small>{p["preview"]}</small>
-        </div>
-        """
+    for bar in bars:
+        width = bar.get_width()
+        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
 
-    html += "<br><h3>Top 5 Models:</h3>" + chart_html + "</div>"
-    return html
 
-# ===============================
-# Gradio UI
-# ===============================
-css = """
-.highlight-ai { color: #FF5733; font-weight: bold; }
-.highlight-human { color: #28a745; font-weight: bold; }
+# -------------------------
+# UI Setup
+# -------------------------
+title = "AI Text Detector"
+description = """
+This tool uses <b>ModernBERT</b> to detect AI-generated text.<br>
+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."
+]
+
+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."
+]
+
+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; }
+    .highlight-human { color: #4CAF50; font-weight: bold; }
+    .highlight-ai { color: #FF5733; font-weight: bold; }
+""")
+
+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")
+    plot_output = gr.Plot(label="Model Probability Distribution")
+
+    analyze_btn.click(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)
 
-with gr.Blocks(css=css, theme="soft") as demo:
-    gr.Markdown("# 🧠 AI vs Human Text Detector")
-    txt = gr.Textbox(label="Paste your article", lines=12, placeholder="Enter your full text here...")
-    btn = gr.Button("Analyze", variant="primary")
-    out = gr.HTML(label="Results", elem_id="result-box")
-    btn.click(classify_text, inputs=txt, outputs=out)
+    gr.Markdown(bottom_text)
 
-demo.launch()
+iface.launch(share=True)