Update app.py

app.py

@@ -1,74 +1,114 @@
 import gradio as gr
 import pandas as pd
-from detoxify import Detoxify
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
 import torch
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+from typing import List, Tuple
 
-# Load models once
-tox_model = Detoxify('multilingual')
-ai_tokenizer = AutoTokenizer.from_pretrained("openai-community/roberta-base-openai-detector")
-ai_model = AutoModelForSequenceClassification.from_pretrained("openai-community/roberta-base-openai-detector")
-
-# Thresholds
-TOXICITY_THRESHOLD = 0.7
+# =========================
+# Configuration
+# =========================
+MODEL_NAME = "openai-community/roberta-base-openai-detector"
 AI_THRESHOLD = 0.5
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+# =========================
+# Model Loading (once)
+# =========================
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME)
+model.to(DEVICE)
+model.eval()
+
+# =========================
+# Core Logic
+# =========================
+@torch.no_grad()
+def detect_ai_probability(texts: List[str]) -> List[float]:
+    """
+    Returns probability that each text is AI-generated.
+    """
+    inputs = tokenizer(
+        texts,
+        return_tensors="pt",
+        padding=True,
+        truncation=True,
+        max_length=512
+    ).to(DEVICE)
+
+    logits = model(**inputs).logits
+    probs = torch.softmax(logits, dim=1)[:, 1]  # AI-generated class
+    return probs.cpu().tolist()
+
+
+def classify_texts(texts: List[str]) -> pd.DataFrame:
+    """
+    Classify texts as AI or Human.
+    """
+    probabilities = detect_ai_probability(texts)
+
+    df = pd.DataFrame({
+        "Comment": texts,
+        "AI Probability": [round(p, 4) for p in probabilities],
+        "Prediction": [
+            "🤖 Likely AI" if p >= AI_THRESHOLD else "🧍 Human"
+            for p in probabilities
+        ]
+    })
 
-def detect_ai(text):
-    with torch.no_grad():
-        inputs = ai_tokenizer(text, return_tensors="pt", truncation=True, padding=True)
-        logits = ai_model(**inputs).logits
-        probs = torch.softmax(logits, dim=1).squeeze().tolist()
-    return round(probs[1], 4)  # Probability of AI-generated
-
-def classify_comments(comment_list):
-    results = tox_model.predict(comment_list)
-    df = pd.DataFrame(results, index=comment_list).round(4)
-    df.columns = [col.replace("_", " ").title().replace(" ", "_") for col in df.columns]
-    df.columns = [col.replace("_", " ") for col in df.columns]
-    df["⚠️ Warning"] = df.apply(
-        lambda row: "⚠️ High Risk" if any(score > TOXICITY_THRESHOLD for score in row) else "✅ Safe",
-        axis=1
-    )
-    df["🧪 AI Probability"] = [detect_ai(c) for c in df.index]
-    df["🧪 AI Detection"] = df["🧪 AI Probability"].apply(
-        lambda x: "🤖 Likely AI" if x > AI_THRESHOLD else "🧍 Human"
-    )
     return df
 
-def run_classification(text_input, csv_file):
-    comment_list = []
+
+def run_detector(text_input: str, csv_file) -> Tuple[pd.DataFrame, Tuple[str, bytes]]:
+    """
+    Handles UI input and output.
+    """
+    texts: List[str] = []
 
     if text_input.strip():
-        comment_list += [c.strip() for c in text_input.strip().split('\n') if c.strip()]
+        texts.extend([t.strip() for t in text_input.split("\n") if t.strip()])
 
     if csv_file:
         df = pd.read_csv(csv_file.name)
-        if 'comment' not in df.columns:
-            return "CSV must contain a 'comment' column.", None
-        comment_list += df['comment'].astype(str).tolist()
+        if "comment" not in df.columns:
+            return pd.DataFrame({"Error": ["CSV must contain a 'comment' column"]}), None
+        texts.extend(df["comment"].astype(str).tolist())
 
-    if not comment_list:
-        return "Please provide comments via text or CSV.", None
+    if not texts:
+        return pd.DataFrame({"Error": ["No input provided"]}), None
 
-    df = classify_comments(comment_list)
-    csv_data = df.copy()
-    csv_data.insert(0, "Comment", df.index)
-    return df, ("toxicity_predictions.csv", csv_data.to_csv(index=False).encode())
+    result_df = classify_texts(texts)
 
-# Build the Gradio UI
-with gr.Blocks(title="🌍 Toxic Comment & AI Detector") as app:
-    gr.Markdown("## 🌍 Toxic Comment & AI Detector")
-    gr.Markdown("Detects multilingual toxicity and whether a comment is AI-generated. Paste comments or upload a CSV.")
+    csv_bytes = result_df.to_csv(index=False).encode("utf-8")
+    return result_df, ("ai_detection_results.csv", csv_bytes)
+
+
+# =========================
+# Gradio UI
+# =========================
+with gr.Blocks(title="🧪 AI Text Detector") as app:
+    gr.Markdown("## 🧪 AI Text Detector")
+    gr.Markdown("Detect whether text is **AI-generated or human-written**.")
 
     with gr.Row():
-        text_input = gr.Textbox(lines=8, label="💬 Paste Comments (one per line)")
-        csv_input = gr.File(label="📥 Upload CSV (must have 'comment' column)")
+        text_input = gr.Textbox(
+            lines=8,
+            label="✍️ Paste Text (one per line)",
+            placeholder="Enter multiple comments, one per line..."
+        )
+        csv_input = gr.File(
+            label="📄 Upload CSV",
+            file_types=[".csv"]
+        )
 
-    submit_button = gr.Button("🔍 Analyze Comments")
-    output_table = gr.Dataframe(label="📊 Prediction Results")
-    download_button = gr.File(label="📤 Download CSV")
+    analyze_btn = gr.Button("🔍 Analyze")
+    output_table = gr.Dataframe(label="📊 Results")
+    download_file = gr.File(label="⬇️ Download CSV")
 
-    submit_button.click(fn=run_classification, inputs=[text_input, csv_input], outputs=[output_table, download_button])
+    analyze_btn.click(
+        fn=run_detector,
+        inputs=[text_input, csv_input],
+        outputs=[output_table, download_file]
+    )
 
 if __name__ == "__main__":
-    app.launch()
+    app.launch()