Update app_8.py

app_8.py

@@ -12,7 +12,6 @@ hf_token = st.secrets["HUGGINGFACE_TOKEN"]
 login(token=hf_token)
 
 
-# Load model & tokenizer once (cached for efficiency)
 @st.cache_resource
 def load_model():
     
@@ -35,20 +34,19 @@ def load_model():
     repo_id = "anamargarida/SpanExtractionWithSignalCls_2"  
     filename = "model.safetensors"  
 
-    # Download the model file
+    
     model_path = hf_hub_download(repo_id=repo_id, filename=filename)
 
-    # Load the model weights
     state_dict = load_file(model_path)
 
     model.load_state_dict(state_dict)
     
     return tokenizer, model
 
-# Load the model and tokenizer
+
 tokenizer, model = load_model()
 
-model.eval()  # Set model to evaluation mode
+model.eval()  
 def extract_arguments(text, tokenizer, model, beam_search=True):
      
     class Args:
@@ -59,14 +57,14 @@ def extract_arguments(text, tokenizer, model, beam_search=True):
     args = Args()
     inputs = tokenizer(text, return_offsets_mapping=True, return_tensors="pt")
     
-    # Get tokenized words (for reconstruction later)
+    
     word_ids = inputs.word_ids()
     
     with torch.no_grad():
         outputs = model(**inputs)
 
 
-    # Extract logits
+    
     start_cause_logits = outputs["start_arg0_logits"][0]
     end_cause_logits = outputs["end_arg0_logits"][0]
     start_effect_logits = outputs["start_arg1_logits"][0]
@@ -164,11 +162,11 @@ def extract_arguments(text, tokenizer, model, beam_search=True):
     
         
         
-    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
+    cause_text1 = mark_text_by_position(input_text, start_cause1, end_cause1, "#FFD700")  # yellow for cause
+    effect_text1 = mark_text_by_position(input_text, start_effect1, end_effect1, "#90EE90")  # green for effect
     
     if start_signal is not None and end_signal is not None:
-        signal_text = mark_text_by_position(input_text, start_signal, end_signal, "#FF6347")  # Tomato red for signal
+        signal_text = mark_text_by_position(input_text, start_signal, end_signal, "#FF6347")  # red for signal
     else: 
         signal_text = None
         
@@ -225,7 +223,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)
+input_text = st.text_area("Enter your text here:", height=100)
 beam_search = st.radio("Enable Position Selector & Beam Search?", ('Yes', 'No')) == 'Yes'