Update app.py

app.py

@@ -1,27 +1,23 @@
+import json
+
 import gradio as gr
 from pie_modules.models import *  # noqa: F403
 from pie_modules.taskmodules import *  # noqa: F403
-from prettytable import PrettyTable
 from pytorch_ie.annotations import LabeledSpan
 from pytorch_ie.auto import AutoPipeline
 from pytorch_ie.documents import TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions
 from pytorch_ie.models import *  # noqa: F403
 from pytorch_ie.taskmodules import *  # noqa: F403
 
-
-def predict(text):
-    document = TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions(text=text)
-
-    # add single partition from the whole text (the model only considers text in partitions)
-    document.labeled_partitions.append(LabeledSpan(start=0, end=len(text), label="text"))
-
-    # execute NER pipeline
-    pipeline(document)
+def render_pretty_table(
+    document: TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions, **render_kwargs
+):
+    from prettytable import PrettyTable
 
     t = PrettyTable()
     t.field_names = ["head", "tail", "relation"]
     t.align = "l"
-    for relation in document.binary_relations.predictions:
+    for relation in list(document.binary_relations) + list(document.binary_relations.predictions):
         t.add_row([str(relation.head), str(relation.tail), relation.label])
 
     html = t.get_html_string(format=True)
@@ -30,9 +26,159 @@ def predict(text):
     return html
 
 
+def render_spacy(
+    document: TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions,
+    style="ent",
+    inject_relations=True,
+    **render_kwargs,
+):
+    from spacy import displacy
+
+    spans = list(document.labeled_spans) + list(document.labeled_spans.predictions)
+    spacy_doc = {
+        "text": document.text,
+        "ents": [
+            {"start": entity.start, "end": entity.end, "label": entity.label} for entity in spans
+        ],
+        "title": None,
+    }
+
+    html = displacy.render(
+        spacy_doc, page=True, manual=True, minify=True, style=style, **render_kwargs
+    )
+    html = "<div style='max-width:100%; max-height:360px; overflow:auto'>" + html + "</div>"
+    if inject_relations:
+        print("Injecting relation data")
+        binary_relations = list(document.binary_relations) + list(
+            document.binary_relations.predictions
+        )
+        sorted_entities = sorted(spans, key=lambda x: (x.start, x.end))
+        html = inject_relation_data(
+            html, sorted_entities=sorted_entities, binary_relations=binary_relations
+        )
+    else:
+        print("Not injecting relation data")
+    return html
+
+
+def inject_relation_data(html: str, sorted_entities, binary_relations) -> str:
+    from bs4 import BeautifulSoup
+
+    # Parse the HTML using BeautifulSoup
+    soup = BeautifulSoup(html, "html.parser")
+
+    # Add unique IDs to each entity
+    entities = soup.find_all(class_="entity")
+    entity2id = {}
+    for idx, entity in enumerate(entities):
+        entity["id"] = f"entity-{idx}"
+        entity["data-original-color"] = (
+            entity["style"].split("background:")[1].split(";")[0].strip()
+        )
+        entity_annotation = sorted_entities[idx]
+        # sanity check
+        if str(entity_annotation) != entity.next:
+            raise ValueError(f"Entity text mismatch: {entity_annotation} != {entity.text}")
+        entity2id[sorted_entities[idx]] = f"entity-{idx}"
+
+    # Generate prefixed relations
+    prefixed_relations = [
+        {
+            "head": entity2id[relation.head],
+            "tail": entity2id[relation.tail],
+            "label": relation.label,
+        }
+        for relation in binary_relations
+    ]
+
+    # Create the JavaScript function to handle mouse over and mouse out events
+    script = (
+        """
+    <script>
+    function highlightRelations(entityId, relations) {
+        // Reset all entities' styles
+        const entities = document.querySelectorAll('.entity');
+        entities.forEach(entity => {
+            entity.style.backgroundColor = entity.getAttribute('data-original-color');
+            entity.style.color = '';
+        });
+
+        // If an entity is hovered, highlight it and its related entities with different colors
+        if (entityId !== null) {
+            const selectedEntity = document.getElementById(entityId);
+            if (selectedEntity) {
+                selectedEntity.style.backgroundColor = '#ffa';
+                selectedEntity.style.color = '#000';
+            }
+
+            relations.forEach(relation => {
+                if (relation.head === entityId) {
+                    const tailEntity = document.getElementById(relation.tail);
+                    if (tailEntity) {
+                        tailEntity.style.backgroundColor = '#aff';
+                        tailEntity.style.color = '#000';
+                    }
+                }
+                if (relation.tail === entityId) {
+                    const headEntity = document.getElementById(relation.head);
+                    if (headEntity) {
+                        headEntity.style.backgroundColor = '#faf';
+                        headEntity.style.color = '#000';
+                    }
+                }
+            });
+        }
+    }
+
+    // Event listeners for mouse over and mouse out on each entity
+    document.addEventListener('DOMContentLoaded', (event) => {
+        const relations = %s;
+        const entities = document.querySelectorAll('.entity');
+        entities.forEach(entity => {
+            entity.addEventListener('mouseover', () => {
+                highlightRelations(entity.id, relations);
+            });
+            entity.addEventListener('mouseout', () => {
+                highlightRelations(null, relations);
+            });
+        });
+    });
+    </script>
+    """
+        % prefixed_relations
+    )
+
+    # Inject the script into the HTML
+    soup.body.append(BeautifulSoup(script, "html.parser"))
+
+    # Return the modified HTML as a string
+    return str(soup)
+
+
+def predict(text, render_as, render_kwargs_json):
+    document = TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions(text=text)
+
+    # add single partition from the whole text (the model only considers text in partitions)
+    document.labeled_partitions.append(LabeledSpan(start=0, end=len(text), label="text"))
+
+    # execute prediction pipeline
+    pipeline(document)
+
+    render_kwargs = json.loads(render_kwargs_json)
+    if render_as == "Pretty Table":
+        html = render_pretty_table(document, **render_kwargs)
+    elif render_as == "spaCy":
+        html = render_spacy(document, **render_kwargs)
+    else:
+        raise ValueError(f"Unknown render_as value: {render_as}")
+
+    return html
+
+
 if __name__ == "__main__":
 
     model_name_or_path = "ArneBinder/sam-pointer-bart-base-v0.3"
+    # local path
     # model_name_or_path = "models/dataset-sciarg/task-ner_re/v0.3/2024-03-01_18-25-32"
 
     pipeline = AutoPipeline.from_pretrained(model_name_or_path, device=-1, num_workers=0)
@@ -43,14 +189,34 @@ if __name__ == "__main__":
         # taskmodule_kwargs=dict(create_relation_candidates=True),
     )
 
+    default_render_kwargs = {
+        "style": "ent",
+        "options": {
+            "colors": {"own_claim": "#009933", "background_claim": "#0033cc", "data": "#993399"}
+        },
+    }
+
     iface = gr.Interface(
         fn=predict,
         inputs=[
             gr.Textbox(
                 lines=20,
                 value="Scholarly Argumentation Mining (SAM) has recently gained attention due to its potential to help scholars with the rapid growth of published scientific literature. It comprises two subtasks: argumentative discourse unit recognition (ADUR) and argumentative relation extraction (ARE), both of which are challenging since they require e.g. the integration of domain knowledge, the detection of implicit statements, and the disambiguation of argument structure. While previous work focused on dataset construction and baseline methods for specific document sections, such as abstract or results, full-text scholarly argumentation mining has seen little progress. In this work, we introduce a sequential pipeline model combining ADUR and ARE for full-text SAM, and provide a first analysis of the performance of pretrained language models (PLMs) on both subtasks. We establish a new SotA for ADUR on the Sci-Arg corpus, outperforming the previous best reported result by a large margin (+7% F1). We also present the first results for ARE, and thus for the full AM pipeline, on this benchmark dataset. Our detailed error analysis reveals that non-contiguous ADUs as well as the interpretation of discourse connectors pose major challenges and that data annotation needs to be more consistent.",
-            )
+            ),
+        ],
+        additional_inputs=[
+            gr.Dropdown(
+                label="Render as",
+                choices=["Pretty Table", "spaCy"],
+                value="spaCy",
+            ),
+            gr.Textbox(
+                label="Render Arguments",
+                lines=5,
+                value=json.dumps(default_render_kwargs, indent=2),
+            ),
         ],
+        additional_inputs_accordion=gr.Accordion(label="Render Options", open=False),
         outputs=["html"],
     )
     iface.launch()