Update app.py

app.py

@@ -243,10 +243,13 @@ def set_intro(css):
     # )
 
 def write_candidates_to_file(text, candidates, selected_candidates):
+    # Write the selected candidates to a file
+    # Ensure the directory exists
+    os.makedirs("logs", exist_ok=True)
     # Write the selected candidates to a file
     with open("logs/selected_candidates.tsv", "a") as file:
-        file.write(text + "\t" + str(candidates) + str(selected_candidates) + "\n")
-
+        file.write(f"{text}\t{str(candidates)}\t{str(selected_candidates)}\n")
+   
 
 def run_client():
     with open(Path(__file__).parent / "style.css") as f:
@@ -352,15 +355,36 @@ 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=50, batch_size=400, progress_bar=False)
+                        response = relik_model.retrieve(text, k=30, 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[:10]
+                        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 = []
@@ -369,24 +393,6 @@ def run_client():
         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 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"], st.session_state.selected_candidates)
-            st.success("Selected candidates have been saved to file.")
-
+ 
 if __name__ == "__main__":
     run_client()