import os 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 from streamlit_extras.badges import badge from streamlit_extras.stylable_container import stylable_container import logging # Configure logging logging.basicConfig(level=logging.INFO, format='%(asctime)s %(message)s') logger = logging.getLogger() # RELIK = os.getenv("RELIK", "localhost:8000/api/entities") 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 = {} random_colors = generate_pastel_colors(len(ents)) for ent in ents: colors[ent] = random_colors.pop(random.randint(0, len(random_colors) - 1)) return colors def floatrange(start, stop, steps): if int(steps) == 1: return [stop] return [ start + float(i) * (stop - start) / (float(steps) - 1) for i in range(steps) ] def hsl_to_rgb(h, s, l): def hue_2_rgb(v1, v2, v_h): while v_h < 0.0: v_h += 1.0 while v_h > 1.0: v_h -= 1.0 if 6 * v_h < 1.0: return v1 + (v2 - v1) * 6.0 * v_h if 2 * v_h < 1.0: return v2 if 3 * v_h < 2.0: return v1 + (v2 - v1) * ((2.0 / 3.0) - v_h) * 6.0 return v1 # if not (0 <= s <= 1): raise ValueError, "s (saturation) parameter must be between 0 and 1." # if not (0 <= l <= 1): raise ValueError, "l (lightness) parameter must be between 0 and 1." r, b, g = (l * 255,) * 3 if s != 0.0: if l < 0.5: var_2 = l * (1.0 + s) else: var_2 = (l + s) - (s * l) var_1 = 2.0 * l - var_2 r = 255 * hue_2_rgb(var_1, var_2, h + (1.0 / 3.0)) g = 255 * hue_2_rgb(var_1, var_2, h) b = 255 * hue_2_rgb(var_1, var_2, h - (1.0 / 3.0)) return int(round(r)), int(round(g)), int(round(b)) def generate_pastel_colors(n): """Return different pastel colours. Input: n (integer) : The number of colors to return Output: A list of colors in HTML notation (eg.['#cce0ff', '#ffcccc', '#ccffe0', '#f5ccff', '#f5ffcc']) Example: >>> print generate_pastel_colors(5) ['#cce0ff', '#f5ccff', '#ffcccc', '#f5ffcc', '#ccffe0'] """ if n == 0: return [] # To generate colors, we use the HSL colorspace (see http://en.wikipedia.org/wiki/HSL_color_space) start_hue = 0.0 # 0=red 1/3=0.333=green 2/3=0.666=blue saturation = 1.0 lightness = 0.9 # We take points around the chromatic circle (hue): # (Note: we generate n+1 colors, then drop the last one ([:-1]) because # it equals the first one (hue 0 = hue 1)) return [ "#%02x%02x%02x" % hsl_to_rgb(hue, saturation, lightness) for hue in floatrange(start_hue, start_hue + 1, n + 1) ][:-1] def set_sidebar(css): with st.sidebar: st.markdown(f"", unsafe_allow_html=True) st.image( "https://upload.wikimedia.org/wikipedia/commons/8/87/The_World_Bank_logo.svg", use_column_width=True, ) st.markdown("### World Bank") st.markdown("### DIME") def get_el_annotations(response): i_link_wrapper = " Intervention {}" o_link_wrapper = " Outcome: {}" # swap labels key with ents ents = [ { "start": l.start, "end": l.end, "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.spans ] 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 def get_retriever_annotations(response): el_link_wrapper = " {}" # 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 def get_retriever_annotations_candidates(text, ents): el_link_wrapper = " {}" # 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: for line in 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(): retriever_question = GoldenRetriever( question_encoder="/home/user/app/models/retriever/question_encoder", document_index="/home/user/app/models/retriever/document_index/questions" ) retriever_intervention_gpt_taxonomy = GoldenRetriever( question_encoder="models/retriever/intervention/gpt+llama/taxonomy/question_encoder", document_index="models/retriever/intervention/gpt+llama/taxonomy/document_index" ) retriever_intervention_gpt_db = GoldenRetriever( question_encoder="models/retriever/intervention/gpt+llama/db/question_encoder", document_index="models/retriever/intervention/gpt+llama/db/document_index" ) retriever_outcome_gpt_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+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") return [relik_question, retriever_intervention_gpt_db, retriever_outcome_gpt_db, retriever_intervention_gpt_taxonomy, retriever_outcome_gpt_taxonomy] def set_intro(css): # intro st.markdown("# ImpactAI") st.image( "http://35.237.102.64/public/logo.png", ) st.markdown( "### 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 " # "for WSD, SRL and Semantic Parsing](https://www.researchgate.net/publication/360671045_Universal_Semantic_Annotator_the_First_Unified_API_for_WSD_SRL_and_Semantic_Parsing), which will be presented at LREC 2022 by " # "[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/)." # ) from datetime import datetime from pathlib import Path from huggingface_hub import HfApi, CommitScheduler from uuid import uuid4 # Access token from environment variable # api = HfApi() # api.set_access_token(os.getenv("HF_TOKEN")) JSON_DATASET_DIR = Path("json_demo_selected_io") JSON_DATASET_DIR.mkdir(parents=True, exist_ok=True) JSON_DATASET_PATH = JSON_DATASET_DIR / f"train-{uuid4()}.json" scheduler = CommitScheduler( repo_id="demo-retriever", repo_type="dataset", folder_path=JSON_DATASET_DIR, path_in_repo="data", token=os.getenv("HF_TOKEN") ) def write_candidates_to_file(text, candidates, selected_candidates): logger.info(f"Text: {text}\tCandidates: {str(candidates)}\tSelected Candidates: {str(selected_candidates)}\n") with scheduler.lock: with JSON_DATASET_PATH.open("a") as f: json.dump({"text": text, "Candidates": [candidate for candidate in candidates], "Selected Candidates": [candidate for candidate in selected_candidates], "datetime": datetime.now().isoformat()}, f) f.write("\n") def run_client(): with open(Path(__file__).parent / "style.css") as f: css = f.read() st.set_page_config( page_title="ImpactAI", page_icon="🦮", layout="wide", ) set_sidebar(css) set_intro(css) # Radio button selection analysis_type = st.radio( "Choose analysis type:", options=["Retriever", "Entity Linking"], index=0 # Default to 'question' ) selection_options = ["DB Intervention", "DB Outcome", "Taxonomy Intervention", "Taxonomy Outcome", "Top-k DB in Taxonomy Intervention", "Top-k DB in Taxonmy Outcome", ] if analysis_type == "Retriever": # Selection list using selectbox selection_list = st.selectbox( "Select an option:", options=selection_options ) # text input text = st.text_area( "Enter Text Below:", value="", height=200, max_chars=1500, ) with stylable_container( key="annotate_button", css_styles=""" button { background-color: #a8ebff; color: black; border-radius: 25px; } """, ): submit = st.button("Annotate") # submit = st.button("Run") if "relik_model" not in st.session_state.keys(): st.session_state["relik_model"] = load_model() relik_model = st.session_state["relik_model"][0] if 'candidates' not in st.session_state: st.session_state['candidates'] = [] if 'selected_candidates' not in st.session_state: st.session_state['selected_candidates'] = [] # ReLik API call if submit: if analysis_type == "Entity Linking": relik_model = st.session_state["relik_model"][0] else: model_idx = selection_options.index(selection_list) if selection_list == "Top-k DB in Taxonomy Intervention" or selection_list == "Top-k DB in Taxonmy Outcome": relik_model = st.session_state["relik_model"][model_idx-1] else: relik_model = st.session_state["relik_model"][model_idx+1] text = text.strip() if text: st.markdown("####") with st.spinner(text="In progress"): if analysis_type == "Entity Linking": response = relik_model(text) dict_of_ents, options = get_el_annotations(response=response) dict_of_ents_candidates, options_candidates = get_retriever_annotations(response=response) st.markdown("#### Entity Linking") 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) candidate_text = "".join(f"
  • Intervention: {candidate}
  • " if io_map[candidate] == "intervention" else f"
  • Outcome: {candidate}
  • " for candidate in dict_of_ents_candidates["ents"][0:10]) text = """

    Possible Candidates:

    " st.markdown(text, unsafe_allow_html=True) else: if selection_list == "Top-k DB in Taxonomy Intervention" or selection_list == "Top-k DB in Taxonomy Outcome": response = relik_model.retrieve(text, k=50, batch_size=400, progress_bar=False) candidates_text = [pred.document.text for pred in response[0] if pred.document.text in db_set] else: response = relik_model.retrieve(text, k=10, batch_size=400, progress_bar=False) candidates_text = [pred.document.text for pred in response[0]] if candidates_text: st.session_state.candidates = candidates_text[:10] else: st.session_state.candidates = [] st.session_state.selected_candidates = [] st.markdown("

    No Candidates Found

    ", unsafe_allow_html=True) else: st.error("Please enter some text.") # Ensure the candidates list is displayed even after interactions if st.session_state.candidates and analysis_type != "Entity Linking": dict_of_ents_candidates, options_candidates = get_retriever_annotations_candidates(text, st.session_state.candidates) st.markdown("

    Possible Candidates:

    ", unsafe_allow_html=True) for candidate in dict_of_ents_candidates["ents"]: checked = candidate in st.session_state.selected_candidates if st.checkbox(candidate, key=candidate, value=checked): if candidate not in st.session_state.selected_candidates: st.session_state.selected_candidates.append(candidate) else: if candidate in st.session_state.selected_candidates: st.session_state.selected_candidates.remove(candidate) if st.button("Save Selected Candidates"): if write_candidates_to_file(text, dict_of_ents_candidates["ents"], st.session_state.selected_candidates): st.success("Selected candidates have been saved to file.") if __name__ == "__main__": run_client()