Update app.py

app.py

@@ -1,11 +1,11 @@
 import json
 from collections import defaultdict
-from typing import Tuple
+from typing import Dict, List, Optional, Tuple, Union
 
 import gradio as gr
 from pie_modules.models import *  # noqa: F403
 from pie_modules.taskmodules import *  # noqa: F403
-from pytorch_ie.annotations import LabeledSpan
+from pytorch_ie.annotations import BinaryRelation, LabeledSpan
 from pytorch_ie.auto import AutoPipeline
 from pytorch_ie.documents import TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions
 from pytorch_ie.models import *  # noqa: F403
@@ -33,6 +33,7 @@ def render_spacy(
     document: TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions,
     style="ent",
     inject_relations=True,
+    colors_hover=None,
     **render_kwargs,
 ):
     from spacy import displacy
@@ -51,20 +52,25 @@ def render_spacy(
     )
     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
+            html,
+            sorted_entities=sorted_entities,
+            binary_relations=binary_relations,
+            additional_colors=colors_hover,
         )
-    else:
-        print("Not injecting relation data")
     return html
 
 
-def inject_relation_data(html: str, sorted_entities, binary_relations) -> str:
+def inject_relation_data(
+    html: str,
+    sorted_entities,
+    binary_relations: List[BinaryRelation],
+    additional_colors: Optional[Dict[str, Union[str, dict]]] = None,
+) -> str:
     from bs4 import BeautifulSoup
 
     # Parse the HTML using BeautifulSoup
@@ -82,13 +88,18 @@ def inject_relation_data(html: str, sorted_entities, binary_relations) -> str:
     entities = soup.find_all(class_="entity")
     for idx, entity in enumerate(entities):
         entity["id"] = f"entity-{idx}"
-        entity["data-original-color"] = (
-            entity["style"].split("background:")[1].split(";")[0].strip()
-        )
+        original_color = entity["style"].split("background:")[1].split(";")[0].strip()
+        entity["data-color-original"] = original_color
+        if additional_colors is not None:
+            for key, color in additional_colors.items():
+                entity[f"data-color-{key}"] = (
+                    json.dumps(color) if isinstance(color, dict) else color
+                )
         entity_annotation = sorted_entities[idx]
         # sanity check
         if str(entity_annotation) != entity.next:
             raise ValueError(f"Entity text mismatch: {entity_annotation} != {entity.text}")
+        entity["data-label"] = entity_annotation.label
         entity["data-relation-tails"] = json.dumps(
             [
                 {"entity-id": entity2id[tail], "label": label}
@@ -165,10 +176,24 @@ if __name__ == "__main__":
             # we need to convert the keys to uppercase because the spacy rendering function expects them in uppercase
             "colors": {
                 "own_claim".upper(): "#009933",
-                "background_claim".upper(): "#0033cc",
+                "background_claim".upper(): "#99ccff",
                 "data".upper(): "#993399",
             }
         },
+        "colors_hover": {
+            "selected": "#ffa",
+            # "tail": "#aff",
+            "tail": {
+                # green
+                "supports": "#9f9",
+                # red
+                "contradicts": "#f99",
+                # do not highlight
+                "parts_of_same": None,
+            },
+            "head": None,  # "#faf",
+            "other": None,
+        },
     }
 
     with gr.Blocks() as demo:
@@ -216,26 +241,47 @@ if __name__ == "__main__":
 
         js = """
         () => {
-            function highlightRelations(entityId) {
+            function maybeSetColor(entity, colorAttributeKey, colorDictKey) {
+                var color = entity.getAttribute('data-color-' + colorAttributeKey);
+                // if color is a json string, parse it and use the value at colorDictKey
+                try {
+                    const colors = JSON.parse(color);
+                    color = colors[colorDictKey];
+                } catch (e) {}
+                if (color) {
+                    console.log('setting color', color);
+                    console.log('entity', entity);
+                    entity.style.backgroundColor = color;
+                    entity.style.color = '#000';
+                }
+            }
+
+            function highlightRelationArguments(entityId) {
                 const entities = document.querySelectorAll('.entity');
+                // reset all entities
                 entities.forEach(entity => {
-                    entity.style.backgroundColor = entity.getAttribute('data-original-color');
+                    const color = entity.getAttribute('data-color-original');
+                    entity.style.backgroundColor = color;
                     entity.style.color = '';
                 });
 
                 if (entityId !== null) {
+                    var visitedEntities = new Set();
+                    // highlight selected entity
                     const selectedEntity = document.getElementById(entityId);
                     if (selectedEntity) {
-                        selectedEntity.style.backgroundColor = '#ffa';
-                        selectedEntity.style.color = '#000';
+                        const label = selectedEntity.getAttribute('data-label');
+                        maybeSetColor(selectedEntity, 'selected', label);
+                        visitedEntities.add(selectedEntity);
                     }
                     // highlight tails
                     const relationTailsAndLabels = JSON.parse(selectedEntity.getAttribute('data-relation-tails'));
                     relationTailsAndLabels.forEach(relationTail => {
                         const tailEntity = document.getElementById(relationTail['entity-id']);
                         if (tailEntity) {
-                            tailEntity.style.backgroundColor = '#aff';
-                            tailEntity.style.color = '#000';
+                            const label = relationTail['label'];
+                            maybeSetColor(tailEntity, 'tail', label);
+                            visitedEntities.add(tailEntity);
                         }
                     });
                     // highlight heads
@@ -243,8 +289,16 @@ if __name__ == "__main__":
                     relationHeadsAndLabels.forEach(relationHead => {
                         const headEntity = document.getElementById(relationHead['entity-id']);
                         if (headEntity) {
-                            headEntity.style.backgroundColor = '#faf';
-                            headEntity.style.color = '#000';
+                            const label = relationHead['label'];
+                            maybeSetColor(headEntity, 'head', label);
+                            visitedEntities.add(headEntity);
+                        }
+                    });
+                    // highlight other entities
+                    entities.forEach(entity => {
+                        if (!visitedEntities.has(entity)) {
+                            const label = entity.getAttribute('data-label');
+                            maybeSetColor(entity, 'other', label);
                         }
                     });
                 }
@@ -252,12 +306,17 @@ if __name__ == "__main__":
 
             const entities = document.querySelectorAll('.entity');
             entities.forEach(entity => {
+                const alreadyHasListener = entity.getAttribute('data-has-listener');
+                if (alreadyHasListener) {
+                    return;
+                }
                 entity.addEventListener('mouseover', () => {
-                    highlightRelations(entity.id);
+                    highlightRelationArguments(entity.id);
                 });
                 entity.addEventListener('mouseout', () => {
-                    highlightRelations(null);
+                    highlightRelationArguments(null);
                 });
+                entity.setAttribute('data-has-listener', 'true');
             });
         }
         """