Update app.py

app.py

@@ -1,90 +1,66 @@
-import gradio as gr
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
-import torch
 import re
-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 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()
+import torch
+from fastapi import FastAPI, Request
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+from pydantic import BaseModel
+from typing import List
+import uvicorn
 
-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()
+# ========== CONFIG ==========
+MODEL_PATH = "roberta-base-openai-detector"  # or your preferred detector
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
+tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
+model_1 = AutoModelForSequenceClassification.from_pretrained(MODEL_PATH).to(device)
+model_2 = AutoModelForSequenceClassification.from_pretrained(MODEL_PATH).to(device)
+model_3 = AutoModelForSequenceClassification.from_pretrained(MODEL_PATH).to(device)
 
-# ---- 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'
+    0: "gpt2", 1: "gpt3", 2: "gpt4", 3: "chatgpt", 4: "dolly", 5: "human", 24: "human"
 }
 
-# ---- Text Cleaning ----
+app = FastAPI(title="AI Text Classifier API", version="1.0.0")
+
+
+# ========== HELPERS ==========
 def clean_text(text: str) -> str:
-    text = re.sub(r'\s{2,}', ' ', text)
-    text = re.sub(r'\s+([,.;:?!])', r'\1', text)
-    return text
+    text = re.sub(r'\s+', ' ', text)
+    return text.strip()
+
 
-newline_to_space = Replace(Regex(r'\s*\n\s*'), " ")
-tokenizer.backend_tokenizer.normalizer = Sequence([
-    tokenizer.backend_tokenizer.normalizer,
-    newline_to_space,
-    Strip()
-])
+# ========== INPUT MODEL ==========
+class TextInput(BaseModel):
+    text: str
 
 
-# ---- Classification Function ----
-def classify_text(text):
-    cleaned_text = clean_text(text)
+# ========== MAIN LOGIC ==========
+@app.post("/analyze")
+async def analyze_text(data: TextInput):
+    cleaned_text = clean_text(data.text)
     if not cleaned_text.strip():
-        return "**Error:** Please enter some text to analyze.", None
+        return {"success": False, "error": "Empty text provided"}
 
-    # Split into paragraphs
     paragraphs = [p.strip() for p in re.split(r'\n{2,}', cleaned_text) if p.strip()]
+    if not paragraphs:
+        paragraphs = [cleaned_text]
+
     chunk_scores = []
-    all_probabilities = []
+    all_probs = []
 
-    for i, paragraph in enumerate(paragraphs):
+    for paragraph in 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
             logits_3 = model_3(**inputs).logits
-
             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 = (softmax_1 + softmax_2 + softmax_3) / 3
+            probs = averaged[0]
+            all_probs.append(probs.cpu())
 
-            avg_probs = (softmax_1 + softmax_2 + softmax_3) / 3
-            probs = avg_probs[0]
-            all_probabilities.append(probs.cpu())
-
-        human_prob = probs[24].item()
+        human_prob = probs[24].item() if 24 in label_mapping else probs[-1].item()
         ai_probs_clone = probs.clone()
         ai_probs_clone[24] = 0
         ai_total = ai_probs_clone.sum().item()
@@ -94,86 +70,30 @@ def classify_text(text):
         ai_model = label_mapping[torch.argmax(ai_probs_clone).item()]
 
         chunk_scores.append({
-            "human": human_pct,
-            "ai": ai_pct,
+            "human": round(human_pct, 2),
+            "ai": round(ai_pct, 2),
             "model": ai_model,
-            "text": paragraph[:200].replace('\n', ' ') + ("..." if len(paragraph) > 200 else "")
+            "text_preview": paragraph[:250].replace('\n', ' ') + ("..." if len(paragraph) > 250 else "")
         })
 
-    # ---- Overall Averages ----
+    # ---- 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:
-        result_message = f"**Overall Result:** <span class='highlight-human'>{avg_human:.2f}% Human-written</span>"
+    if avg_ai > avg_human:
+        top_model = max(chunk_scores, key=lambda c: c["ai"])["model"]
+        overall = {"result": f"{avg_ai:.2f}% AI-generated", "model": top_model}
     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 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_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_idx]
-
-    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()
-    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 + "\n\n" + paragraph_text, fig
-
-
-# ---- UI Setup ----
-title = "AI Text Detector"
-description = """
-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. 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%; 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; }
-""")
-
-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)
+        overall = {"result": f"{avg_human:.2f}% Human-written", "model": "human"}
+
+    return {
+        "success": True,
+        "overall": overall,
+        "paragraphs": chunk_scores,
+        "total_paragraphs": len(chunk_scores)
+    }
+
+
+# ========== RUN LOCALLY ==========
+if __name__ == "__main__":
+    uvicorn.run("app:app", host="0.0.0.0", port=8000)