Update app_3.py

app_3.py

@@ -136,19 +136,26 @@ def extract_arguments(text, tokenizer, model, beam_search=True):
     def mark_text_by_position(original_text, start_token, end_token, color):
         """Marks text in the original string based on character positions."""
         # Inserts tags into the original text based on token offsets.
+
+        if start_token is not None and end_token is not None and start_token <= end_token:
         
+            start_idx, end_idx = offset_mapping[start_token][0], offset_mapping[end_token][1]
+             
+            if start_idx is not None and end_idx is not None and start_idx < end_idx:
+                return (
+                    original_text[:start_idx]
+                    + f"<mark style='background-color:{color}; padding:2px; border-radius:4px;'>"
+                    + original_text[start_idx:end_idx]
+                    + "</mark>"
+                    + original_text[end_idx:]
+                )
+    
+        if start_token > end_token:
+            st.write("The prediction is not correct: The position of the predicted end token comes before the position of the start token.")
+            
         
-        start_idx, end_idx = offset_mapping[start_token][0], offset_mapping[end_token][1]
-         
-        if start_idx is not None and end_idx is not None and start_idx <= end_idx:
-            return (
-                original_text[:start_idx]
-                + f"<mark style='background-color:{color}; padding:2px; border-radius:4px;'>"
-                + original_text[start_idx:end_idx]
-                + "</mark>"
-                + original_text[end_idx:]
-            )
         return original_text 
+        
 
     cause_text1 = mark_text_by_position(input_text, start_cause1, end_cause1, "#FFD700")  # Gold for cause
     effect_text1 = mark_text_by_position(input_text, start_effect1, end_effect1, "#90EE90")  # Light green for effect
@@ -163,7 +170,7 @@ def extract_arguments(text, tokenizer, model, beam_search=True):
 
 st.title("Causal Relation Extraction")
 input_text = st.text_area("Enter your text here:", height=300)
-beam_search = st.radio("Enable Beam Search?", ('No', 'Yes')) == 'Yes'
+beam_search = st.radio("Enable Position Selector & Beam Search?", ('No', 'Yes')) == 'Yes'
 
 
 if st.button("Extract_Now_1"):