Update predictor.py

predictor.py

@@ -1,5 +1,6 @@
 import streamlit as st
 import pandas as pd
+import re 
 from utilities import predict
 
 st.set_page_config(layout="wide")
@@ -32,6 +33,14 @@ def result_form(result, user_label):
         predicted_label = max(probabilities, key=probabilities.get)
         return probabilities[user_label] < 0.35 or predicted_label != user_label
 
+def split_sentences(context):
+    # Regular expression to split text into sentences
+    pattern = r'(?<!\.\.\.)(?<=[.!?]) +|\.\.\.(?= )'
+    sentences = re.split(pattern, context)
+    # Remove leading/trailing whitespaces
+    sentences = [sentence.strip() for sentence in sentences if sentence.strip()]
+    return sentences
+
 def create_expander_with_check_button(label, title, context, predict_func):
     claim_key = f"{label}_input"
     evidence_key = f"{label}_evidence_selected"
@@ -50,10 +59,8 @@ def create_expander_with_check_button(label, title, context, predict_func):
                 result = predict_func(context, claim)
                 if result_form(result, label):
                     if label != 'NEI':  # NEI does not require evidence
-                        # Split context into sentences, handling ellipses ("...") properly
-                        sentences = re.split(r'(?<=[.!?]) +', context)
-                        # Remove leading/trailing whitespaces
-                        sentences = [sentence.strip() for sentence in sentences if sentence.strip()]
+                        # Split context into sentences with special handling for ellipses
+                        sentences = split_sentences(context)
                         # Display available sentences as options
                         st.multiselect(f"Selected evidence for {label}", sentences, default=st.session_state[label_e_ops], key=evidence_key)
                 else: