Update app.py

app.py

@@ -166,6 +166,20 @@ with open("/home/user/app/models/retriever/document_index/documents.jsonl", "r")
         element = json.loads(line)
         io_map[element["text"]] = element["metadata"]["type"]
 
+            
+import json
+db_set = set()
+with open("models/retriever/intervention/gpt/db/document_index/documents.jsonl", "r") as r:
+    for line in r:
+        element = json.loads(line)
+        db_set.add(element["text"])
+
+with open("models/retriever/outcome/gpt/db/document_index/documents.jsonl", "r") as r:
+    for line in r:
+        element = json.loads(line)
+        db_set.add(element["text"])
+
+
 @st.cache_resource()
 def load_model():
 
@@ -249,7 +263,7 @@ def run_client():
         index=0  # Default to 'question'
     )
 
-    selection_options = ["DB Intervention (GPT)", "DB Outcome (GPT)", "Taxonomy Intervention (GPT)", "Taxonomy Outcome (GPT)"]
+    selection_options = ["DB Intervention", "DB Outcome", "Top-k DB Intervention", "Top-k Soft DB Outcome", "Taxonomy Intervention", "Taxonomy Outcome"]
     
     if analysis_type == "Retriever":
         # Selection list using selectbox
@@ -328,20 +342,38 @@ def run_client():
     
                     st.markdown(text, unsafe_allow_html=True)
                 else:
-                    response = relik_model.retrieve(text, k=10, batch_size=100, progress_bar=False)
+                    if "Top-k" in selection_options:
+                        response = relik_model.retrieve(text, k=100, 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)
 
-                    candidates_text = []
-                    for pred in response[0]:
-                        candidates_text.append(pred.document.text)
                     
-                    dict_of_ents_candidates, options_candidates = get_retriever_annotations_candidates(text, candidates_text)
+                    else:
+                        response = relik_model.retrieve(text, k=10, batch_size=400, progress_bar=False)
 
-                    text = """
-                    <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"][0:10]) + "</ul>"
+                        candidates_text = []
+                        for pred in response[0]:
+                            candidates_text.append(pred.document.text)
+
+                    if candidates_text:
+                        candidates_text = candidates_text[:5]
     
-                    st.markdown(text, unsafe_allow_html=True)
+                        dict_of_ents_candidates, options_candidates = get_retriever_annotations_candidates(text, candidates_text)
+    
+                        
+                        text = """
+                        <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)
+                    else:
+                        
+                        text = "<h2 style='color: black;'>No Candidates Found</h2>"
+
         else:
             st.error("Please enter some text.")