Update app.py

app.py

@@ -301,6 +301,11 @@ def run_client():
         st.session_state["relik_model"] = load_model()
     relik_model = st.session_state["relik_model"][0]
 
+    if 'candidates' not in st.session_state:
+        st.session_state['candidates'] = []
+    if 'selected_candidates' not in st.session_state:
+        st.session_state['selected_candidates'] = []
+
     # ReLik API call
     if submit:
 
@@ -344,47 +349,40 @@ def run_client():
                     <ul style='color: black;'>
                     """ + candidate_text + "</ul>"
     
-                    st.markdown(text, unsafe_allow_html=True)
+                    st.markdown(text, unsafe_allow_html=True)    
                 else:
-                    if selection_list == "Top-k DB in Taxonomy Intervention" or selection_list == "Top-k DB in Taxonmy Outcome":
-                        response = relik_model.retrieve(text, k=30, batch_size=400, progress_bar=False)
-                           
-                        candidates_text = []
-                        for pred in response[0]:
-                            if pred.document.text in db_set:
-                                candidates_text.append(pred.document.text)
-
-                    
+                    if selection_list == "Top-k DB in Taxonomy Intervention" or selection_list == "Top-k DB in Taxonomy Outcome":
+                        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 = []
-                        for pred in response[0]:
-                            candidates_text.append(pred.document.text)
-
+                        candidates_text = [pred.document.text for pred in response[0]]
+    
                     if candidates_text:
-                        candidates_text = candidates_text[:5]
+                        st.session_state.candidates = candidates_text[:10]
     
-                        dict_of_ents_candidates, options_candidates = get_retriever_annotations_candidates(text, candidates_text)
+                        dict_of_ents_candidates, options_candidates = get_retriever_annotations_candidates(text, st.session_state.candidates)
     
-
-                        text = "<h2 style='color: black;'>Possible Candidates:</h2>"
-                        #<h2 style='color: black;'>Possible Candidates:</h2>
-                        # <ul style='color: black;'>
-                        # """ + "".join(f"<li style='color: black;'>{candidate}</li>" for candidate in dict_of_ents_candidates["ents"]) + "</ul>"
-        
-                        st.markdown(text, unsafe_allow_html=True)
-                       # Multiselect widget for selecting candidates
-                        selected_candidates = st.multiselect(
-                            "Select Candidates:",
-                            dict_of_ents_candidates["ents"],  # Options for the multiselect
-                            help="Select one or more candidates from the list."
-                        )
-                       # Button to save the selected candidates to a file
+                        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 st.checkbox(candidate, key=candidate):
+                                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"], 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)
+    
                     else:
                         
                         text = "<h2 style='color: black;'>No Candidates Found</h2>"