Update app.py

app.py

@@ -355,44 +355,46 @@ def run_client():
                     st.markdown(text, unsafe_allow_html=True)    
                 else:
                     if selection_list == "Top-k DB in Taxonomy Intervention" or selection_list == "Top-k DB in Taxonomy Outcome":
-                        response = relik_model.retrieve(text, k=30, batch_size=400, progress_bar=False)
+                        response = relik_model.retrieve(text, k=50, batch_size=400, progress_bar=False)
                         candidates_text = [pred.document.text for pred in response[0] if pred.document.text in db_set]
                     else:
                         response = relik_model.retrieve(text, k=10, batch_size=400, progress_bar=False)
                         candidates_text = [pred.document.text for pred in response[0]]
     
                     if candidates_text:
-                        st.session_state.candidates = candidates_text[:5]
-    
-                        dict_of_ents_candidates, options_candidates = get_retriever_annotations_candidates(text, st.session_state.candidates)
-    
-                        st.markdown("<h2 style='color: black;'>Possible Candidates:</h2>", unsafe_allow_html=True)
-    
-                        # Display the candidates with checkboxes
-                        for candidate in dict_of_ents_candidates["ents"]:
-                            if candidate in st.session_state.selected_candidates:
-                                checked = True
-                            else:
-                                checked = False
-                            if st.checkbox(candidate, key=candidate, value=checked):
-                                if candidate not in st.session_state.selected_candidates:
-                                    st.session_state.selected_candidates.append(candidate)
-                            else:
-                                if candidate in st.session_state.selected_candidates:
-                                    st.session_state.selected_candidates.remove(candidate)
-    
-                        # Button to save the selected candidates to a file
-                        if st.button("Save Selected Candidates"):
-                            write_candidates_to_file(text, dict_of_ents_candidates["ents"], st.session_state.selected_candidates)
-                            st.success("Selected candidates have been saved to file.")
-                    else:
-                        st.session_state.candidates = []
-                        st.session_state.selected_candidates = []
-                        st.markdown("<h2 style='color: black;'>No Candidates Found</h2>", unsafe_allow_html=True)
+                        st.session_state.candidates = candidates_text[:10]
 
         else:
             st.error("Please enter some text.")
 
+    if st.session_state.candidates:
+        dict_of_ents_candidates, options_candidates = get_retriever_annotations_candidates(text, st.session_state.candidates)
+
+        st.markdown("<h2 style='color: black;'>Possible Candidates:</h2>", unsafe_allow_html=True)
+
+        # Display the candidates with checkboxes
+        for candidate in dict_of_ents_candidates["ents"]:
+            if candidate in st.session_state.selected_candidates:
+                checked = True
+            else:
+                checked = False
+            if st.checkbox(candidate, key=candidate, value=checked):
+                if candidate not in st.session_state.selected_candidates:
+                    st.session_state.selected_candidates.append(candidate)
+            else:
+                if candidate in st.session_state.selected_candidates:
+                    st.session_state.selected_candidates.remove(candidate)
+
+        # Button to save the selected candidates to a file
+        if st.button("Save Selected Candidates"):
+            write_candidates_to_file(text, dict_of_ents_candidates["ents"], st.session_state.selected_candidates)
+            st.success("Selected candidates have been saved to file.")
+
+    else:
+        st.session_state.candidates = []
+        st.session_state.selected_candidates = []
+        st.markdown("<h2 style='color: black;'>No Candidates Found</h2>", unsafe_allow_html=True)
+
  
 if __name__ == "__main__":
     run_client()