Update app.py

app.py

@@ -221,7 +221,7 @@ relik_models = {
 }
 
 
-def get_span_annotations(response, doc):
+def get_span_annotations(response, doc, ner=False):
     dict_ents = {}
     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' style='color: #414141'></i> {}</a>"
@@ -236,7 +236,7 @@ def get_span_annotations(response, doc):
                 span.end,
                 el_link_wrapper.format(
                     span.label.replace(" ", "_"), span.label
-                ) if span.label != "--NME--" else "--NME--",
+                ) if span.label != "--NME--" and not ner else span.label,
                 # kb_id=span.label.replace(" ", "_")
             )
         )
@@ -249,7 +249,7 @@ def get_span_annotations(response, doc):
     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"):
+def generate_graph(spans, response, colors, dict_ents, bgcolor="#111827", font_color="white", ner=False):
     g = Network(
         width="720px",
         height="600px",
@@ -269,7 +269,7 @@ def generate_graph(spans, response, colors, dict_ents, bgcolor="#111827", font_c
 
     for ent in spans:
         # if not NME use title:
-        if dict_ents[(ent.start, ent.end)][2] != "--NME--":
+        if dict_ents[(ent.start, ent.end)][2] != "--NME--" or not ner:
             g.add_node(
                 dict_ents[(ent.start, ent.end)][2],
                 label=dict_ents[(ent.start, ent.end)][2],
@@ -289,22 +289,23 @@ def generate_graph(spans, response, colors, dict_ents, bgcolor="#111827", font_c
             )
     seen_rels = set()
     for rel in response.triplets:
-        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
+        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--" 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--" else dict_ents[(rel.object.start, rel.object.end)][1],
+            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,
         )
@@ -344,7 +345,7 @@ def text_analysis(Text, Model, Relation_Threshold, Window_Size, Window_Stride):
     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)
+    spans, colors, dict_ents = get_span_annotations(response=annotated_text, doc=doc, ner="ner" in Model)
     doc.spans["sc"] = spans
 
     # build the EL display
@@ -362,7 +363,7 @@ def text_analysis(Text, Model, Relation_Threshold, Window_Size, Window_Stride):
     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)
+        display_re = generate_graph(spans, annotated_text, colors, dict_ents, ner="ner" in Model)
     return display_el, display_re