Update app_3.py

app_3.py

@@ -131,44 +131,41 @@ def extract_arguments(text, tokenizer, model, beam_search=True):
     
     tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])
     token_ids = inputs["input_ids"][0]
-    #offset_mapping = inputs["offset_mapping"][0].tolist()
-    
-
-
-    return start_cause1, end_cause1, start_cause2, end_cause2, start_effect1, end_effect1, start_effect2, end_effect2, start_signal, end_signal
-
-
-   
-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.
-    
     offset_mapping = inputs["offset_mapping"][0].tolist()
+   
+    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.
+        
+        
+        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
+    signal_text = mark_text_by_position(input_text, start_signal, end_signal, "#FF6347")  # Tomato red for signal
+    cause_text2 = mark_text_by_position(input_text, start_cause2, end_cause2, "#FFD700")
+    effect_text2 = mark_text_by_position(input_text, start_effect2, end_effect2, "#90EE90")
     
-    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  
+    return cause_text1, effect_text1, signal_text, cause_text2, effect_text2
 
 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'
 
 
-if st.button("Extract"):
+if st.button("Extract_Now"):
     if input_text:
-        start_cause1, end_cause1, start_cause2, end_cause2, start_effect1, end_effect1, start_effect2, end_effect2, start_signal, end_signal = extract_arguments(input_text, tokenizer, model, beam_search=beam_search)
-
-        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
-        signal_text = mark_text_by_position(input_text, start_signal, end_signal, "#FF6347")  # Tomato red for signal
+        cause_text1, effect_text1, signal_text, cause_text2, effect_text2 = extract_arguments(input_text, tokenizer, model, beam_search=beam_search)
 
         # Display first relation
         st.markdown(f"<strong>Relation 1:</strong>", unsafe_allow_html=True)
@@ -178,8 +175,6 @@ if st.button("Extract"):
 
         # Display second relation if beam search is enabled
         if beam_search:
-            cause_text2 = mark_text_by_position(input_text, start_cause2, end_cause2, "#FFD700")
-            effect_text2 = mark_text_by_position(input_text, start_effect2, end_effect2, "#90EE90")
 
             st.markdown(f"<strong>Relation 2:</strong>", unsafe_allow_html=True)
             st.markdown(f"**Cause:** {cause_text2}", unsafe_allow_html=True)