Update app.py

app.py

@@ -3,6 +3,12 @@ import re
 import time
 from pathlib import Path
 
+from relik.retriever import GoldenRetriever
+
+from relik.retriever.indexers.inmemory import InMemoryDocumentIndex
+from relik.retriever.indexers.document import DocumentStore
+from relik.retriever import GoldenRetriever
+from relik.reader.pytorch_modules.span import RelikReaderForSpanExtraction
 import requests
 import streamlit as st
 from spacy import displacy
@@ -14,7 +20,13 @@ from streamlit_extras.stylable_container import stylable_container
 import random
 
 from relik.inference.annotator import Relik
-
+from relik.inference.data.objects import (
+    AnnotationType,
+    RelikOutput,
+    Span,
+    TaskType,
+    Triples,
+)
 
 def get_random_color(ents):
     colors = {}
@@ -93,44 +105,31 @@ def generate_pastel_colors(n):
 
 
 def set_sidebar(css):
-    white_link_wrapper = "<link rel='stylesheet' href='https://cdnjs.cloudflare.com/ajax/libs/font-awesome/6.4.2/css/all.min.css'><a href='{}'>{}</a>"
     with st.sidebar:
         st.markdown(f"<style>{css}</style>", unsafe_allow_html=True)
         st.image(
-            "http://nlp.uniroma1.it/static/website/sapienza-nlp-logo-wh.svg",
+            "https://upload.wikimedia.org/wikipedia/commons/8/87/The_World_Bank_logo.svg",
             use_column_width=True,
         )
-        st.markdown("## ReLiK")
-        st.write(
-            f"""
-                - {white_link_wrapper.format("#", "<i class='fa-solid fa-file'></i>&nbsp; Paper")}
-                - {white_link_wrapper.format("https://github.com/SapienzaNLP/relik", "<i class='fa-brands fa-github'></i>&nbsp; GitHub")}
-                - {white_link_wrapper.format("https://hub.docker.com/repository/docker/sapienzanlp/relik", "<i class='fa-brands fa-docker'></i>&nbsp; Docker Hub")}
-                """,
-            unsafe_allow_html=True,
-        )
-        st.markdown("## Sapienza NLP")
-        st.write(
-            f"""
-                - {white_link_wrapper.format("https://nlp.uniroma1.it", "<i class='fa-solid fa-globe'></i>&nbsp; Webpage")}
-                - {white_link_wrapper.format("https://github.com/SapienzaNLP", "<i class='fa-brands fa-github'></i>&nbsp; GitHub")}
-                - {white_link_wrapper.format("https://twitter.com/SapienzaNLP", "<i class='fa-brands fa-twitter'></i>&nbsp; Twitter")}
-                - {white_link_wrapper.format("https://www.linkedin.com/company/79434450", "<i class='fa-brands fa-linkedin'></i>&nbsp; LinkedIn")}
-                """,
-            unsafe_allow_html=True,
-        )
-
+        st.markdown("### World Bank")
+        st.markdown("### DIME")
 
 def get_el_annotations(response):
-    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>"
+    i_link_wrapper = "<link rel='stylesheet' href='https://cdnjs.cloudflare.com/ajax/libs/font-awesome/6.4.2/css/all.min.css'><a href='https://developmentevidence.3ieimpact.org/taxonomy-search-detail/intervention/disaggregated-intervention/{}' style='color: #414141'> <span style='font-size: 1.0em; font-family: monospace'> Intervention {}</span></a>"
+    o_link_wrapper = "<link rel='stylesheet' href='https://cdnjs.cloudflare.com/ajax/libs/font-awesome/6.4.2/css/all.min.css'><a href='https://developmentevidence.3ieimpact.org/taxonomy-search-detail/intervention/disaggregated-outcome/{}' style='color: #414141'><span style='font-size: 1.0em; font-family: monospace'> Outcome: {}</span></a>"
     # swap labels key with ents
     ents = [
         {
             "start": l.start,
             "end": l.end,
-            "label": el_link_wrapper.format(l.label.replace(" ", "_"), l.label),
+            "label": i_link_wrapper.format(l.label[0].upper() + l.label[1:].replace("/", "%2").replace(" ", "%20").replace("&","%26"), l.label),
+        } if io_map[l.label] == "intervention" else 
+        {
+            "start": l.start,
+            "end": l.end,
+            "label": o_link_wrapper.format(l.label[0].upper() + l.label[1:].replace("/", "%2").replace(" ", "%20").replace("&","%26"), l.label), 
         }
-        for l in response.labels
+        for l in response.spans
     ]
     dict_of_ents = {"text": response.text, "ents": ents}
     label_in_text = set(l["label"] for l in dict_of_ents["ents"])
@@ -138,24 +137,58 @@ def get_el_annotations(response):
     return dict_of_ents, options
 
 
+
+def get_retriever_annotations(response):
+    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
+    ents = [l.text
+        for l in response.candidates[TaskType.SPAN]
+    ]
+    dict_of_ents = {"text": response.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/model/retriever/document_index/documents.jsonl", "r") as r:
+    for line in r:
+        element = json.loads(line)
+        io_map[element["text"]] = element["metadata"]["type"]
+
 @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",
+
+    retriever = GoldenRetriever(
+        question_encoder="/home/user/app/model/retriever/question_encoder",
+        document_index=InMemoryDocumentIndex(
+                documents=DocumentStore.from_file(
+                    "/home/user/app/model/retriever/document_index/documents.jsonl"
+                ),
+                metadata_fields=["definition"],
+                separator=' <def> ',
+                device="cuda"
+            ),
+            devide="cuda"
+
     )
+    retriever.index()
 
+    reader = RelikReaderForSpanExtraction("/home/user/app/model/reader",
+                                        dataset_kwargs={"use_nme": True})
+
+    relik = Relik(reader=reader, retriever=retriever, window_size="none", top_k=100, task="span", device="cuda", document_index_device="cpu")
+
+    return relik
 
 def set_intro(css):
     # intro
-    st.markdown("# ReLik")
+
+    st.markdown("# CausalAI")
+    st.image(
+    "http://35.237.102.64/public/logo.png",
+)
     st.markdown(
-        "### Retrieve, Read and LinK: Fast and Accurate Entity Linking and Relation Extraction on an Academic Budget"
+        "### 3ie taxonomy level 4 Intervention/Outcome candidate retriever with Entity Linking"
     )
     # st.markdown(
     #     "This is a front-end for the paper [Universal Semantic Annotator: the First Unified API "
@@ -163,16 +196,13 @@ def set_intro(css):
     #     "[Riccardo Orlando](https://riccorl.github.io), [Simone Conia](https://c-simone.github.io/), "
     #     "[Stefano Faralli](https://corsidilaurea.uniroma1.it/it/users/stefanofaralliuniroma1it), and [Roberto Navigli](https://www.diag.uniroma1.it/navigli/)."
     # )
-    badge(type="github", name="sapienzanlp/relik")
-    badge(type="pypi", name="relik")
-
 
 def run_client():
     with open(Path(__file__).parent / "style.css") as f:
         css = f.read()
 
     st.set_page_config(
-        page_title="ReLik",
+        page_title="CausalAI",
         page_icon="🦮",
         layout="wide",
     )
@@ -182,7 +212,7 @@ def run_client():
     # text input
     text = st.text_area(
         "Enter Text Below:",
-        value="Michael Jordan was one of the best players in the NBA.",
+        value="How does unconditional cash transver affect to reduce poverty?",
         height=200,
         max_chars=1500,
     )
@@ -191,8 +221,8 @@ def run_client():
         key="annotate_button",
         css_styles="""
             button {
-                background-color: #802433;
-                color: white;
+                background-color: #a8ebff;
+                color: black;
                 border-radius: 25px;
             }
             """,
@@ -212,6 +242,7 @@ def run_client():
             st.markdown("#### Entity Linking")
             with st.spinner(text="In progress"):
                 response = relik_model(text)
+
                 # response = requests.post(RELIK, json=text)
                 # if response.status_code != 200:
                 #     st.error("Error: {}".format(response.status_code))
@@ -220,14 +251,23 @@ def run_client():
 
                 # st.markdown("##")
                 dict_of_ents, options = get_el_annotations(response=response)
+                dict_of_ents_candidates, options_candidates = get_retriever_annotations(response=response)
+
                 display = displacy.render(
                     dict_of_ents, manual=True, style="ent", options=options
                 )
+
+
                 display = display.replace("\n", " ")
+
                 # heurstic, prevents split of annotation decorations
                 display = display.replace("border-radius: 0.35em;", "border-radius: 0.35em; white-space: nowrap;")
+
                 with st.container():
                     st.write(display, unsafe_allow_html=True)
+                
+                text = "## Possible Candidates:\n- " + "\n- ".join([candidate for candidate in dict_of_ents_candidates["ents"][:5]])
+                st.markdown(text)
 
         else:
             st.error("Please enter some text.")