Update app.py

app.py

@@ -148,6 +148,17 @@ def get_retriever_annotations(response):
     label_in_text = set(l for l in dict_of_ents["ents"])
     options = {"ents": label_in_text, "colors": get_random_color(label_in_text)}
     return dict_of_ents, options
+    
+
+def get_retriever_annotations_candidates(text, ents):
+    el_link_wrapper = "<link rel='stylesheet' href='https://cdnjs.cloudflare.com/ajax/libs/font-awesome/6.4.2/css/all.min.css'><a href='https://en.wikipedia.org/wiki/{}' style='color: #414141'><i class='fa-brands fa-wikipedia-w fa-xs'></i> <span style='font-size: 1.0em; font-family: monospace'> {}</span></a>"
+    # swap labels key with ents
+    dict_of_ents = {"text": text, "ents": ents}
+    label_in_text = set(l for l in dict_of_ents["ents"])
+    options = {"ents": label_in_text, "colors": get_random_color(label_in_text)}
+    return dict_of_ents, options
+
+    
 import json
 io_map = {}
 with open("/home/user/app/models/retriever/document_index/documents.jsonl", "r") as r:
@@ -164,48 +175,66 @@ def load_model():
 
     )
 
-    retriever_intervention = GoldenRetriever(
-        question_encoder="models/retriever/level-4-small-no-negative-interventions/question_encoder",
-        document_index="models/retriever/level-4-small-no-negative-interventions/document_index"
+    retriever_intervention_gpt_taxonomy = GoldenRetriever(
+        question_encoder="models/retriever/intervention/gpt/taxonomy/question_encoder",
+        document_index="models/retriever/intervention/gpt/taxonomy/document_index"
 
     )
 
-    retriever_outcome = GoldenRetriever(
-        question_encoder="models/retriever/level-4-small-no-negative-outcomes/question_encoder",
-        document_index="models/retriever/level-4-small-no-negative-outcomes/document_index"
+    retriever_intervention_gpt_llama_taxonomy = GoldenRetriever(
+        question_encoder="models/retriever/intervention/gpt+llama/taxonomy/question_encoder",
+        document_index="models/retriever/intervention/gpt+llama/taxonomy/document_index"
 
     )
 
-    retriever_question_db = GoldenRetriever(
-        question_encoder="/home/user/app/models/retriever/level-4-small-no-negatives/question_encoder",
-        document_index="/home/user/app/models/retriever/level-4-small-no-negatives/document_index"
+    
+    retriever_intervention_gpt_db = GoldenRetriever(
+        question_encoder="models/retriever/intervention/gpt/db/question_encoder",
+        document_index="models/retriever/intervention/gpt/db/document_index"
+
+    )
+
+    retriever_intervention_gpt_llama_db = GoldenRetriever(
+        question_encoder="models/retriever/intervention/gpt+llama/db/question_encoder",
+        document_index="models/retriever/intervention/gpt+llama/db/document_index"
 
     )
 
-    retriever_intervention_db = GoldenRetriever(
-        question_encoder="models/retriever/level-4-small-no-negative-interventions/question_encoder",
-        document_index="models/retriever/level-4-small-no-negative-interventions/document_index_db"
+
+    
+    retriever_outcome_gpt_taxonomy = GoldenRetriever(
+        question_encoder="models/retriever/outcome/gpt/taxonomy/question_encoder",
+        document_index="models/retriever/outcome/gpt/taxonomy/document_index"
 
     )
 
-    retriever_outcome_db = GoldenRetriever(
-        question_encoder="models/retriever/level-4-small-no-negative-outcomes/question_encoder",
-        document_index="models/retriever/level-4-small-no-negative-outcomes/document_index_db"
+    retriever_outcome_gpt_llama_taxonomy = GoldenRetriever(
+        question_encoder="models/retriever/outcome/gpt+llama/taxonomy/question_encoder",
+        document_index="models/retriever/outcome/gpt+llama/taxonomy/document_index"
 
     )
 
+    
+    retriever_outcome_gpt_db = GoldenRetriever(
+        question_encoder="models/retriever/outcome/gpt/db/question_encoder",
+        document_index="models/retriever/outcome/gpt/db/document_index"
+
+    )
+
+    retriever_outcome_gpt_llama_db = GoldenRetriever(
+        question_encoder="models/retriever/outcome/gpt+llama/db/question_encoder",
+        document_index="models/retriever/outcome/gpt+llama/db/document_index"
+
+    )
 
     reader = RelikReaderForSpanExtraction("/home/user/app/models/small-extended-large-batch",
                                         dataset_kwargs={"use_nme": True})
 
     relik_question = Relik(reader=reader, retriever=retriever_question, window_size="none", top_k=100, task="span", device="cpu", document_index_device="cpu")
-    relik_intervention = Relik(reader=reader, retriever=retriever_intervention, window_size="none", top_k=100, task="span", device="cpu", document_index_device="cpu")
-    relik_outcome = Relik(reader=reader, retriever=retriever_outcome, window_size="none", top_k=100, task="span", device="cpu", document_index_device="cpu")
-    relik_question_db = Relik(reader=reader, retriever=retriever_question_db, window_size="none", top_k=100, task="span", device="cpu", document_index_device="cpu")
-    relik_intrervention_db = Relik(reader=reader, retriever=retriever_intervention_db, window_size="none", top_k=100, task="span", device="cpu", document_index_device="cpu")
-    relik_outcome_db = Relik(reader=reader, retriever=retriever_outcome_db, window_size="none", top_k=100, task="span", device="cpu", document_index_device="cpu")
 
-    return [relik_question, relik_intervention, relik_outcome, relik_question_db, relik_intrervention_db, relik_outcome_db]
+        selection_options = ["DB Intervention (GPT)", "DB Outcome (GPT)", "DB Intervention (GPT+Llama)", "DB Outcome (GPT+Llama)", "Taxonomy Intervention (GPT)", "Taxonomy Outcome (GPT)", "Taxonomy Intervention (GPT+Llama)", "Taxonomy Outcome (GPT+Llama)"]
+
+    return [relik_question, retriever_intervention_gpt_db, retriever_outcome_gpt_db, retriever_intervention_gpt_llama_db, retriever_outcome_gpt_llama_db, retriever_intervention_gpt_taxonomy, retriever_outcome_gpt_taxonomy, retriever_intervention_gpt_llama_taxonomy, retriever_outcome_gpt_llama_taxonomy]
 
 def set_intro(css):
     # intro
@@ -239,10 +268,19 @@ def run_client():
         # Radio button selection
     analysis_type = st.radio(
         "Choose analysis type:",
-        options=["question", "intervention", "outcome", "db intervention", "db outcome"],
+        options=["Retriever", "Entity Linking"],
         index=0  # Default to 'question'
     )
+
+    selection_options = ["DB Intervention (GPT)", "DB Outcome (GPT)", "DB Intervention (GPT+Llama)", "DB Outcome (GPT+Llama)", "Taxonomy Intervention (GPT)", "Taxonomy Outcome (GPT)", "Taxonomy Intervention (GPT+Llama)", "Taxonomy Outcome (GPT+Llama)"]
     
+    if analysis_type == "Retriever"
+        # Selection list using selectbox
+        selection_list = st.selectbox(
+            "Select an option:",
+            options=options
+        )
+
     # text input
     text = st.text_area(
         "Enter Text Below:",
@@ -273,40 +311,24 @@ def run_client():
         entity_linking_bool = False
 
 
-        if analysis_type == "question":
+        if analysis_type == "Entity Linking":
             relik_model = st.session_state["relik_model"][0]
             entity_linking_bool = True
-        elif analysis_type == "intervention":
-            relik_model = st.session_state["relik_model"][1]
-        elif  analysis_type == "outcome":
-            relik_model = st.session_state["relik_model"][2]
-
-        elif  analysis_type == "db intervention":
-            relik_model = st.session_state["relik_model"][4]
-            
-        elif  analysis_type == "db outcome":
-            relik_model = st.session_state["relik_model"][5]
-
         else:
-            relik_model = st.session_state["relik_model"][3]
+            model_idx = selection_options.index(selection_list)
+            relik_model = st.session_state["relik_model"][model_idx+1]
+            
             
         text = text.strip()
         if text:
             st.markdown("####")
             with st.spinner(text="In progress"):
-                response = relik_model(text)
+                if entity_linking_bool:
+                    response = relik_model(text)
+                    
+                    dict_of_ents, options = get_el_annotations(response=response)
+                    dict_of_ents_candidates, options_candidates = get_retriever_annotations(response=response)
 
-                # response = requests.post(RELIK, json=text)
-                # if response.status_code != 200:
-                #     st.error("Error: {}".format(response.status_code))
-                # else:
-                #     response = response.json()
-
-                # st.markdown("##")
-                dict_of_ents, options = get_el_annotations(response=response)
-                dict_of_ents_candidates, options_candidates = get_retriever_annotations(response=response)
-
-                if entity_linking_bool: 
                     st.markdown("#### Entity Linking")
 
                     display = displacy.render(
@@ -329,10 +351,18 @@ def run_client():
     
                     st.markdown(text, unsafe_allow_html=True)
                 else:
+                    response = relik_model.retrieve(text, k=10, batch_size=100, progress_bar=False)
+
+                    candidates_text = []
+                    for pred in response[0]:
+                        candidates.append(pred.document.text)
+                    
+                    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"][2:12]) + "</ul>"
+                    """ + "".join(f"<li style='color: black;'>{candidate}</li>" for candidate in dict_of_ents_candidates["ents"][0:10]) + "</ul>"
     
                     st.markdown(text, unsafe_allow_html=True)
         else: