fix relation highligthing

app.py

@@ -1,4 +1,5 @@
 import json
+from collections import defaultdict
 from typing import Tuple
 
 import gradio as gr
@@ -10,6 +11,7 @@ from pytorch_ie.documents import TextDocumentWithLabeledSpansBinaryRelationsAndL
 from pytorch_ie.models import *  # noqa: F403
 from pytorch_ie.taskmodules import *  # noqa: F403
 
+
 def render_pretty_table(
     document: TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions, **render_kwargs
 ):
@@ -68,9 +70,16 @@ def inject_relation_data(html: str, sorted_entities, binary_relations) -> str:
     # Parse the HTML using BeautifulSoup
     soup = BeautifulSoup(html, "html.parser")
 
+    entity2tails = defaultdict(list)
+    entity2heads = defaultdict(list)
+    for relation in binary_relations:
+        entity2heads[relation.tail].append((relation.head, relation.label))
+        entity2tails[relation.head].append((relation.tail, relation.label))
+
+    entity2id = {entity: f"entity-{idx}" for idx, entity in enumerate(sorted_entities)}
+
     # 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"] = (
@@ -80,77 +89,18 @@ def inject_relation_data(html: str, sorted_entities, binary_relations) -> str:
         # 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"))
+        entity["data-relation-tails"] = json.dumps(
+            [
+                {"entity-id": entity2id[tail], "label": label}
+                for tail, label in entity2tails.get(entity_annotation, [])
+            ]
+        )
+        entity["data-relation-heads"] = json.dumps(
+            [
+                {"entity-id": entity2id[head], "label": label}
+                for head, label in entity2heads.get(entity_annotation, [])
+            ]
+        )
 
     # Return the modified HTML as a string
     return str(soup)
@@ -264,5 +214,54 @@ if __name__ == "__main__":
         )
         render_btn.click(**render_button_kwargs, api_name="render")
 
-    demo.launch()
+        js = """
+        () => {
+            function highlightRelations(entityId) {
+                const entities = document.querySelectorAll('.entity');
+                entities.forEach(entity => {
+                    entity.style.backgroundColor = entity.getAttribute('data-original-color');
+                    entity.style.color = '';
+                });
+
+                if (entityId !== null) {
+                    const selectedEntity = document.getElementById(entityId);
+                    if (selectedEntity) {
+                        selectedEntity.style.backgroundColor = '#ffa';
+                        selectedEntity.style.color = '#000';
+                    }
+                    // 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';
+                        }
+                    });
+                    // highlight heads
+                    const relationHeadsAndLabels = JSON.parse(selectedEntity.getAttribute('data-relation-heads'));
+                    relationHeadsAndLabels.forEach(relationHead => {
+                        const headEntity = document.getElementById(relationHead['entity-id']);
+                        if (headEntity) {
+                            headEntity.style.backgroundColor = '#faf';
+                            headEntity.style.color = '#000';
+                        }
+                    });
+                }
+            }
+
+            const entities = document.querySelectorAll('.entity');
+            entities.forEach(entity => {
+                entity.addEventListener('mouseover', () => {
+                    highlightRelations(entity.id);
+                });
+                entity.addEventListener('mouseout', () => {
+                    highlightRelations(null);
+                });
+            });
+        }
+        """
 
+        rendered_output.change(fn=None, js=js, inputs=[], outputs=[])
+
+    demo.launch()