Update app.py

app.py

@@ -1,81 +1,42 @@
-import gradio as gr
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
 import torch
-import re
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
 import matplotlib.pyplot as plt
-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")
-
-# Load Model 1 (local)
-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()
-
-# Load Model 2 (URL)
-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()
-
-# Load Model 3 (URL)
-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',
-    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
+import re
+import gradio as gr
+
+# --- Load models ---
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+model_1 = AutoModelForSequenceClassification.from_pretrained("roberta-base-openai-detector").to(device)
+model_2 = AutoModelForSequenceClassification.from_pretrained("roberta-large-openai-detector").to(device)
+model_3 = AutoModelForSequenceClassification.from_pretrained("Hello-SimpleAI/chatgpt-detector-roberta").to(device)
+tokenizer = AutoTokenizer.from_pretrained("roberta-base-openai-detector")
 
-newline_to_space = Replace(Regex(r'\s*\n\s*'), " ")
-tokenizer.backend_tokenizer.normalizer = Sequence([
-    tokenizer.backend_tokenizer.normalizer,
-    newline_to_space,
-    Strip()
-])
+# --- Label Mapping (example) ---
+label_mapping = {i: f"model_{i}" for i in range(25)}
+label_mapping[24] = "Human"
+
+def clean_text(text):
+    text = re.sub(r'\s+', ' ', text).strip()
+    return text
 
-# --- Classification Function (Per Paragraph) ---
 def classify_text(text):
-    """
-    Classifies each paragraph separately and provides per-paragraph scores
-    + overall result.
-    """
     cleaned_text = clean_text(text)
     if not cleaned_text.strip():
         return "", None
 
-    # Split text into paragraphs (2+ newlines)
-    paragraphs = [p.strip() for p in re.split(r'\n{2,}', cleaned_text) if p.strip()]
-    chunk_scores = []
+    # Split into paragraphs (two newlines)
+    paragraphs = re.split(r'\n{2,}', cleaned_text)
+    if len(paragraphs) == 1 and len(cleaned_text.split()) > 300:
+        # Fallback: split by ~300 words
+        words = cleaned_text.split()
+        paragraphs = [' '.join(words[i:i + 300]) for i in range(0, len(words), 300)]
+
+    paragraph_scores = []
     all_probabilities = []
 
-    for i, paragraph in enumerate(paragraphs):
-        inputs = tokenizer(paragraph, return_tensors="pt", truncation=True, padding=True).to(device)
+    for i, para in enumerate(paragraphs):
+        inputs = tokenizer(para, return_tensors="pt", truncation=True, padding=True).to(device)
 
         with torch.no_grad():
             logits_1 = model_1(**inputs).logits
@@ -95,41 +56,48 @@ def classify_text(text):
         ai_probs_clone[24] = 0
         ai_total_prob = ai_probs_clone.sum().item()
 
-        total_decision_prob = human_prob + ai_total_prob
-        human_percentage = (human_prob / total_decision_prob) * 100
-        ai_percentage = (ai_total_prob / total_decision_prob) * 100
-        ai_argmax_index = torch.argmax(ai_probs_clone).item()
-        ai_argmax_model = label_mapping[ai_argmax_index]
-
-        chunk_scores.append({
-            "paragraph": paragraph[:150] + ("..." if len(paragraph) > 150 else ""),
-            "human": human_percentage,
-            "ai": ai_percentage,
-            "model": ai_argmax_model
+        total = human_prob + ai_total_prob
+        human_pct = (human_prob / total) * 100
+        ai_pct = (ai_total_prob / total) * 100
+        ai_index = torch.argmax(ai_probs_clone).item()
+        ai_model = label_mapping[ai_index]
+
+        short_preview = (para[:180] + "...") if len(para) > 180 else para
+
+        paragraph_scores.append({
+            "id": i + 1,
+            "human": human_pct,
+            "ai": ai_pct,
+            "model": ai_model,
+            "preview": short_preview
         })
 
-    # --- Overall Average ---
-    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)
+    # --- Averages ---
+    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)
 
     if avg_human > avg_ai:
-        result_message = f"**Overall Result:** <span class='highlight-human'>{avg_human:.2f}% Human-written</span>"
+        result_message = f"<b>Overall Result:</b> <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 Table ---
-    paragraph_table = "\n\n**Paragraph Analysis:**\n"
-    for idx, c in enumerate(chunk_scores, 1):
-        color = "#4CAF50" if c['human'] > c['ai'] else "#FF5733"
-        paragraph_table += (
-            f"<div style='margin-bottom:10px; border-left:4px solid {color}; padding-left:10px;'>"
-            f"<b>Paragraph {idx}</b>: {c['human']:.2f}% Human | {c['ai']:.2f}% AI → <i>{c['model']}</i><br>"
-            f"<small>{c['paragraph']}</small>"
-            f"</div>\n"
-        )
-
-    # --- Plot (Top 5 Models Overall) ---
+        top_model = max(paragraph_scores, key=lambda p: p['ai'])['model']
+        result_message = f"<b>Overall Result:</b> <span class='highlight-ai'>{avg_ai:.2f}% AI-generated (likely {top_model})</span>"
+
+    # --- Paragraph Analysis HTML ---
+    html_output = f"<div style='font-family: Arial, sans-serif; line-height:1.6;'>{result_message}<br><br>"
+    html_output += "<h3>Paragraph Analysis:</h3>"
+
+    for p in paragraph_scores:
+        color = "#28a745" if p["human"] > p["ai"] else "#FF5733"
+        html_output += 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>
+        """
+
+    html_output += "</div>"
+
+    # --- Top 5 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_probs.cpu().numpy()
@@ -147,46 +115,21 @@ def classify_text(text):
     ax.set_xlim(0, max(top_5_probs) * 1.18)
     plt.tight_layout()
 
-    return result_message + "\n\n" + paragraph_table, fig
-
+    return html_output, fig
 
-# --- UI ---
-title = "AI Text Detector"
-description = """
-This tool uses <b>ModernBERT</b> to detect AI-generated text.  
-Each paragraph is analyzed separately to show which parts are likely AI-generated.
+# --- Gradio UI ---
+css = """
+.highlight-ai { color: #FF5733; font-weight: bold; }
+.highlight-human { color: #28a745; font-weight: bold; }
 """
-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. 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. 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');
-    #text_input_box { border-radius: 10px; border: 2px solid #4CAF50; font-size: 18px; padding: 15px; margin-bottom: 20px; width: 60%; box-sizing: border-box; margin: auto; }
-    #result_output_box { border-radius: 10px; border: 2px solid #4CAF50; font-size: 16px; padding: 15px; margin-top: 20px; width: 80%; box-sizing: border-box; margin: auto; }
-    body { font-family: 'Roboto Mono', sans-serif !important; padding: 20px; display: block; justify-content: center; align-items: center; overflow-y: auto; }
-    .gradio-container { border: 1px solid #4CAF50; border-radius: 15px; padding: 30px; box-shadow: 0px 0px 10px rgba(0,255,0,0.6); max-width: 900px; margin: auto; }
-    .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...", elem_id="text_input_box", lines=10)
-    result_output = gr.HTML("", elem_id="result_output_box")
-    plot_output = gr.Plot(label="Model Probability Distribution")
-    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)
+
+with gr.Blocks(css=css, theme="soft") as demo:
+    gr.Markdown("# 🧠 AI/Human Text Detector")
+    text_input = gr.Textbox(label="Paste your text here", lines=12, placeholder="Paste your article or essay...")
+    output_html = gr.HTML(label="Analysis Results")
+    output_plot = gr.Plot(label="Top 5 Models")
+    analyze_btn = gr.Button("🔍 Analyze Text", variant="primary")
+
+    analyze_btn.click(classify_text, inputs=text_input, outputs=[output_html, output_plot])
+
+demo.launch()