Upload 3 files

app.py

@@ -1,75 +1,21 @@
 import json
-from collections import defaultdict
-from typing import Dict, List, Optional, Tuple, Union
+from typing import Tuple
 
 import gradio as gr
 from pie_modules.models import *  # noqa: F403
 from pie_modules.taskmodules import *  # noqa: F403
-from pytorch_ie.annotations import BinaryRelation, LabeledSpan
+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
 
-from .rendering_utils import render_pretty_table, render_spacy
-
+from rendering_utils import render_pretty_table, render_spacy
 
 RENDER_WITH_DISPLACY = "displaCy + highlighted arguments"
 RENDER_WITH_PRETTY_TABLE = "Pretty Table"
 
 
-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
-    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")
-    for idx, entity in enumerate(entities):
-        entity["id"] = f"entity-{idx}"
-        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}
-                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)
-
-
 def predict(text: str) -> Tuple[dict, str]:
     document = TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions(text=text)
 

rendering_utils.py

@@ -5,7 +5,7 @@ from typing import Dict, List, Optional, Union
 from pytorch_ie.annotations import BinaryRelation
 from pytorch_ie.documents import TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions
 
-from .rendering_utils_displacy import EntityRenderer
+from rendering_utils_displacy import EntityRenderer
 
 
 def render_pretty_table(