Rename app_7.py to app_8.py

app_7.py → app_8.py

@@ -65,7 +65,6 @@ def extract_arguments(text, tokenizer, model, beam_search=True):
     with torch.no_grad():
         outputs = model(**inputs)
 
-    #st.write("Model output keys:", outputs.keys())
 
     # Extract logits
     start_cause_logits = outputs["start_arg0_logits"][0]
@@ -132,9 +131,8 @@ def extract_arguments(text, tokenizer, model, beam_search=True):
         end_signal_logits[start_signal + 5:] = -1e4
         end_signal = end_signal_logits.argmax().item()
 
-    if not has_signal:
-        start_signal = 'NA'
-        end_signal = 'NA'
+    If not has_signal:
+        start_signal, end_signal = None, None
         
 
     tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])
@@ -161,21 +159,20 @@ def extract_arguments(text, tokenizer, model, beam_search=True):
         signal = 'NA'
     list1 = [start_cause1, end_cause1, start_effect1, end_effect1, start_signal, end_signal]
     list2 = [start_cause2, end_cause2, start_effect2, end_effect2, start_signal, end_signal]
-    return cause1, cause2, effect1, effect2, signal, list1, list2
-
-
-def mark_text(original_text, span, color):
-    """Replace extracted span with a colored background marker, starting at the exact first token."""
-    if span:
-        # Get the starting index of the span in the original text
-        start_idx = original_text.find(span)
-        
-        # Ensure the span is found and replace it with <mark> tags directly at the first token's position
-        if start_idx != -1:
-            # Replace the span with the mark tag at the exact first token position
-            return re.sub(re.escape(span), f"<mark style='background-color:{color}; padding:2px; border-radius:4px;'>{span}</mark>", original_text, flags=re.IGNORECASE)
-        
-    return original_text  # Return unchanged text if no span is found
+    #return cause1, cause2, effect1, effect2, signal, list1, list2
+    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_idx, end_idx, color):
+    """Marks text in the original string based on character positions."""
+    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 unchanged if indices are invalidt  # Return unchanged text if no span is found
 
 st.title("Causal Relation Extraction")
 input_text = st.text_area("Enter your text here:", height=300)
@@ -184,11 +181,11 @@ beam_search = st.radio("Enable Beam Search?", ('No', 'Yes')) == 'Yes'
 
 if st.button("Extract1"):
     if input_text:
-        cause1, cause2, effect1, effect2, signal, list1, list2 = extract_arguments(input_text, tokenizer, model, beam_search=beam_search)
+        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(input_text, cause1, "#FFD700")  # Gold for cause
-        effect_text1 = mark_text(input_text, effect1, "#90EE90")  # Light green for effect
-        signal_text = mark_text(input_text, signal, "#FF6347")  # Tomato red for signal
+        cause_text1 = mark_text(input_text, start_cause1, end_cause1, "#FFD700")  # Gold for cause
+        effect_text1 = mark_text(input_text, start_effect1, end_effect1, "#90EE90")  # Light green for effect
+        signal_text = mark_text(input_text, start_signal, end_signal, "#FF6347")  # Tomato red for signal
 
         st.markdown(f"<span style='font-size: 24px;'><strong>Relation 1:</strong></span>", unsafe_allow_html=True)
         st.markdown(f"**Cause:**<br>{cause_text1}", unsafe_allow_html=True)
@@ -199,9 +196,9 @@ if st.button("Extract1"):
 
         if beam_search:
 
-            cause_text2 = mark_text(input_text, cause2, "#FFD700")  # Gold for cause
-            effect_text2 = mark_text(input_text, effect2, "#90EE90")  # Light green for effect
-            signal_text = mark_text(input_text, signal, "#FF6347")  # Tomato red for signal
+            cause_text2 = mark_text(input_text, start_cause2, end_cause2, "#FFD700")  # Gold for cause
+            effect_text2 = mark_text(input_text, start_effect2, end_effect2, "#90EE90")  # Light green for effect
+            signal_text = mark_text(input_text, start_signal, end_signal, "#FF6347")  # Tomato red for signal
     
             st.markdown(f"<span style='font-size: 24px;'><strong>Relation 2:</strong></span>", unsafe_allow_html=True)
             st.markdown(f"**Cause:**<br>{cause_text2}", unsafe_allow_html=True)