Spaces:
Runtime error
Runtime error
| import os | |
| from typing import List, Optional, Union | |
| import gradio as gr | |
| import spacy | |
| from spacy.tokens import Doc, Span | |
| from relik import Relik | |
| from relik.inference.data.objects import TaskType, RelikOutput | |
| from pyvis.network import Network | |
| # RELIK Models Setup | |
| def setup_relik_model(model_name: str, device: str): | |
| return Relik.from_pretrained(model_name, device=device) | |
| relik_models = { | |
| "sapienzanlp/relik-entity-linking-large": setup_relik_model("sapienzanlp/relik-entity-linking-large", "cuda"), | |
| "relik-ie/relik-relation-extraction-small": setup_relik_model("relik-ie/relik-relation-extraction-small", "cuda") | |
| } | |
| def get_span_annotations(response, doc): | |
| spans = [] | |
| for span in response.spans: | |
| spans.append(Span(doc, span.start, span.end, span.label)) | |
| colors = {span.label_: '#ff5733' for span in spans} # Simple fixed color for demonstration | |
| return spans, colors | |
| def generate_graph(spans, response, colors): | |
| g = Network(width="720px", height="600px", directed=True) | |
| for ent in spans: | |
| g.add_node(ent.text, label=ent.text, color=colors[ent.label_], size=15) | |
| seen_rels = set() | |
| for rel in response.triplets: | |
| if (rel.subject.text, rel.object.text, rel.label) in seen_rels: | |
| continue | |
| g.add_edge(rel.subject.text, rel.object.text, label=rel.label) | |
| seen_rels.add((rel.subject.text, rel.object.text, rel.label)) | |
| html = g.generate_html() | |
| return f"""<iframe style="width: 100%; height: 600px;margin:0 auto" srcdoc='{html.replace("'", '"')}'></iframe>""" | |
| def text_analysis(Text, Model, Relation_Threshold, Window_Size, Window_Stride): | |
| if Model not in relik_models: | |
| raise ValueError(f"Model {Model} not found.") | |
| relik = relik_models[Model] | |
| nlp = spacy.blank("xx") | |
| annotated_text = relik(Text, annotation_type="word", 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 = get_span_annotations(annotated_text, doc) | |
| doc.spans["sc"] = spans | |
| display_el = spacy.displacy.render(doc, style="span", options={"colors": colors}).replace("\n", " ") | |
| display_el = display_el.replace("border-radius: 0.35em;", "border-radius: 0.35em; white-space: nowrap;").replace("span style", "span id='el' style") | |
| display_re = generate_graph(spans, annotated_text, colors) if annotated_text.triplets else "" | |
| return display_el, display_re | |
| theme = gr.themes.Base(primary_hue="rose", secondary_hue="rose", text_size="lg") | |
| 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: | |
| gr.Markdown("# ReLiK with P-FAF Integration") | |
| gr.Interface( | |
| text_analysis, | |
| [ | |
| gr.Textbox(label="Input Text", placeholder="Enter sentence here..."), | |
| gr.Dropdown(list(relik_models.keys()), value="sapienzanlp/relik-entity-linking-large", label="Relik Model"), | |
| gr.Slider(minimum=0, maximum=1, step=0.05, value=0.5, label="Relation Threshold"), | |
| gr.Slider(minimum=16, maximum=128, step=16, value=32, label="Window Size"), | |
| gr.Slider(minimum=8, maximum=64, step=8, value=16, label="Window Stride") | |
| ], | |
| [gr.HTML(label="Entities"), gr.HTML(label="Relations")], | |
| examples=[ | |
| ["Michael Jordan was one of the best players in the NBA."], | |
| ["Noam Chomsky is a renowned linguist and cognitive scientist."] | |
| ], | |
| allow_flagging="never" | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch() |