Update app.py

app.py

@@ -1,5 +1,6 @@
 import re
 import string
+from collections import defaultdict
 
 import seaborn as sns
 import streamlit as st
@@ -12,10 +13,12 @@ from transformers import (
 )
 
 
-# Setup & Constants
+# Setup
 def setup_page():
     st.set_page_config(
-        page_title="Juristische Anonymisierung", page_icon="⚖️", layout="wide"
+        page_title="Juristische Anonymisierung",
+        page_icon="⚖️",
+        layout="wide",
     )
     logging.set_verbosity(logging.ERROR)
     st.markdown(
@@ -115,11 +118,12 @@ def load_ner_model():
 
 
 @st.cache_data(show_spinner=False)
-def ner_merge_lines(text):
+def ner_merge_lines(text: str):
     ner = load_ner_model()
     merged_lines = []
     for line in text.splitlines():
         if not line.strip():
+            merged_lines.append((line, []))
             continue
         tokens = ner(line)
         merged = merge_entities(tokens)
@@ -134,7 +138,6 @@ def merge_entities(entities):
     merged = [ents[0].copy()]
     merged[0]["score_sum"] = ents[0]["score"]
     merged[0]["count"] = 1
-
     for ent in ents[1:]:
         prev = merged[-1]
         if ent["index"] == prev["index"] + 1:
@@ -151,11 +154,9 @@ def merge_entities(entities):
             new_ent["score_sum"] = ent["score"]
             new_ent["count"] = 1
             merged.append(new_ent)
-
     if "score_sum" in merged[-1]:
         merged[-1]["score"] = merged[-1]["score_sum"] / merged[-1]["count"]
         del merged[-1]["score_sum"], merged[-1]["count"]
-
     final = []
     for ent in merged:
         w = ent["word"].strip()
@@ -175,6 +176,34 @@ def truncate(number, decimals=2):
     return int(number * factor) / factor
 
 
+# Canonical grouping
+def canonical_key(text: str, label: str):
+    s = text.casefold().strip()
+
+    if label == "RS":
+        m = re.search(r"(ecli:[a-z]{2}:[a-z0-9]+:\d{4}:[a-z0-9.\-]+)", s)
+        if m:
+            original = text[m.start() : m.end()]
+            canon = m.group(1).replace(" ", "")
+            return (canon, label, original)
+
+        m = re.search(
+            r"((?:[ivxlcdm]+|\d{1,3})\s*[a-zäöüß]{1,3}\s*\d{1,6}\s*/\s*\d{2,4})", s
+        )
+        if m:
+            original = text[m.start() : m.end()].strip()
+            canon = re.sub(r"\s+", "", m.group(1))
+            return (canon, label, original)
+
+        cleaned = re.sub(r"[^\w]+", "", s)
+        return (cleaned, label, text.strip())
+
+    cleaned_generic = re.sub(r"[^\w]+", " ", s)
+    cleaned_generic = re.sub(r"\s+", " ", cleaned_generic).strip()
+    return (cleaned_generic, label, text.strip())
+
+
+# Rendering
 def highlight_entities(
     line,
     merged_entities,
@@ -198,11 +227,13 @@ def highlight_entities(
         truncated_score = truncate(ent["score"], 2)
         tooltip = f"{label_desc} ({truncated_score:.2f})"
         color = ENTITY_COLORS.get(label, "#cccccc")
+
         html += line[last_end:start]
 
         should_anonymize = any(
             label in entity_importance[level] for level in importance_levels
         )
+
         if should_anonymize:
             key = (ent["word"].lower(), label)
             if key not in anonymized_map:
@@ -210,11 +241,14 @@ def highlight_entities(
                 suffix = chr(ord("A") + count)
                 label_counters[label] = count + 1
                 anonymized_map[key] = suffix
+
             suffix = anonymized_map[key]
             display = f"{label_desc} {suffix}"
+
             normalized_word = ent["word"].strip().lower()
             display_key = f"{label_desc} {suffix} : {normalized_word}"
-            if display_key not in allowed_keys:
+
+            if allowed_keys and display_key not in allowed_keys:
                 display = ent["word"]
                 style = ""
                 css_class = "entity"
@@ -226,7 +260,7 @@ def highlight_entities(
             style = ""
             css_class = "entity"
 
-        html += f'<span class="{css_class}" style="{style}">{display}<span class="tooltip">{tooltip}</span></span>'
+        html += f'<span class="{css_class}" style="{style}" title="{tooltip}">{display}</span>'
         last_end = end
 
     html += line[last_end:]
@@ -243,6 +277,7 @@ def main():
         st.session_state.manual_phrases = []
 
     st.markdown("#### Juristische Anonymisierung")
+
     uploaded_file = st.file_uploader(
         "Bitte laden Sie eine .txt-Datei hoch:", type="txt"
     )
@@ -263,7 +298,8 @@ def main():
             ent_list = [entity_labels[k] for k in entity_importance[level]]
             st.markdown(f"**{label}**: {', '.join(ent_list)}")
 
-    threshold = st.slider("Schwellenwert für das Modellvertrauen:", 0.0, 1.0, 0.8, 0.01)
+    threshold = st.slider("Schwellenwert für das Modellvertrauen:", 0.0, 1.0, 0.5, 0.01)
+
     st.markdown("---")
 
     if uploaded_file:
@@ -277,18 +313,16 @@ def main():
         with st.spinner("Modell wird einmalig auf die Datei angewendet..."):
             merged_all_lines = ner_merge_lines(text)
 
+        # Manual phrases to RED
         manual_phrases = st.session_state.manual_phrases
         overlap_warnings = set()
-
         for idx, (line, merged) in enumerate(merged_all_lines):
             for phrase in manual_phrases:
                 for match in re.finditer(re.escape(phrase), line.lower()):
                     start, end = match.start(), match.end()
-
                     if any(start < e["end"] and end > e["start"] for e in merged):
                         overlap_warnings.add(phrase)
                         continue
-
                     merged.append(
                         {
                             "start": start,
@@ -299,75 +333,140 @@ def main():
                             "index": 9999,
                         }
                     )
-
             merged_all_lines[idx] = (line, sorted(merged, key=lambda x: x["start"]))
 
-        label_counters = {}
-        anonymized_map = {}
-        all_display_keys = []
+        # Grouping layer for the sidebar
+        groups = defaultdict(
+            lambda: {"variants": set(), "displays": set(), "rep": None}
+        )
 
         for _, merged in merged_all_lines:
             for ent in merged:
                 label = ent["entity"].split("-")[-1]
                 if any(label in entity_importance[lvl] for lvl in importance_levels):
-                    key = (ent["word"].lower(), label)
-                    if key not in anonymized_map:
-                        count = label_counters.get(label, 0)
-                        suffix = chr(ord("A") + count)
-                        label_counters[label] = count + 1
-                        anonymized_map[key] = suffix
-                    suffix = anonymized_map[key]
-                    normalized_word = ent["word"].strip().lower()
-                    display = f"{entity_labels.get(label, label)} {suffix} : {normalized_word}"
-                    if display not in all_display_keys:
-                        all_display_keys.append(display)
-
-        all_display_keys.sort(key=lambda tag: tag.lower())
+
+                    variant_norm = ent["word"].strip().lower()
+                    canon_key, canon_label, display_key = canonical_key(
+                        ent["word"], label
+                    )
+
+                    g = groups[(canon_key, canon_label)]
+                    g["variants"].add(variant_norm)
+                    g["displays"].add(display_key)
+
+        # Suffix per canonical group
+        label_counters_for_groups = {}
+        for (canon_text, label), data in groups.items():
+            count = label_counters_for_groups.get(label, 0)
+            suffix = chr(ord("A") + count)
+            label_counters_for_groups[label] = count + 1
+            data["suffix"] = suffix
+
+        for key, data in groups.items():
+            if data["displays"]:
+                data["rep"] = max(data["displays"], key=len)
+            else:
+                data["rep"] = ""
+
+        anonymized_map = {}
+        for (canon_text, label), data in groups.items():
+            suffix = data["suffix"]
+            for v in data["variants"]:
+                anonymized_map[(v, label)] = suffix
+
+        entity_labels_map = entity_labels
+        display_to_variants = {}
+        groups_by_label_desc = defaultdict(list)
+        all_display_keys = set()
+
+        for (canon_text, label), data in groups.items():
+            label_desc = entity_labels_map.get(label, label)
+            suffix = data["suffix"]
+            shown = f"{label_desc} {suffix} : {data['rep']}"
+            groups_by_label_desc[label_desc].append(shown)
+            display_keys = [f"{label_desc} {suffix} : {v}" for v in data["variants"]]
+            display_to_variants[shown] = display_keys
+            all_display_keys.update(display_keys)
+
+        label_order = [
+            "RS",
+            "GS",
+            "PER",
+            "AN",
+            "GRT",
+            "VO",
+            "VS",
+            "VT",
+            "EUN",
+            "LIT",
+            "UN",
+            "INN",
+            "ORG",
+            "MRK",
+            "RR",
+            "LD",
+            "LDS",
+            "ST",
+            "STR",
+            "RED",
+        ]
+        label_order_desc = [entity_labels_map.get(x, x) for x in label_order]
 
         with st.sidebar:
             st.markdown("### Neue Phrase schwärzen:")
-
             if "manual_phrases" not in st.session_state:
                 st.session_state.manual_phrases = []
 
             with st.form("manual_add_form"):
                 new_phrase = st.text_input("Neue Phrase:")
                 submitted = st.form_submit_button("Hinzufügen")
-                with st.sidebar.expander(
-                    "Hinweise zu manuellen Phrasen", expanded=False
-                ):
-                    st.markdown("**Noch in Entwicklung**")
-                    st.markdown(
-                        "_Manuelle Schwärzungen können fehlschlagen, wenn sich die Phrase mit bereits erkannten Entitäten überschneidet oder über mehrere Zeilen erstreckt._"
-                    )
 
-                if submitted and new_phrase.strip():
-                    cleaned = new_phrase.strip().lower()
-                    if cleaned not in st.session_state.manual_phrases:
-                        st.session_state.manual_phrases.append(cleaned)
-                        st.rerun()
+            with st.sidebar.expander("Hinweise zu manuellen Phrasen", expanded=False):
+                st.markdown("**Noch in Entwicklung**")
+                st.markdown(
+                    "_Manuelle Schwärzungen können fehlschlagen, wenn sich die Phrase "
+                    "mit bereits erkannten Entitäten überschneidet oder über mehrere "
+                    "Zeilen erstreckt._"
+                )
+
+            if submitted and new_phrase.strip():
+                cleaned = new_phrase.strip().lower()
+                if cleaned not in st.session_state.manual_phrases:
+                    st.session_state.manual_phrases.append(cleaned)
+                st.rerun()
 
             st.markdown("---")
             st.markdown("### Anonymisierte Entitäten verwalten:")
-            selected_keys = []
-            for label_code in sorted(
-                set(k[1] for k in anonymized_map.keys()),
-                key=lambda x: entity_labels.get(x, x),
-            ):
-                group = [k for k in anonymized_map if k[1] == label_code]
-                label_name = entity_labels[label_code]
-                st.markdown(f"**{label_name}**")
-                for key in sorted(group, key=lambda k: anonymized_map[k]):
-                    suffix = anonymized_map[key]
-                    normalized_word = key[0].strip().lower()
-                    entity_display = f"{label_name} {suffix} : {normalized_word}"
-                    if st.checkbox(entity_display, value=True, key=entity_display):
-                        selected_keys.append(entity_display)
+
+            selected_canon = []
+            for lab_desc in label_order_desc:
+                items = groups_by_label_desc.get(lab_desc, [])
+                if not items:
+                    continue
+                st.markdown(f"**{lab_desc}**")
+                for shown in sorted(items, key=str.lower):
+                    checked = st.checkbox(shown, value=True, key=f"chk::{shown}")
+                    if checked:
+                        selected_canon.append(shown)
+
+        if not selected_canon and groups_by_label_desc:
+            selected_canon = [
+                x for items in groups_by_label_desc.values() for x in items
+            ]
+
+        allowed_keys = set()
+        for shown in selected_canon:
+            allowed_keys.update(display_to_variants.get(shown, []))
+
+        if not allowed_keys and all_display_keys:
+            allowed_keys = set(all_display_keys)
+
+        label_counters_runtime = {}
 
         anonymized_lines = []
         for line, merged in merged_all_lines:
             if not line.strip():
-                st.markdown("<br>", unsafe_allow_html=True)
+                st.markdown("<br/>", unsafe_allow_html=True)
                 anonymized_lines.append("")
                 continue
 
@@ -376,18 +475,20 @@ def main():
                 merged,
                 importance_levels,
                 threshold,
-                label_counters,
+                label_counters_runtime,
                 anonymized_map,
-                selected_keys,
+                allowed_keys,
                 entity_labels,
                 entity_importance,
                 ENTITY_COLORS,
             )
+
             st.markdown(
-                f'<div style="margin-bottom:0.8rem; line-height:1.8;">{html_line}</div>',
+                f'<div style="white-space: pre-wrap;">{html_line}</div>',
                 unsafe_allow_html=True,
             )
-            cleaned = re.sub(r'<span class="tooltip">.*?</span>', "", html_line)
+
+            cleaned = re.sub(r"<!--.*?-->", "", html_line, flags=re.DOTALL)
             text_only = re.sub(r"<[^>]+>", "", cleaned)
             anonymized_lines.append(text_only.strip())