Demo notebook NE display using HTML

mailcom/parse.py

@@ -58,6 +58,12 @@ def __init__(self):
         # records the already replaced names in an email
         self.used_first_names = {}
 
+        # records NEs in the last email
+        self.per_list = []
+        self.org_list = []
+        self.loc_list = []
+        self.misc_list = []
+
     def init_spacy(self, language: str, model="default"):
         if model == "default":
             model = self.spacy_default_model_dict[language]
@@ -95,6 +101,11 @@ def init_transformers(
     def reset(self):
         # reset used names for processing a new email
         self.used_first_names.clear()
+        # reset NEs
+        self.per_list.clear()
+        self.org_list.clear()
+        self.loc_list.clear()
+        self.misc_list.clear()
 
     def get_sentences(self, input_text):
         doc = self.nlp_spacy(input_text)
@@ -107,8 +118,8 @@ def get_ner(self, sentence):
         ner = self.ner_recognizer(sentence)
         return ner
 
-    def pseudonymize_per(self, new_sentence, nelist):
-        unique_ne_list = list(dict.fromkeys(nelist))
+    def pseudonymize_per(self, new_sentence):
+        unique_ne_list = list(dict.fromkeys(self.per_list))
         for ne in unique_ne_list:
             # choose the pseudonym
             nm_list = self.used_first_names
@@ -140,7 +151,6 @@ def pseudonymize_per(self, new_sentence, nelist):
     def pseudonymize_ne(self, ner, sentence):
         # remove any named entities
         entlist = []
-        nelist = []
         new_sentence = sentence
         for i in range(len(ner)):
             entity = ner[i]
@@ -157,18 +167,21 @@ def pseudonymize_ne(self, ner, sentence):
             # replace PER
             if ent_string == "PER":
                 # add the name of this entity to list
-                nelist.append(ent_word)
+                self.per_list.append(ent_word)
             # replace LOC
             elif ent_string == "LOC":
                 new_sentence = new_sentence.replace(ent_word, "[location]")
+                self.loc_list.append(ent_word)
             # replace ORG
             elif ent_string == "ORG":
                 new_sentence = new_sentence.replace(ent_word, "[organization]")
+                self.org_list.append(ent_word)
             # replace MISC
             elif ent_string == "MISC":
                 new_sentence = new_sentence.replace(ent_word, "[misc]")
+                self.misc_list.append(ent_word)
         # replace all unique PER now
-        new_sentence = self.pseudonymize_per(new_sentence, nelist)
+        new_sentence = self.pseudonymize_per(new_sentence)
 
         newlist = [new_sentence]
         return newlist

mailcom/test/test_parse.py

@@ -105,7 +105,8 @@ def test_get_sentences_with_punctuation(get_default_fr):
 def test_pseudonymize_per(get_default_fr):
     sentence = "Francois and Agathe are friends."
     nelist = ["Francois", "Agathe"]
-    pseudonymized_sentence = get_default_fr.pseudonymize_per(sentence, nelist)
+    get_default_fr.per_list = nelist
+    pseudonymized_sentence = get_default_fr.pseudonymize_per(sentence)
     assert "Francois" not in pseudonymized_sentence
     assert "Agathe" not in pseudonymized_sentence
     assert any(

notebook/demo.ipynb

@@ -27,14 +27,51 @@
    "source": [
     "import mailcom.inout\n",
     "import mailcom.parse\n",
-    "import pandas as pd"
+    "import pandas as pd\n",
+    "from IPython.display import display, HTML"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "Below, the input files are loaded from the given `input_dir` directory. You can provide relative or absolute paths to the directory that contains your `eml` or `html` files. All files of the `eml` or `htlm` file type in that directory will be considered input files."
+    "The cell below defines a function used to display the result in the end, and highlight all named entities found in the text. It is used for demonstration purposes in this example."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# function for displaying the result using HTML\n",
+    "def highlight_ne(text, per_list, org_list, loc_list, misc_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",
+    "    \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",
+    "    \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",
+    "    \n",
+    "    return text"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Below, the input files are loaded from the given `input_dir` directory. You can provide relative or absolute paths to the directory that contains your `eml` or `html` files. All files of the `eml` or `html` file type in that directory will be considered input files."
    ]
   },
   {
@@ -69,7 +106,9 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "In the cell below, the emails are looped over and the text is extracted. The text is then split into sentences and the sentences are pseudonymized. The pseudonymized sentences are then joined back into a text and saved to a new file."
+    "In the cell below, the emails are looped over and the text is extracted. The text is then split into sentences and the sentences are pseudonymized. The pseudonymized sentences are then joined back into a text and saved to a new file.\n",
+    "\n",
+    "The input text is displayed and the found named entities are highlighted for demonstration. Note that emails (all words containing '@') are filtered out seperately and thus not highlighted here."
    ]
   },
   {
@@ -86,11 +125,18 @@
     "    # after this function was called, the email metadata can be accessed via io.email_content\n",
     "    # the dict already has the entries content, date, attachments, attachment type\n",
     "    email_dict = io.email_content.copy()\n",
-    "    text = io.get_html_text(text)\n",
+    "    html_text = io.get_html_text(text)\n",
+    "    email_dict[\"html_text\"] = html_text\n",
     "    if not text:\n",
     "        continue\n",
     "    # Test functionality of Pseudonymize class\n",
-    "    output_text = ps.pseudonymize(text)\n",
+    "    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",
+    "    display(HTML(highlighted_html))\n",
+    "\n",
+    "    # add pseudonymized text to dict\n",
     "    email_dict[\"pseudo_content\"] = output_text\n",
     "    out_list.append(email_dict)"
    ]
@@ -128,7 +174,7 @@
     "# print results\n",
     "for idx, mail in df.iterrows():\n",
     "    print(\"Email\", idx)\n",
-    "    print(\"Original Text:\\n\", mail[\"content\"])\n",
+    "    print(\"Original Text:\\n\", mail[\"html_text\"])\n",
     "    print(\"Pseudonymized Text:\\n\", mail[\"pseudo_content\"])\t"
    ]
   },
@@ -156,7 +202,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.9"
+   "version": "3.11.10"
   }
  },
  "nbformat": 4,