Optimize sentence-level analysis

app.py

@@ -39,7 +39,15 @@ def classify_text(text):
 
 def highlight_suspicious_sentences(text):
     sentences = sent_tokenize(text)
-    scores = [classify_text(sentence) for sentence in sentences]
+    inputs = tokenizer(
+        sentences, return_tensors="pt", truncation=True, max_length=512, padding=True
+    )
+    with torch.no_grad():
+        logits = model(**inputs).logits
+    probabilities = torch.softmax(logits, dim=1)
+    scores = probabilities[
+        :, 1
+    ].tolist()  # Probability of being AI-generated for each sentence
     return sentences, scores
 
 
@@ -69,21 +77,22 @@ if text_input != st.session_state.text_input:
 
 if st.button("Analyze and Highlight"):
     if text_input:
-        overall_probability = classify_text(text_input)
-        st.markdown(
-            f"<h3>Overall probability of being AI-generated: <span style='color: {'red' if overall_probability > 0.5 else 'green'};'>{overall_probability:.2%}</span></h3>",
-            unsafe_allow_html=True,
-        )
-
-        st.markdown("### Sentence-level analysis:")
-        sentences, scores = highlight_suspicious_sentences(text_input)
-
-        for sentence, score in zip(sentences, scores):
-            color = get_color(score)
+        with st.spinner("Analyzing text..."):
+            overall_probability = classify_text(text_input)
             st.markdown(
-                f"<div style='background-color: {color}; padding: 10px; margin: 5px 0; border-radius: 5px;'><strong>{score:.2%}</strong> - {sentence}</div>",
+                f"<h3>Overall probability of being AI-generated: <span style='color: {'red' if overall_probability > 0.5 else 'green'};'>{overall_probability:.2%}</span></h3>",
                 unsafe_allow_html=True,
             )
+
+            st.markdown("### Sentence-level analysis:")
+            sentences, scores = highlight_suspicious_sentences(text_input)
+
+            for sentence, score in zip(sentences, scores):
+                color = get_color(score)
+                st.markdown(
+                    f"<div style='background-color: {color}; padding: 10px; margin: 5px 0; border-radius: 5px;'><strong>{score:.2%}</strong> - {sentence}</div>",
+                    unsafe_allow_html=True,
+                )
     else:
         st.warning("Please enter some text to analyze.")