Update src/streamlit_app.py

src/streamlit_app.py

@@ -3,28 +3,43 @@ import streamlit as st
 import torch
 import time
 
-# Limit CPU threads
+# Limit CPU threads for better responsiveness
 torch.set_num_threads(2)
 
 st.title("AI Humanizer Lite (CPU Friendly)")
 
-if "model" not in st.session_state:
+# Load pipelines once per session
+if "detect_pipe" not in st.session_state or "humanizer_pipe" not in st.session_state:
     with st.spinner("Loading models (takes ~10s on CPU)..."):
-        st.session_state.detect_pipe = pipeline("text-classification", model="roberta-base-openai-detector", device=-1)
-        st.session_state.humanizer_pipe = pipeline("text2text-generation", model="sshleifer/distilbart-cnn-12-6", device=-1)
-
-def detect_ai(text):
-    outputs = st.session_state.detect_pipe(text[:256])
-    # Assume label "AI" means AI-generated with a score
-    for out in outputs:
-        if out["label"].lower() == "ai":
-            return out["score"]
+        st.session_state.detect_pipe = pipeline(
+            "text-classification",
+            model="roberta-base-openai-detector",
+            device=-1,
+        )
+        st.session_state.humanizer_pipe = pipeline(
+            "text2text-generation",
+            model="sshleifer/distilbart-cnn-12-6",
+            device=-1,
+        )
+
+def detect_ai(text: str) -> float:
+    try:
+        outputs = st.session_state.detect_pipe(text[:256])
+        for out in outputs:
+            if out["label"].lower() == "ai":
+                return out["score"]
+    except Exception as e:
+        st.error(f"Error during AI detection: {e}")
     return 0.0
 
-def humanize_text(text):
+def humanize_text(text: str) -> str:
     prompt = f"Rewrite this text naturally: {text}"
-    result = st.session_state.humanizer_pipe(prompt, max_length=128, num_beams=3)
-    return result[0]["generated_text"]
+    try:
+        result = st.session_state.humanizer_pipe(prompt, max_length=128, num_beams=3)
+        return result[0]["generated_text"]
+    except Exception as e:
+        st.error(f"Error during text rewriting: {e}")
+        return text
 
 text = st.text_area("Paste AI-generated text here", height=200)