Update app.py

app.py

@@ -56,33 +56,10 @@ def classify_text_likelihood(text: str) -> float:
     net = ai_score - og_score
     return 1 / (1 + math.exp(-SIGMOID_K * net))
 
-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__)
 
-
-
 st.markdown("""
 <link href="https://fonts.googleapis.com/css2?family=Roboto&display=swap" rel="stylesheet">
 <style>
@@ -139,19 +116,29 @@ if st.button("Classify", type="primary"):
                 probs = F.softmax(logits, dim=-1).cpu()
                 preds = torch.argmax(probs, dim=-1).cpu()
 
-            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)
+            # Display sentences with Streamlit components instead of HTML
+            for i, sentence in enumerate(sentences):
+                is_ai = preds[i].item() == 0
+                confidence = probs[i, preds[i].item()].item()
+                
+                # Use Streamlit components for display
+                container = st.container(border=True)
+                with container:
+                    if is_ai:
+                        st.markdown(f"🤖 **AI content** (confidence: {confidence:.2f})")
+                    else:
+                        st.markdown(f"👤 **Human content** (confidence: {confidence:.2f})")
+                    st.text(sentence)
+                    
+                    # Show word indicators
+                    words = tokenize(sentence)
+                    ai_indicators = [w for w in words if w in AI_WEIGHTS]
+                    human_indicators = [w for w in words if w in OG_WEIGHTS]
+                    
+                    if ai_indicators:
+                        st.text(f"AI indicators: {', '.join(ai_indicators)}")
+                    if human_indicators:
+                        st.text(f"Human indicators: {', '.join(human_indicators)}")
 
             avg = torch.mean(probs, dim=0)
             model_ai = avg[0].item()
@@ -161,3 +148,8 @@ if st.button("Classify", type="primary"):
             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}%")
+            
+            # Add progress bars for visual representation
+            st.progress(model_ai, text="Model AI Score")
+            st.progress(heuristic_ai, text="Heuristic Score") 
+            st.progress(combined, text="Combined Score")