import spaces
from functools import partial
import os
from typing import Any, Dict, List, Tuple, Optional, Union
import gradio as gr
import spacy
from pyvis.network import Network
from spacy import displacy
from spacy.tokens import Doc, Span
from relik.common.utils import CONFIG_NAME, from_cache
from relik.inference.annotator import Relik
from relik.inference.serve.frontend.utils import get_random_color
from relik.retriever.pytorch_modules.model import GoldenRetriever
from relik.retriever.indexers.inmemory import InMemoryDocumentIndex
from relik.inference.data.objects import TaskType
from relik.retriever.pytorch_modules import RetrievedSample
from relik.retriever.indexers.document import Document, DocumentStore
from relik.retriever.indexers.base import BaseDocumentIndex
LOGO = """
"""
DESCRIPTION = """
Retrieve, Read and LinK: Fast and Accurate Entity Linking and Relation Extraction on an Academic Budget
A blazing fast and lightweight Information Extraction model for Entity Linking and Relation Extraction.
Riccardo Orlando,
Pere-Lluís Huguet Cabot,
Edoardo Barba,
and Roberto Navigli
"""
INSTRUCTION = """
## Use it locally
Installation from PyPI
```bash
pip install relik
```
ReLiK is a lightweight and fast model for **Entity Linking** and **Relation Extraction**.
It is composed of two main components: a **retriever** and a **reader**.
The retriever is responsible for retrieving relevant documents from a large collection of documents,
while the reader is responsible for extracting entities and relations from the retrieved documents.
ReLiK can be used with the `from_pretrained` method to load a pre-trained pipeline.
Here is an example of how to use ReLiK for Entity Linking:
```python
from relik import Relik
from relik.inference.data.objects import RelikOutput
relik = Relik.from_pretrained("sapienzanlp/relik-entity-linking-large")
relik_out: RelikOutput = relik("Michael Jordan was one of the best players in the NBA.")
# RelikOutput(
# text="Michael Jordan was one of the best players in the NBA.",
# tokens=['Michael', 'Jordan', 'was', 'one', 'of', 'the', 'best', 'players', 'in', 'the', 'NBA', '.'],
# id=0,
# spans=[
# Span(start=0, end=14, label="Michael Jordan", text="Michael Jordan"),
# Span(start=50, end=53, label="National Basketball Association", text="NBA"),
# ],
# triplets=[],
# candidates=Candidates(
# span=[
# [
# [
# {"text": "Michael Jordan", "id": 4484083},
# {"text": "National Basketball Association", "id": 5209815},
# {"text": "Walter Jordan", "id": 2340190},
# {"text": "Jordan", "id": 3486773},
# {"text": "50 Greatest Players in NBA History", "id": 1742909},
# ...
# ]
# ]
# ]
# ),
# )
```
and for Relation Extraction:
```python
from relik import Relik
from relik.inference.data.objects import RelikOutput
relik = Relik.from_pretrained("sapienzanlp/relik-relation-extraction-large")
relik_out: RelikOutput = relik("Michael Jordan was one of the best players in the NBA.")
```
For more information, please refer to the [source code](https://github.com/SapienzaNLP/relik/).
"""
class GoldenSillyRetriever(GoldenRetriever):
def __init__(self, documents: List[str], *args, **kwargs):
self.documents = DocumentStore([Document(doc) for doc in documents])
self.document_index = BaseDocumentIndex(self.documents)
def retrieve(self,
text: Optional[Union[str, List[str]]] = None,
k: int = 100,
*args,
**kwargs,
) -> List[List[RetrievedSample]]:
if isinstance(text, str):
text = [text]
elif text is None:
text = []
return [
[RetrievedSample(score=1.0, document=doc) for doc in self.documents[:k]]
for _ in text
]
def index(self):
pass
def eval(self):
pass
def save_pretrained(self):
pass
def to(self, device):
pass
wikipedia_retriever = GoldenRetriever("relik-ie/encoder-e5-base-v2-wikipedia", device="cuda")
wikipedia_index = InMemoryDocumentIndex.from_pretrained("relik-ie/encoder-e5-base-v2-wikipedia-index", index_precision="bf16", device="cuda")
wikidata_retriever = GoldenRetriever("relik-ie/encoder-e5-small-v2-wikipedia-relations", device="cuda")
wikidata_index = InMemoryDocumentIndex.from_pretrained("relik-ie/encoder-e5-small-v2-wikipedia-relations-index", index_precision="bf16", device="cuda")
ner_type_retriever = GoldenSillyRetriever(
documents=['media', 'disease', 'miscellaneous', 'event', 'person', 'location', 'time', 'celestial', 'organization', 'concept']
)
relik_available_models = [
"relik-ie/relik-cie-xl",
"relik-ie/relik-cie-small",
"sapienzanlp/relik-entity-linking-large",
"relik-ie/relik-entity-linking-large-robust",
"relik-ie/relik-relation-extraction-small",
"relik-ie/relik-relation-extraction-large",
"relik-ie/relik-relation-extraction-small-wikipedia-ner",
]
relik_models = {
"sapienzanlp/relik-entity-linking-large": Relik.from_pretrained(
"sapienzanlp/relik-entity-linking-large",
device="cuda",
index={
TaskType.SPAN: wikipedia_index,
},
retriever={
TaskType.SPAN:wikipedia_retriever,
},
reader_kwargs={"dataset_kwargs": {"use_nme": True}},
),
"relik-ie/relik-cie-small": Relik.from_pretrained(
"relik-ie/relik-cie-small",
device="cuda",
index={
TaskType.SPAN: wikipedia_index,
TaskType.TRIPLET: wikidata_index,
},
retriever={
TaskType.SPAN: wikipedia_retriever,
TaskType.TRIPLET: wikidata_retriever,
},
reader_kwargs={"dataset_kwargs": {"use_nme": True}},
),
"relik-ie/relik-cie-xl": Relik.from_pretrained(
"relik-ie/relik-cie-xl",
device="cuda",
index={
TaskType.SPAN: wikipedia_index,
TaskType.TRIPLET: wikidata_index,
},
retriever={
TaskType.SPAN: wikipedia_retriever,
TaskType.TRIPLET: wikidata_retriever,
},
reader_kwargs={"dataset_kwargs": {"use_nme": True}},
),
"relik-ie/relik-relation-extraction-small-wikipedia-ner": Relik.from_pretrained(
"relik-ie/relik-relation-extraction-small-wikipedia-ner",
device="cuda",
use_nme=True,
retriever={
TaskType.SPAN: ner_type_retriever,
TaskType.TRIPLET: wikidata_retriever,
},
index={
TaskType.SPAN: ner_type_retriever.document_index,
TaskType.TRIPLET: wikidata_index,
}
),
"relik-ie/relik-relation-extraction-small": Relik.from_pretrained(
"relik-ie/relik-relation-extraction-small",
index={
TaskType.TRIPLET:wikidata_index,
},
device="cuda",
retriever={
TaskType.TRIPLET: wikidata_retriever,
},
),
"relik-ie/relik-relation-extraction-large": Relik.from_pretrained(
"relik-ie/relik-relation-extraction-large",
index={
TaskType.TRIPLET:wikidata_index,
},
device="cuda",
retriever={
TaskType.TRIPLET: wikidata_retriever,
},
),
"relik-ie/relik-entity-linking-large-robust": Relik.from_pretrained(
"relik-ie/relik-entity-linking-large-robust",
index={
TaskType.SPAN: wikipedia_index,
},
device="cuda",
retriever={
TaskType.SPAN: wikipedia_retriever,
},
reader_kwargs={"dataset_kwargs": {"use_nme": True}},
),
}
def get_span_annotations(response, doc, ner=False):
dict_ents = {}
el_link_wrapper = (
" {}"
)
spans = []
for idx, span in enumerate(response.spans):
spans.append(
Span(
doc,
span.start,
span.end,
el_link_wrapper.format(
span.label.replace(" ", "_"), span.label
) if (span.label != "--NME--" and not ner) else span.label,
# kb_id=span.label.replace(" ", "_")
)
)
dict_ents[(span.start, span.end)] = (
span.label + str(idx),
doc[span.start : span.end].text,
span.label,
span.label.replace(" ", "_"),
)
colors = get_random_color(set([span.label_ for span in spans]))
return spans, colors, dict_ents
def generate_graph(spans, response, colors, dict_ents, bgcolor="#111827", font_color="white", ner=False):
g = Network(
width="720px",
height="600px",
directed=True,
notebook=False,
bgcolor=bgcolor,
font_color=font_color,
)
g.barnes_hut(
gravity=-3000,
central_gravity=0.3,
spring_length=50,
spring_strength=0.001,
damping=0.09,
overlap=0,
)
for ent in spans:
# if not NME use title:
if dict_ents[(ent.start, ent.end)][2] != "--NME--" and not ner:
g.add_node(
dict_ents[(ent.start, ent.end)][2],
label=dict_ents[(ent.start, ent.end)][2],
color=colors[ent.label_],
title=dict_ents[(ent.start, ent.end)][2],
size=15,
labelHighlightBold=True,
)
else:
g.add_node(
ent.text,
label=ent.text,
color=colors[ent.label_],
title=ent.text,
size=15,
labelHighlightBold=True,
)
seen_rels = set()
for rel in response.triplets:
if not ner:
if dict_ents[(rel.subject.start, rel.subject.end)][2] == "--NME--" and dict_ents[(rel.object.start, rel.object.end)][2] == "--NME--":
if (rel.subject.text, rel.object.text, rel.label) in seen_rels:
continue
elif dict_ents[(rel.subject.start, rel.subject.end)][2] == "--NME--" and dict_ents[(rel.object.start, rel.object.end)][2] != "--NME--":
if (rel.subject.text, dict_ents[(rel.object.start, rel.object.end)][2], rel.label) in seen_rels:
continue
elif dict_ents[(rel.subject.start, rel.subject.end)][2] != "--NME--" and dict_ents[(rel.object.start, rel.object.end)][2] == "--NME--":
if (dict_ents[(rel.subject.start, rel.subject.end)][2], rel.object.text, rel.label) in seen_rels:
continue
else:
if (dict_ents[(rel.subject.start, rel.subject.end)][2], dict_ents[(rel.object.start, rel.object.end)][2], rel.label) in seen_rels:
continue
g.add_edge(
dict_ents[(rel.subject.start, rel.subject.end)][2] if dict_ents[(rel.subject.start, rel.subject.end)][2] != "--NME--" and not ner else dict_ents[(rel.subject.start, rel.subject.end)][1],
dict_ents[(rel.object.start, rel.object.end)][2] if dict_ents[(rel.object.start, rel.object.end)][2] != "--NME--" and not ner else dict_ents[(rel.object.start, rel.object.end)][1],
label=rel.label,
title=rel.label,
)
if dict_ents[(rel.subject.start, rel.subject.end)][2] != "--NME--" and dict_ents[(rel.object.start, rel.object.end)][2] != "--NME--":
seen_rels.add((dict_ents[(rel.subject.start, rel.subject.end)][2], dict_ents[(rel.object.start, rel.object.end)][2], rel.label))
elif dict_ents[(rel.subject.start, rel.subject.end)][2] != "--NME--" and dict_ents[(rel.object.start, rel.object.end)][2] == "--NME--":
seen_rels.add((dict_ents[(rel.subject.start, rel.subject.end)][2], rel.object.text, rel.label))
elif dict_ents[(rel.subject.start, rel.subject.end)][2] == "--NME--" and dict_ents[(rel.object.start, rel.object.end)][2] != "--NME--":
seen_rels.add((rel.subject.text, dict_ents[(rel.object.start, rel.object.end)][2], rel.label))
else:
seen_rels.add((rel.subject.text, rel.object.text, rel.label))
# g.show(filename, notebook=False)
html = g.generate_html()
# need to remove ' from HTML
html = html.replace("'", '"')
return f""""""
@spaces.GPU
def text_analysis(Text, Model, Relation_Threshold, Window_Size, Window_Stride):
global loaded_model
if Model is None:
return "", ""
# if loaded_model is None or loaded_model["key"] != Model:
# relik = Relik.from_pretrained(Model, index_precision="bf16")
# loaded_model = {"key": Model, "model": relik}
# else:
# relik = loaded_model["model"]
if Model not in relik_models:
raise ValueError(f"Model {Model} not found.")
relik = relik_models[Model]
# spacy for span visualization
nlp = spacy.blank("xx")
annotated_text = relik(Text, annotation_type="word", num_workers=0, remove_nmes= False, relation_threshold = Relation_Threshold, window_size=Window_Size, window_stride=Window_Stride)
doc = Doc(nlp.vocab, words=[token.text for token in annotated_text.tokens])
spans, colors, dict_ents = get_span_annotations(response=annotated_text, doc=doc, ner="ner" in Model)
doc.spans["sc"] = spans
# build the EL display
display_el = displacy.render(doc, style="span", options={"colors": colors})#, "kb_url_template": "https://en.wikipedia.org/wiki/{}"})
display_el = display_el.replace("\n", " ")
# heuristic, prevents split of annotation decorations
display_el = display_el.replace(
"border-radius: 0.35em;",
"border-radius: 0.35em; white-space: nowrap;",
)
display_el = display_el.replace(
"span style",
"span id='el' style",
)
display_re = ""
if annotated_text.triplets:
# background_color should be the same as the background of the page
display_re = generate_graph(spans, annotated_text, colors, dict_ents, ner="ner" in Model)
return display_el, display_re
theme = theme = gr.themes.Base(
primary_hue="rose",
secondary_hue="rose",
text_size="lg",
# font=[gr.themes.GoogleFont("Montserrat"), "Arial", "sans-serif"],
)
css = """
h1 {
text-align: center;
display: block;
}
mark {
color: black;
}
#el {
white-space: nowrap;
}
"""
with gr.Blocks(fill_height=True, css=css, theme=theme) as demo:
# check if demo is running in dark mode
gr.Markdown(LOGO)
gr.Markdown(DESCRIPTION)
gr.Interface(
text_analysis,
[
gr.Textbox(label="Input Text", placeholder="Enter sentence here..."),
gr.Dropdown(
relik_available_models,
value=relik_available_models[0],
label="Relik Model",
),
gr.Slider(
minimum=0,
maximum=1,
step=0.05,
value=0.5,
label="Relation Threshold",
info="Minimum confidence for relation extraction (Only for RE and cIE)",
),
gr.Slider(
minimum=16,
maximum=128,
step=16,
value=32,
label="Window Size",
info="Window size for the sliding window",
),
gr.Slider(
minimum=8,
maximum=64,
step=8,
value=16,
label="Window Stride",
info="Window stride for the sliding window",
),
],
[gr.HTML(label="Entities"), gr.HTML(label="Relations")],
examples=[
["Avram Noam Chomsky born December 7, 1928) is an American professor and public intellectual known for his work in linguistics, political activism, and social criticism. Sometimes called 'the father of modern linguistics', Chomsky is also a major figure in analytic philosophy and one of the founders of the field of cognitive science. He is a laureate professor of linguistics at the University of Arizona and an institute professor emeritus at the Massachusetts Institute of Technology (MIT). Among the most cited living authors, Chomsky has written more than 150 books on topics such as linguistics, war, and politics. In addition to his work in linguistics, since the 1960s Chomsky has been an influential voice on the American left as a consistent critic of U.S. foreign policy, contemporary capitalism, and corporate influence on political institutions and the media."],
["'Bella ciao' (Italian pronunciation: [ˈbɛlla ˈtʃaːo]; 'Goodbye beautiful') is an Italian song dedicated to the partisans of the Italian resistance, which fought against the occupying troops of Nazi Germany and the collaborationist Fascist forces during the liberation of Italy. It was based on a folk song of the late 19th century, sung by female workers of the paddy fields in Northern Italy (mondine) in protest against harsh working conditions. Versions of 'Bella ciao' continue to be sung worldwide as a hymn of resistance."],
],
allow_flagging="never",
)
gr.Markdown("")
gr.Markdown(INSTRUCTION)
if __name__ == "__main__":
demo.launch()
