Update app.py

app.py

@@ -3,26 +3,115 @@ import torch
 import torch.nn.functional as F
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 import re
-import logging  # Optional: Add logging for better debugging
-
-# Set up logging (optional but helpful)
+import math
+import logging
+
+# --- Heuristic weights & config (from your PHP) ---
+SIGMOID_K = 0.5
+
+AI_WEIGHTS = {
+    'users': 1.0000, 'significant': 0.8061, 'despite': 0.5934, 'additionally': 0.5313,
+    'potential': 0.4900, 'features': 0.4824, 'various': 0.4785, 'regarding': 0.4689,
+    'remains': 0.4403, 'featuring': 0.4031, 'experience': 0.4026, 'including': 0.3963,
+    'challenges': 0.3548, 'allowing': 0.3528, 'enhance': 0.3437, 'aims': 0.3238,
+    'leading': 0.3064, 'user': 0.3054, 'recent': 0.2717, 'concerns': 0.2707,
+    'capabilities': 0.2684, 'technology': 0.2636, 'devices': 0.2616, 'following': 0.2551,
+    'anticipated': 0.2484, 'unique': 0.2418, 'expressed': 0.2403, 'innovative': 0.2383,
+    'design': 0.2379, 'remain': 0.2371, 'previous': 0.2331, 'priced': 0.2325,
+    'launch': 0.2324, 'enhancing': 0.2319, 'showcasing': 0.2305, 'feature': 0.2288,
+    'particularly': 0.2287, 'set': 0.2228, 'aimed': 0.2216, 'highlighted': 0.2192,
+    'ongoing': 0.2188, 'access': 0.2182, 'available': 0.2159, 'alongside': 0.2144,
+    'introduced': 0.2133, 'previously': 0.2122, 'highlighting': 0.2113, 'models': 0.2081,
+    'faced': 0.2057, 'platforms': 0.2055, 'updates': 0.2037, 'offers': 0.2032,
+    'significantly': 0.2027, 'issues': 0.2021, 'emphasized': 0.1977, 'initially': 0.1955,
+    'content': 0.1926, 'emphasizing': 0.1924, 'options': 0.1874, 'performance': 0.1864,
+    'initial': 0.1832, 'notable': 0.1821, 'additional': 0.1812, 'individuals': 0.1804,
+    'initiative': 0.1802, 'enhanced': 0.1797, 'release': 0.1797, 'currently': 0.1790,
+    'traditional': 0.1769, 'future': 0.1731, 'expected': 0.1725, 'applications': 0.1707,
+    'indicating': 0.1699, 'notably': 0.1658, 'insights': 0.1656, 'noted': 0.1645,
+    'players': 0.1645, 'narrative': 0.1642, 'landscape': 0.1640, 'upcoming': 0.1634,
+    'providing': 0.1631, 'offering': 0.1615, 'enabling': 0.1610, 'gaming': 0.1595,
+    'compared': 0.1553, 'indicated': 0.1539, 'extensive': 0.1530, 'approach': 0.1521,
+    'allows': 0.1519, 'stated': 0.1519, 'development': 0.1515, 'commitment': 0.1495,
+    'highlights': 0.1493, 'essential': 0.1483, 'experiences': 0.1480, 'recently': 0.1471,
+    'suggesting': 0.1457, 'market': 0.1447, 'uncertain': 0.1440, 'potentially': 0.1433
+}
+
+OG_WEIGHTS = {
+    'says': 1.0000, 'just': 0.9623, 'people': 0.8774, 'said': 0.8259, 'company': 0.7645,
+    'll': 0.6372, 'make': 0.6237, 'time': 0.5634, 'way': 0.5374, 've': 0.5039,
+    'want': 0.4435, 'like': 0.4426, 'don': 0.4338, 'going': 0.4160, 'really': 0.4126,
+    'use': 0.3769, 'good': 0.3718, 'lot': 0.3710, 'able': 0.3611, 'things': 0.3595,
+    'big': 0.3483, 'doesn': 0.3470, 'right': 0.3453, 'work': 0.3443, 'new': 0.3381,
+    'know': 0.3355, 'think': 0.3218, 'today': 0.3209, 'isn': 0.3039, 'look': 0.3013,
+    'world': 0.2907, 'say': 0.2875, 'best': 0.2825, 'used': 0.2758, 'little': 0.2735,
+    'actually': 0.2724, 'phone': 0.2551, 'thing': 0.2477, 'year': 0.2461, 'come': 0.2328,
+    'told': 0.2315, 'far': 0.2250, 'better': 0.2245, 'didn': 0.2244, 'getting': 0.2209,
+    'help': 0.2193, 'makes': 0.2141, 'got': 0.2139, 'won': 0.2096, 'called': 0.2078,
+    'different': 0.2010, 'verge': 0.2001, 'game': 0.1956, 'looks': 0.1954, 'comes': 0.1953,
+    'years': 0.1935, 'working': 0.1924, 'kind': 0.1899, 'let': 0.1891, 'great': 0.1878,
+    'read': 0.1876, 'number': 0.1868, 'long': 0.1852, 'according': 0.1795, 'coming': 0.1784,
+    'day': 0.1750, 'pretty': 0.1734, 'looking': 0.1685, 'bit': 0.1682, 'place': 0.1677,
+    'start': 0.1667, 'trying': 0.1661, 'sure': 0.1655, 'means': 0.1642, 'course': 0.1641,
+    'week': 0.1637, 'story': 0.1610, 'buy': 0.1589, 'probably': 0.1581, 'play': 0.1561,
+    'using': 0.1554, 'doing': 0.1551, 'hard': 0.1525, 'did': 0.1509, 'money': 0.1497,
+    'point': 0.1472, 'idea': 0.1429, 'end': 0.1425, 'aren': 0.1396, 'fact': 0.1371,
+    'run': 0.1363, 'does': 0.1362, 'case': 0.1331, 'built': 0.1301, 'biggest': 0.1300,
+    'started': 0.1286, 'exactly': 0.1279, 'screen': 0.1277, 'deal': 0.1264, 'apps': 0.1234
+}
+
+def tokenize(text):
+    return re.findall(r'\b[a-z]{2,}\b', text.lower())
+
+def classify_text_likelihood(text: str) -> float:
+    tokens = tokenize(text)
+    if not tokens:
+        return 0.5
+    ai_score = og_score = matched = 0
+    for t in tokens:
+        aw = AI_WEIGHTS.get(t, 0)
+        ow = OG_WEIGHTS.get(t, 0)
+        if aw or ow:
+            matched += 1
+            ai_score += aw
+            og_score += ow
+    if matched == 0:
+        return 0.5
+    net = ai_score - og_score
+    return 1 / (1 + math.exp(-SIGMOID_K * net))
+
+# Wrap words in thick colored underlines based on heuristic
+def highlight_heuristic_words(text: str) -> str:
+    parts = re.split(r'(\b[a-z]{2,}\b)', text)
+    out = []
+    for part in parts:
+        lower = part.lower()
+        if lower in AI_WEIGHTS:
+            out.append(
+                f"<span style='text-decoration: underline; "
+                f"text-decoration-color: darkred; text-decoration-thickness: 2px;'>"
+                f"{part}</span>"
+            )
+        elif lower in OG_WEIGHTS:
+            out.append(
+                f"<span style='text-decoration: underline; "
+                f"text-decoration-color: darkgreen; text-decoration-thickness: 2px;'>"
+                f"{part}</span>"
+            )
+        else:
+            out.append(part)
+    return ''.join(out)
+
+# --- Logging & Streamlit setup ---
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# Streamlit page config
 st.set_page_config(
     page_title="AI Article Detection by DEJAN",
     page_icon="🧠",
     layout="wide"
 )
 
-# Logo as provided
-st.logo(
-    image="https://dejan.ai/wp-content/uploads/2024/02/dejan-300x103.png",
-    link="https://dejan.ai/",
-)
-
-# Font styling
 st.markdown("""
 <link href="https://fonts.googleapis.com/css2?family=Roboto&display=swap" rel="stylesheet">
 <style>
@@ -32,22 +121,13 @@ st.markdown("""
 </style>
 """, unsafe_allow_html=True)
 
-@st.cache_resource  # Cache the model and tokenizer to avoid reloading on every interaction
+@st.cache_resource
 def load_model_and_tokenizer(model_name):
-    """Loads the model and tokenizer."""
-    logger.info(f"Loading 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
-    logger.info(f"Using device: {device} with dtype: {dtype}")
-
-    logger.info(f"Loading model: {model_name}")
+    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)
-    model.eval()
-    logger.info("Model loaded successfully.")
-
+    model.to(device).eval()
     return tokenizer, model, device
 
 MODEL_NAME = "dejanseo/ai-detection-small"
@@ -55,73 +135,60 @@ 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 or tokenizer: {e}", exc_info=True)
+    logger.error(f"Failed to load model: {e}", exc_info=True)
     st.stop()
 
-# Labels
-LABELS = ["AI Content", "Human Content"]
-
-# Regex-based sentence splitter
 def sent_tokenize(text):
-    sentences = re.split(r'(?<=[\.!?])\s+', text.strip())
-    return [s for s in sentences if s]
+    return [s for s in re.split(r'(?<=[\.!?])\s+', text.strip()) if s]
 
-# UI
 st.title("AI Article Detection")
-text = st.text_area("Enter text to classify", height=200, placeholder="Paste your text here...")
+
+text = st.text_area("Enter text to classify", height=200, placeholder="Paste your text here…")
 
 if st.button("Classify", type="primary"):
-    if not text or not text.strip():
+    if not text.strip():
         st.warning("Please enter some text.")
     else:
-        with st.spinner("Analyzing... Please wait."):
-            try:
-                sentences = sent_tokenize(text)
-                if not sentences:
-                    st.warning("No sentences detected.")
-                    st.stop()
-
-                # Tokenize sentences
-                inputs = tokenizer(
-                    sentences,
-                    return_tensors="pt",
-                    padding=True,
-                    truncation=True,
-                    max_length=model.config.max_position_embeddings
-                ).to(device)
-
-                # Inference
-                with torch.no_grad():
-                    outputs = model(**inputs)
-                    logits = outputs.logits
-                    probs = F.softmax(logits, dim=-1).cpu()  # [n_sentences, 2]
-                    preds = torch.argmax(probs, dim=-1).cpu()
-
-                # Build inline styled text
-                styled_chunks = []
-                for i, sent in enumerate(sentences):
-                    pred = preds[i].item()
-                    # red for AI (class 0), green for Human (class 1)
-                    r, g = (255, 0) if pred == 0 else (0, 255)
-                    confidence = probs[i, pred].item()  # 0.0–1.0
-                    alpha = confidence  # opacity
-                    span = (
-                        f"<span "
-                        f"style='background-color: rgba({r},{g},0,{alpha:.2f}); "
-                        f"padding:2px; margin:0 2px; border-radius:3px;'>"
-                        f"{sent}"
-                        f"</span>"
-                    )
-                    styled_chunks.append(span)
-
-                full_text_html = "".join(styled_chunks)
-                st.markdown(full_text_html, unsafe_allow_html=True)
-
-                # Overall AI likelihood (class 0)
-                avg_probs = torch.mean(probs, dim=0)
-                ai_likelihood = avg_probs[0].item() * 100
-                st.subheader(f"🤖 AI Likelihood: {ai_likelihood:.1f}%")
-
-            except Exception as e:
-                st.error(f"An error occurred during analysis: {e}")
-                logger.error("Analysis failed", exc_info=True)
+        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()
+
+            # Highlight each sentence and underline heuristic words
+            chunks = []
+            for i, s in enumerate(sentences):
+                inner = highlight_heuristic_words(s)
+                p = preds[i].item()
+                r, g = (255, 0) if p == 0 else (0, 255)
+                conf = probs[i, p].item()
+                alpha = conf
+                span = (
+                    f"<span style='background-color: rgba({r},{g},0,{alpha:.2f}); "
+                    f"padding:2px; margin:0 2px; border-radius:3px;'>{inner}</span>"
+                )
+                chunks.append(span)
+            st.markdown("".join(chunks), unsafe_allow_html=True)
+
+            # Scores
+            avg = torch.mean(probs, dim=0)
+            model_ai = avg[0].item()
+            heuristic_ai = classify_text_likelihood(text)
+            combined = min(model_ai + heuristic_ai, 1.0)
+
+            st.subheader(f"🤖 Model AI Likelihood: {model_ai*100:.1f}%")
+            st.subheader(f"🛠️ Heuristic AI Likelihood: {heuristic_ai*100:.1f}%")
+            st.subheader(f"⚖️ Combined AI Likelihood: {combined*100:.1f}%")