Update app.py

app.py

@@ -1,125 +1,79 @@
-import streamlit as st
 import torch
 import torch.nn.functional as F
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 import re
-import logging
 import pandas as pd
-
-st.set_page_config(
-    page_title="AI Article Detection by Writenix",
-    page_icon="🧠",
-    layout="wide"
-)
-
-st.logo(
-    image="https://dejan.ai/wp-content/uploads/2024/02/dejan-300x103.png",
-    link="https://dejan.ai/",
-)
-
-# --- Logging & Streamlit setup ---
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-st.markdown("""
-<link href="https://fonts.googleapis.com/css2?family=Roboto&display=swap" rel="stylesheet">
-<style>
-    html, body, [class*="css"] {
-        font-family: 'Roboto', sans-serif;
-    }
-</style>
-""", unsafe_allow_html=True)
-
-@st.cache_resource
-def load_model_and_tokenizer(model_name):
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    dtype = torch.bfloat16 if (device.type=="cuda" and torch.cuda.is_bf16_supported()) else torch.float32
-    model = AutoModelForSequenceClassification.from_pretrained(model_name, torch_dtype=dtype)
-    model.to(device).eval()
-    return tokenizer, model, device
+import gradio as gr
 
 MODEL_NAME = "dejanseo/ai-cop"
-try:
-    tokenizer, model, device = load_model_and_tokenizer(MODEL_NAME)
-except Exception as e:
-    st.error(f"Error loading model: {e}")
-    logger.error(f"Failed to load model: {e}", exc_info=True)
-    st.stop()
+
+# --- Load model ---
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+dtype = torch.bfloat16 if (device.type=="cuda" and torch.cuda.is_bf16_supported()) else torch.float32
+model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME, torch_dtype=dtype)
+model.to(device).eval()
 
 def sent_tokenize(text):
     return [s for s in re.split(r'(?<=[\.!?])\s+', text.strip()) if s]
 
-st.title("AI Article Detection")
-
-text = st.text_area("Enter text to classify", height=200, placeholder="Paste your text here…")
-
-if st.button("Classify", type="primary"):
+# --- Inference function ---
+def classify_text(text):
     if not text.strip():
-        st.warning("Please enter some text.")
-    else:
-        with st.spinner("Analyzing…"):
-            sentences = sent_tokenize(text)
-            if not sentences:
-                st.warning("No sentences detected.")
-                st.stop()
-
-            inputs = tokenizer(
-                sentences,
-                return_tensors="pt",
-                padding=True,
-                truncation=True,
-                max_length=model.config.max_position_embeddings
-            ).to(device)
-
-            with torch.no_grad():
-                logits = model(**inputs).logits
-                probs = F.softmax(logits, dim=-1).cpu()
-                preds = torch.argmax(probs, dim=-1).cpu()
-
-            # Create dataframe for sentences
-            sentences_data = []
-            highlighted_sentences = []
-            for i, s in enumerate(sentences):
-                p = preds[i].item()
-                conf = probs[i, p].item()
-                label = "AI" if p == 0 else "Human"
-
-                sentences_data.append({
-                    "sentence": s,
-                    "classification": label,
-                    "confidence": conf
-                })
-
-                if label == "AI":
-                    highlighted_sentences.append(f"<span style='color:red; font-weight:bold'>{s}</span>")
-                else:
-                    highlighted_sentences.append(f"<span style='color:green; font-weight:bold'>{s}</span>")
-
-            # Display dataframe
-            df = pd.DataFrame(sentences_data)
-            st.dataframe(
-                df,
-                column_config={
-                    "sentence": st.column_config.TextColumn("Sentence"),
-                    "classification": st.column_config.TextColumn("Classification"),
-                    "confidence": st.column_config.ProgressColumn(
-                        "Confidence",
-                        help="Model's confidence in the classification",
-                        format="%.2f",
-                        min_value=0,
-                        max_value=1,
-                    ),
-                },
-                hide_index=True,
-            )
-
-            # Highlighted text output
-            st.markdown("### 🔍 Highlighted Text")
-            st.markdown(" ".join(highlighted_sentences), unsafe_allow_html=True)
-
-            # Overall score (just model avg)
-            avg = torch.mean(probs, dim=0)
-            model_ai = avg[0].item()
-
-            st.subheader(f"⚖️ AI Likelihood: {model_ai*100:.1f}%")
+        return "⚠️ Please enter some text.", None, None
+
+    sentences = sent_tokenize(text)
+    if not sentences:
+        return "⚠️ No sentences detected.", None, None
+
+    inputs = tokenizer(
+        sentences,
+        return_tensors="pt",
+        padding=True,
+        truncation=True,
+        max_length=model.config.max_position_embeddings
+    ).to(device)
+
+    with torch.no_grad():
+        logits = model(**inputs).logits
+        probs = F.softmax(logits, dim=-1).cpu()
+        preds = torch.argmax(probs, dim=-1).cpu()
+
+    results = []
+    highlighted_sentences = []
+    for i, s in enumerate(sentences):
+        p = preds[i].item()
+        conf = probs[i, p].item()
+        label = "AI" if p == 0 else "Human"
+
+        results.append([s, label, f"{conf:.2f}"])
+        if label == "AI":
+            highlighted_sentences.append(f"<span style='color:red; font-weight:bold'>{s}</span>")
+        else:
+            highlighted_sentences.append(f"<span style='color:green; font-weight:bold'>{s}</span>")
+
+    # Overall AI likelihood
+    avg = torch.mean(probs, dim=0)
+    model_ai = avg[0].item() * 100
+
+    highlighted_text = " ".join(highlighted_sentences)
+    df = pd.DataFrame(results, columns=["Sentence", "Classification", "Confidence"])
+    return f"⚖️ AI Likelihood: {model_ai:.1f}%", highlighted_text, df
+
+# --- Gradio Interface ---
+with gr.Blocks() as demo:
+    gr.Markdown("## 🧠 AI Article Detection by Writenix")
+
+    with gr.Row():
+        text_input = gr.Textbox(label="Enter text", lines=10, placeholder="Paste your text here…")
+
+    classify_btn = gr.Button("Classify")
+
+    ai_score = gr.Label(label="Overall AI Likelihood")
+    highlighted = gr.HTML()
+    table = gr.Dataframe(headers=["Sentence", "Classification", "Confidence"], wrap=True)
+
+    classify_btn.click(classify_text, inputs=text_input, outputs=[ai_score, highlighted, table])
+
+if __name__ == "__main__":
+    demo.launch()