Rename app_23.py to app_24.py

app_23.py → app_24.py

@@ -60,7 +60,8 @@ def extract_arguments(text, tokenizer, model, beam_search=True):
             self.pretrained_signal_detector = False
         
     args = Args()
-    inputs = tokenizer(text, return_tensors="pt")
+    inputs = tokenizer(text, return_offsets_mapping=True, return_tensors="pt")
+    
     # Get tokenized words (for reconstruction later)
     word_ids = inputs.word_ids()
     
@@ -133,9 +134,14 @@ 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()
     
     for i, (token, word_id) in enumerate(zip(tokens, word_ids)):
         st.write(f"Token {i}: {token}, Word ID: {word_id}")
+
+    st.write("Token & offset:")
+    for i, (token, offset) in enumerate(zip(tokens, offset_mapping)):
+        st.write(f"Token {i}: {token}, Offset: {offset}")
     
 
     st.write("Token Positions, IDs, and Corresponding Tokens:")
@@ -250,11 +256,50 @@ def extract_arguments(text, tokenizer, model, beam_search=True):
         # Join tokens back into a string
         return ' '.join(this_space_splitted_tokens)
 
+    def add_tags_offset(text, start_cause, end_cause, start_effect, end_effect, start_signal, end_signal):
+        """
+        Inserts tags into the original text based on token offsets.
+        
+        Args:
+            text (str): The original input text.
+            tokenizer: The tokenizer used for tokenization.
+            start_cause (int): Start token index of the cause span.
+            end_cause (int): End token index of the cause span.
+            start_effect (int): Start token index of the effect span.
+            end_effect (int): End token index of the effect span.
+            start_signal (int, optional): Start token index of the signal span.
+            end_signal (int, optional): End token index of the signal span.
+    
+        Returns:
+            str: The modified text with annotated spans.
+        """
+    
+        
+    
+        # Convert token-based indices to character-based indices
+        start_cause_char, end_cause_char = offset_mapping[start_cause][0], offset_mapping[end_cause][1]
+        start_effect_char, end_effect_char = offset_mapping[start_effect][0], offset_mapping[end_effect][1]
+    
+        # Insert tags into the original text
+        annotated_text = text[:start_cause_char] + "<ARG0>" + text[start_cause_char:end_cause_char] + "</ARG0>" + text[end_cause_char:start_effect_char] + "<ARG1>" + text[start_effect_char:end_effect_char] + "</ARG1>" + text[end_effect_char:]
+    
+        # If signal span exists, insert signal tags
+        if start_signal is not None and end_signal is not None:
+            start_signal_char, end_signal_char = offset_mapping[start_signal][0], offset_mapping[end_signal][1]
+            annotated_text = (
+                annotated_text[:start_signal_char]
+                + "<SIG0>" + annotated_text[start_signal_char:end_signal_char] + "</SIG0>"
+                + annotated_text[end_signal_char:]
+            )
+    
+        return annotated_text
+
+
 
 
     # Apply the tags to the sentence tokens
-    tagged_sentence1 = add_tags_find(input_text, word_ids, start_cause1, end_cause1, start_effect1, end_effect1, start_signal, end_signal)
-    tagged_sentence2 = add_tags_find(input_text, word_ids, start_cause2, end_cause2, start_effect2, end_effect2, start_signal, end_signal)
+    tagged_sentence1 = add_tags_offset(input_text, start_cause1, end_cause1, start_effect1, end_effect1, start_signal, end_signal)
+    tagged_sentence2 = add_tags_offset(input_text, start_cause2, end_cause2, start_effect2, end_effect2, start_signal, end_signal)
     return tagged_sentence1, tagged_sentence2