Update app.py

app.py

@@ -1,34 +1,59 @@
 import torch
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
-import matplotlib.pyplot as plt
 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
 
-# --- Load models ---
+# ===============================
+# Safe model loading
+# ===============================
 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)
+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
+
+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 (example) ---
 label_mapping = {i: f"model_{i}" for i in range(25)}
 label_mapping[24] = "Human"
 
+# ===============================
+# Helper functions
+# ===============================
 def clean_text(text):
-    text = re.sub(r'\s+', ' ', text).strip()
-    return 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
+# ===============================
 def classify_text(text):
     cleaned_text = clean_text(text)
-    if not cleaned_text.strip():
-        return "", None
+    if not cleaned_text:
+        return "<b style='color:red'>Please enter some text.</b>"
 
-    # 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)]
 
@@ -37,99 +62,92 @@ def classify_text(text):
 
     for i, para in enumerate(paragraphs):
         inputs = tokenizer(para, return_tensors="pt", truncation=True, padding=True).to(device)
+        softmax_outputs = []
 
-        with torch.no_grad():
-            logits_1 = model_1(**inputs).logits
-            logits_2 = model_2(**inputs).logits
-            logits_3 = model_3(**inputs).logits
+        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))
 
-            softmax_1 = torch.softmax(logits_1, dim=1)
-            softmax_2 = torch.softmax(logits_2, dim=1)
-            softmax_3 = torch.softmax(logits_3, dim=1)
+        if not softmax_outputs:
+            return "<b style='color:red'>Error: No models loaded successfully.</b>"
 
-            averaged_probabilities = (softmax_1 + softmax_2 + softmax_3) / 3
-            probabilities = averaged_probabilities[0]
-            all_probabilities.append(probabilities.cpu())
+        avg_probs = sum(softmax_outputs) / len(softmax_outputs)
+        probabilities = avg_probs[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()
+        ai_probs = probabilities.clone()
+        ai_probs[24] = 0
+        ai_prob = ai_probs.sum().item()
 
-        total = human_prob + ai_total_prob
+        total = human_prob + ai_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
+        ai_pct = (ai_prob / total) * 100
+        ai_model = label_mapping[torch.argmax(ai_probs).item()]
 
+        preview = para[:180].strip() + ("..." if len(para) > 180 else "")
         paragraph_scores.append({
             "id": i + 1,
             "human": human_pct,
             "ai": ai_pct,
             "model": ai_model,
-            "preview": short_preview
+            "preview": preview
         })
 
-    # --- 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"<b>Overall Result:</b> <span class='highlight-human'>{avg_human:.2f}% Human-written</span>"
+        overall = f"<b>Overall Result:</b> <span class='highlight-human'>{avg_human:.2f}% Human-written</span>"
     else:
-        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>"
+        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 5 chart ---
+    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()
+
+    fig, ax = plt.subplots(figsize=(8, 4))
+    ax.barh(labels, values, color='#4CAF50')
+    ax.set_xlabel("Probability")
+    ax.set_title("Top 5 Model Predictions")
+    ax.invert_yaxis()
+    chart_html = plot_to_base64(fig)
 
-    # --- 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>"
+    # --- 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_output += f"""
+        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>
         """
 
-    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()
-    top_5_labels = [label_mapping[i.item()] for i in top_5_indices]
-
-    fig, ax = plt.subplots(figsize=(10, 5))
-    bars = ax.barh(top_5_labels, top_5_probs, color='#4CAF50', alpha=0.8)
-    ax.set_xlabel('Probability', fontsize=12)
-    ax.set_title('Top 5 Predictions (Averaged)', fontsize=14, fontweight='bold')
-    ax.invert_yaxis()
-    ax.grid(axis='x', linestyle='--', alpha=0.6)
-    for bar in bars:
-        width = bar.get_width()
-        ax.text(width + 0.01, bar.get_y() + bar.get_height() / 2, f'{width:.2%}', va='center')
-    ax.set_xlim(0, max(top_5_probs) * 1.18)
-    plt.tight_layout()
+    html += "<br><h3>Top 5 Models:</h3>" + chart_html + "</div>"
+    return html
 
-    return html_output, fig
-
-# --- Gradio UI ---
+# ===============================
+# Gradio UI
+# ===============================
 css = """
 .highlight-ai { color: #FF5733; font-weight: bold; }
 .highlight-human { color: #28a745; font-weight: bold; }
 """
 
 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])
+    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)
 
 demo.launch()