Upload app

app.py

@@ -123,13 +123,26 @@ def set_sidebar(css):
 
 def get_el_annotations(response):
     # swap labels key with ents
-    dict_of_ents = {"text": response.text, "ents": []}
-    dict_of_ents["ents"] = response.labels
+    ents = [{"start": l.start, "end": l.end, "label": l.label} for l in response.labels]
+    dict_of_ents = {"text": response.text, "ents": ents}
     label_in_text = set(l["label"] for l in dict_of_ents["ents"])
     options = {"ents": label_in_text, "colors": get_random_color(label_in_text)}
     return dict_of_ents, options
 
 
+@st.cache_resource()
+def load_model():
+    return Relik(
+        question_encoder="/home/user/app/models/relik-retriever-small-aida-blink-pretrain-omniencoder/question_encoder",
+        document_index="/home/user/app/models/relik-retriever-small-aida-blink-pretrain-omniencoder/document_index_filtered",
+        reader="/home/user/app/models/relik-reader-aida-deberta-small",
+        top_k=100,
+        window_size=32,
+        window_stride=16,
+        candidates_preprocessing_fn="relik.inference.preprocessing.wikipedia_title_and_openings_preprocessing",
+    )
+
+
 def set_intro(css):
     # intro
     st.markdown("# ReLik")
@@ -161,7 +174,7 @@ def run_client():
     # text input
     text = st.text_area(
         "Enter Text Below:",
-        value="Obama went to Rome for a quick vacation.",
+        value="Michael Jordan was one of the best players in the NBA.",
         height=200,
         max_chars=500,
     )
@@ -179,15 +192,9 @@ def run_client():
         submit = st.button("Annotate")
     # submit = st.button("Run")
 
-    relik = Relik(
-        question_encoder="/home/user/app/models/relik-retriever-small-aida-blink-pretrain-omniencoder/question_encoder",
-        document_index="/home/user/app/models/relik-retriever-small-aida-blink-pretrain-omniencoder/document_index_filtered",
-        reader="/home/user/app/models/relik-reader-aida-deberta-small",
-        top_k=100,
-        window_size=32,
-        window_stride=16,
-        candidates_preprocessing_fn="relik.inference.preprocessing.wikipedia_title_and_openings_preprocessing",
-    )
+    if "relik_model" not in st.session_state.keys():
+        st.session_state["relik_model"] = load_model()
+    relik_model = st.session_state["relik_model"]
 
     # ReLik API call
     if submit:
@@ -196,26 +203,26 @@ def run_client():
             st.markdown("####")
             st.markdown("#### Entity Linking")
             with st.spinner(text="In progress"):
-                response = relik(text)
+                response = relik_model(text)
                 # response = requests.post(RELIK, json=text)
                 # if response.status_code != 200:
                 #     st.error("Error: {}".format(response.status_code))
                 # else:
                 #     response = response.json()
 
-                    # Entity Linking
-                    # with stylable_container(
-                    #     key="container_with_border",
-                    #     css_styles="""
-                    #         {
-                    #             border: 1px solid rgba(49, 51, 63, 0.2);
-                    #             border-radius: 0.5rem;
-                    #             padding: 0.5rem;
-                    #             padding-bottom: 2rem;
-                    #         }
-                    #         """,
-                    # ):
-                    # st.markdown("##")
+                # Entity Linking
+                # with stylable_container(
+                #     key="container_with_border",
+                #     css_styles="""
+                #         {
+                #             border: 1px solid rgba(49, 51, 63, 0.2);
+                #             border-radius: 0.5rem;
+                #             padding: 0.5rem;
+                #             padding-bottom: 2rem;
+                #         }
+                #         """,
+                # ):
+                # st.markdown("##")
                 dict_of_ents, options = get_el_annotations(response=response)
                 display = displacy.render(
                     dict_of_ents, manual=True, style="ent", options=options