Update app.py

app.py

@@ -173,12 +173,56 @@ def load_model():
     )
     retriever.index()
 
+
+
+    retriever_question = GoldenRetriever(
+        question_encoder="/home/user/app/models/retriever/question_encoder",
+        document_index="home/user/app/models/retriever/document_index/questions"
+
+    )
+
+    retriever_intervention = GoldenRetriever(
+        question_encoder="/home/user/app/models/retriever/question_encoder",
+        document_index="home/user/app/models/retriever/document_index/interventions"
+
+    )
+
+    retriever_outcome = GoldenRetriever(
+        question_encoder="/home/user/app/models/retriever/question_encoder",
+        document_index="home/user/app/models/retriever/document_index/outcomes"
+
+    )
+
+    retriever_question_db = GoldenRetriever(
+        question_encoder="/home/user/app/models/retriever/question_encoder",
+        document_index="home/user/app/models/retriever/document_index/question_db"
+
+    )
+
+    retriever_intervention_db = GoldenRetriever(
+        question_encoder="/home/user/app/models/retriever/question_encoder",
+        document_index="home/user/app/models/retriever/document_index/interventions_db"
+
+    )
+
+    retriever_outcome_db = GoldenRetriever(
+        question_encoder="/home/user/app/models/retriever/question_encoder",
+        document_index="home/user/app/models/retriever/document_index/outcomes_db"
+
+    )
+
+
     reader = RelikReaderForSpanExtraction("/home/user/app/models/small-extended-large-batch",
                                         dataset_kwargs={"use_nme": True})
 
-    relik = Relik(reader=reader, retriever=retriever,  window_size=32, window_stride=16, top_k=100, task="span", device="cpu", document_index_device="cpu")
+    relik_question = Relik(reader=reader, retriever=retriever_question, window_size="none", top_k=100, task="span", device="cuda", document_index_device="cpu")
+    relik_intervention = Relik(reader=reader, retriever=retriever_intervention, window_size="none", top_k=100, task="span", device="cuda", document_index_device="cpu")
+    relik_outcome = Relik(reader=reader, retriever=retriever_outcome, window_size="none", top_k=100, task="span", device="cuda", document_index_device="cpu")
+    relik_question_db = Relik(reader=reader, retriever=retriever_question_db, window_size="none", top_k=100, task="span", device="cuda", document_index_device="cpu")
+    relik_intrervention_db = Relik(reader=reader, retriever=retriever_intervention_db, window_size="none", top_k=100, task="span", device="cuda", document_index_device="cpu")
+    relik_outcome_db = Relik(reader=reader, retriever=retriever_outcome_db, window_size="none", top_k=100, task="span", device="cuda", document_index_device="cpu")
 
-    return relik
+    return [relik_question, relik_intervention, relik_outcome, relik_question_db, relik_intrervention_db, relik_outcome_db]
 
 def set_intro(css):
     # intro
@@ -209,6 +253,13 @@ def run_client():
     set_sidebar(css)
     set_intro(css)
 
+        # Radio button selection
+    analysis_type = st.radio(
+        "Choose analysis type:",
+        options=["intervention", "outcome", "question", "db intervention", "db outcome", "db question"],
+        index=2  # Default to 'question'
+    )
+    
     # text input
     text = st.text_area(
         "Enter Text Below:",
@@ -232,10 +283,30 @@ def run_client():
 
     if "relik_model" not in st.session_state.keys():
         st.session_state["relik_model"] = load_model()
-    relik_model = st.session_state["relik_model"]
+    relik_model = st.session_state["relik_model"][0]
 
     # ReLik API call
     if submit:
+        if analysis_type == "question":
+            relik_model = st.session_state["relik_model"][0]
+
+        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 question":
+            relik_model = st.session_state["relik_model"][3]
+
+        elif  analysis_type == "db intervention":
+            relik_model = st.session_state["relik_model"][4]
+
+        elif  analysis_type == "db outcome":
+            print("hola")
+            relik_model = st.session_state["relik_model"][5]
+
+        else:
+            relik_model = st.session_state["relik_model"][0]
+            
         text = text.strip()
         if text:
             st.markdown("####")
@@ -269,7 +340,7 @@ def run_client():
                 text = """
                 <h2 style='color: black;'>Possible Candidates:</h2>
                 <ul style='color: black;'>
-                """ + "".join(f"<li>{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"][2:12]) + "</ul>"
 
                 st.markdown(text, unsafe_allow_html=True)