Adjusted demo notebook

ssciwr · Feb 16, 2025 · 2c28fc3 · 2c28fc3
1 parent 556217a
commit 2c28fc3
Showing 1 changed file with 16 additions and 13 deletions.
diff --git a/notebook/demo.ipynb b/notebook/demo.ipynb
@@ -44,25 +44,28 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "# a dictionary matching colors to the different entity types\n",
+    "colors = {\n",
+    "    \"LOC\": \"green\",\n",
+    "    \"ORG\": \"blue\",\n",
+    "    \"MISC\": \"yellow\",\n",
+    "    \"PER\": \"red\"\n",
+    "}\n",
+    "\n",
     "# function for displaying the result using HTML\n",
-    "def highlight_ne(text, per_list, org_list, loc_list, misc_list):\n",
+    "def highlight_ne(text, ne_list):\n",
     "    # create a list of all entities with their positions\n",
     "    entities = []\n",
-    "    for loc in loc_list:\n",
-    "        entities.append((loc, \"green\"))\n",
-    "    for org in org_list:\n",
-    "        entities.append((org, \"blue\"))\n",
-    "    for misc in misc_list:\n",
-    "        entities.append((misc, \"yellow\"))\n",
-    "    for per in per_list:\n",
-    "        entities.append((per, \"red\"))\n",
+    "    for ne in ne_list:\n",
+    "        entities.append((ne, colors.get(ne[\"entity_group\"])))\n",
     "    \n",
     "    # sort entities by their positions in the text in reverse order\n",
-    "    entities.sort(key=lambda x: text.find(x[0]), reverse=True)\n",
+    "    entities = sorted(entities, key=lambda x: x[0][\"start\"], reverse=True)\n",
     "    \n",
     "    # replace entities with highlighted spans\n",
     "    for entity, color in entities:\n",
-    "        text = text.replace(entity, f\"<span style=\\\"background-color:{color}\\\">{entity}</span>\")\n",
+    "        ent_word = entity[\"word\"]\n",
+    "        text = text.replace(ent_word, f\"<span style=\\\"background-color:{color}\\\">{ent_word}</span>\")\n",
     "    \n",
     "    return text"
    ]
@@ -133,7 +136,7 @@
     "    output_text = ps.pseudonymize(html_text)\n",
     "\n",
     "    # display original text and highlight found and replaced NEs\n",
-    "    highlighted_html = highlight_ne(html_text, ps.per_list, ps.org_list, ps.loc_list, ps.misc_list)\n",
+    "    highlighted_html = highlight_ne(html_text, ps.ne_list)\n",
     "    display(HTML(highlighted_html))\n",
     "\n",
     "    # add pseudonymized text to dict\n",
@@ -202,7 +205,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.10"
+   "version": "3.11.11"
   }
  },
  "nbformat": 4,