Rename app_7.py to app_8.py

app_7.py → app_8.py

@@ -88,7 +88,7 @@ def extract_arguments(text, tokenizer, model, beam_search=True):
     
     # Beam Search for position selection
     if beam_search:
-        indices1, indices2, _, _, _ = model.beam_search_position_selector(
+        indices1, indices2, score1, score2, topk_scores = model.beam_search_position_selector(
             start_cause_logits=start_cause_logits,
             end_cause_logits=end_cause_logits,
             start_effect_logits=start_effect_logits,
@@ -178,19 +178,63 @@ def extract_arguments(text, tokenizer, model, beam_search=True):
     else:
         cause_text2 = None
         effect_text2 = None
-    return cause_text1, effect_text1, signal_text, cause_text2, effect_text2
+
+    if beam_search: 
+        start_cause_probs = torch.softmax(start_cause_logits, dim=-1)
+        end_cause_probs = torch.softmax(end_cause_logits, dim=-1)
+        start_effect_probs = torch.softmax(start_effect_logits, dim=-1)
+        end_effect_probs = torch.softmax(end_effect_logits, dim=-1)
+
+        best_start_cause_score = start_cause_probs[start_cause1].item()
+        best_end_cause_score = end_cause_probs[end_cause1].item()
+        best_start_effect_score = start_effect_probs[start_effect1].item()
+        best_end_effect_score = end_effect_probs[end_effect1].item()
+
+        second_start_cause_score = start_cause_probs[start_cause2].item()
+        second_end_cause_score = end_cause_probs[end_cause2].item()
+        second_start_effect_score = start_effect_probs[start_effect2].item()
+        second_end_effect_score = end_effect_probs[end_effect2].item()
+
+        best_scores = {
+                "Start Cause Score": round(best_start_cause_score, 4),
+                "End Cause Score": round(best_end_cause_score, 4),
+                "Start Effect Score": round(best_start_effect_score, 4),
+                "End Effect Score": round(best_end_effect_score, 4),
+                "Total Best Score [sum of log-probability scores]": round(score1, 4)
+            }
+
+        second_best_scores = {
+            "Start Cause Score": round(second_start_cause_score, 4),
+            "End Cause Score": round(second_end_cause_score, 4),
+            "Start Effect Score": round(second_start_effect_score, 4),
+            "End Effect Score": round(second_end_effect_score, 4),
+            "Total Second Best Score [sum of log-probability scores]": round(score2, 4)
+            }
+
+        top5_scores = sorted(topk_scores.items(), key=lambda x: x[1], reverse=True)[:5]
+        top5_scores = [(k, round(v, 4)) for k, v in top5_scores]
+
+    
+    else:
+        best_scores = {}
+        second_best_scores = {}
+        top5_scores = {}
+        
+        
+    return cause_text1, effect_text1, signal_text, cause_text2, effect_text2, best_scores, second_best_scores, top5_scores 
+    
 
 st.title("Causal Relation Extraction")
 input_text = st.text_area("Enter your text here:", height=300)
-beam_search = st.radio("Enable Position Selector & Beam Search?", ('No', 'Yes')) == 'Yes'
+beam_search = st.radio("Enable Position Selector & Beam Search?", ('Yes', 'No')) == 'Yes'
 
 
 if st.button("Extract"):
     if input_text:
-        cause_text1, effect_text1, signal_text, cause_text2, effect_text2 = extract_arguments(input_text, tokenizer, model, beam_search=beam_search)
+        cause_text1, effect_text1, signal_text, cause_text2, effect_text2, best_scores, second_best_scores, top5_scores = extract_arguments(input_text, tokenizer, model, beam_search=beam_search)
 
         # Display first relation
-        st.markdown(f"<strong>Relation 1:</strong>", unsafe_allow_html=True)
+        st.write("## Relation 1:")
 
         if cause_text1 is None or effect_text1 is None: 
             st.write("The prediction is not correct for at least one span: The position of the predicted end token comes before the position of the start token.")
@@ -199,14 +243,26 @@ if st.button("Extract"):
             st.markdown(f"**Effect:** {effect_text1}", unsafe_allow_html=True)
             st.markdown(f"**Signal:** {signal_text}", unsafe_allow_html=True)
 
+            if beam_search:
+
+                # Display dictionary in Streamlit
+                st.markdown(f"<strong>Best Tuple Scores:</strong>", unsafe_allow_html=True)
+                st.json(best_scores)
+
         
-        # Display second relation if beam search is enabled
-        if beam_search:
+                # Display second relation if beam search is enabled
+                st.write("## Relation 2:")
+                st.markdown(f"**Cause:** {cause_text2}", unsafe_allow_html=True)
+                st.markdown(f"**Effect:** {effect_text2}", unsafe_allow_html=True)
+                st.markdown(f"**Signal:** {signal_text}", unsafe_allow_html=True)
+                
+                st.markdown(f"<strong>Second best Tuple Scores:</strong>", unsafe_allow_html=True)
+                st.json(second_best_scores)
 
-            st.markdown(f"<strong>Relation 2:</strong>", unsafe_allow_html=True)
-            st.markdown(f"**Cause:** {cause_text2}", unsafe_allow_html=True)
-            st.markdown(f"**Effect:** {effect_text2}", unsafe_allow_html=True)
-            st.markdown(f"**Signal:** {signal_text}", unsafe_allow_html=True)
+                st.markdown(f"<strong>top5_scores [sum of log-probability scores]:</strong>", unsafe_allow_html=True)
+                # Unpack top 5 scores 
+                # first, second, third, fourth, fifth = top_5_scores
+                st.json(top5_scores)
 
     else:
         st.warning("Please enter some text before extracting.")