Duplicate from beki/pii-anonymizer

Co-authored-by: Benjamin Kilimnik <beki@users.noreply.huggingface.co>

.gitattributes

@@ -0,0 +1,31 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,14 @@
+---
+title: Presidio with custom PII models trained on PII data generated by Privy
+emoji: 📊
+colorFrom: purple
+colorTo: pink
+sdk: streamlit
+sdk_version: 1.10.0
+app_file: app.py
+pinned: false
+license: mit
+duplicated_from: beki/pii-anonymizer
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,212 @@
+
+"""Streamlit app for Presidio + Privy-trained PII models."""
+
+import spacy
+from spacy_recognizer import CustomSpacyRecognizer
+from presidio_analyzer.nlp_engine import NlpEngineProvider
+from presidio_anonymizer import AnonymizerEngine
+from presidio_analyzer import AnalyzerEngine, RecognizerRegistry
+import pandas as pd
+from annotated_text import annotated_text
+from json import JSONEncoder
+import json
+import warnings
+import streamlit as st
+import os
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+warnings.filterwarnings('ignore')
+# from flair_recognizer import FlairRecognizer
+
+# Helper methods
+@st.cache(allow_output_mutation=True)
+def analyzer_engine():
+    """Return AnalyzerEngine."""
+
+    spacy_recognizer = CustomSpacyRecognizer()
+
+    configuration = {
+        "nlp_engine_name": "spacy",
+        "models": [
+            {"lang_code": "en", "model_name": "en_spacy_pii_distilbert"}],
+    }
+
+    # Create NLP engine based on configuration
+    provider = NlpEngineProvider(nlp_configuration=configuration)
+    nlp_engine = provider.create_engine()
+
+    registry = RecognizerRegistry()
+    # add rule-based recognizers
+    registry.load_predefined_recognizers(nlp_engine=nlp_engine)
+    registry.add_recognizer(spacy_recognizer)
+    # remove the nlp engine we passed, to use custom label mappings
+    registry.remove_recognizer("SpacyRecognizer")
+
+    analyzer = AnalyzerEngine(nlp_engine=nlp_engine,
+                              registry=registry, supported_languages=["en"])
+
+    # uncomment for flair-based NLP recognizer
+    # flair_recognizer = FlairRecognizer()
+    # registry.load_predefined_recognizers()
+    # registry.add_recognizer(flair_recognizer)
+    # analyzer = AnalyzerEngine(registry=registry, supported_languages=["en"])
+    return analyzer
+
+
+@st.cache(allow_output_mutation=True)
+def anonymizer_engine():
+    """Return AnonymizerEngine."""
+    return AnonymizerEngine()
+
+
+def get_supported_entities():
+    """Return supported entities from the Analyzer Engine."""
+    return analyzer_engine().get_supported_entities()
+
+
+def analyze(**kwargs):
+    """Analyze input using Analyzer engine and input arguments (kwargs)."""
+    if "entities" not in kwargs or "All" in kwargs["entities"]:
+        kwargs["entities"] = None
+    return analyzer_engine().analyze(**kwargs)
+
+
+def anonymize(text, analyze_results):
+    """Anonymize identified input using Presidio Abonymizer."""
+    if not text:
+        return
+    res = anonymizer_engine().anonymize(text, analyze_results)
+    return res.text
+
+
+def annotate(text, st_analyze_results, st_entities):
+    tokens = []
+    # sort by start index
+    results = sorted(st_analyze_results, key=lambda x: x.start)
+    for i, res in enumerate(results):
+        if i == 0:
+            tokens.append(text[:res.start])
+
+        # append entity text and entity type
+        tokens.append((text[res.start: res.end], res.entity_type))
+
+        # if another entity coming i.e. we're not at the last results element, add text up to next entity
+        if i != len(results) - 1:
+            tokens.append(text[res.end:results[i+1].start])
+        # if no more entities coming, add all remaining text
+        else:
+            tokens.append(text[res.end:])
+    return tokens
+
+
+st.set_page_config(page_title="Privy + Presidio demo (English)", layout="wide")
+
+# Side bar
+st.sidebar.markdown(
+    """
+Detect and anonymize PII in text using an [NLP model](https://huggingface.co/beki/en_spacy_pii_distilbert) trained on protocol traces (JSON, SQL, XML etc.) generated by 
+[Privy](https://github.com/pixie-io/pixie/tree/main/src/datagen/pii/privy) and rule-based classifiers from [Presidio](https://aka.ms/presidio).
+"""
+)
+
+st_entities = st.sidebar.multiselect(
+    label="Which entities to look for?",
+    options=get_supported_entities(),
+    default=list(get_supported_entities()),
+)
+
+st_threshold = st.sidebar.slider(
+    label="Acceptance threshold", min_value=0.0, max_value=1.0, value=0.35
+)
+
+st_return_decision_process = st.sidebar.checkbox(
+    "Add analysis explanations in json")
+
+st.sidebar.info(
+    "Privy is an open source framework for synthetic data generation in protocol trace formats (json, sql, html etc). Presidio is an open source framework for PII detection and anonymization. "
+    "For more info visit [privy](https://github.com/pixie-io/pixie/tree/main/src/datagen/pii/privy) and [aka.ms/presidio](https://aka.ms/presidio)"
+)
+
+
+# Main panel
+analyzer_load_state = st.info(
+    "Starting Presidio analyzer and loading Privy-trained PII model...")
+engine = analyzer_engine()
+analyzer_load_state.empty()
+
+
+st_text = st.text_area(
+    label="Type in some text",
+    value="SELECT shipping FROM users WHERE shipping = '201 Thayer St Providence RI 02912'"
+    "\n\n"
+    "{user: Willie Porter, ip: 192.168.2.80, email: willie@gmail.com}",
+    height=200,
+)
+
+button = st.button("Detect PII")
+
+if 'first_load' not in st.session_state:
+    st.session_state['first_load'] = True
+
+# After
+st.subheader("Analyzed")
+with st.spinner("Analyzing..."):
+    if button or st.session_state.first_load:
+        st_analyze_results = analyze(
+            text=st_text,
+            entities=st_entities,
+            language="en",
+            score_threshold=st_threshold,
+            return_decision_process=st_return_decision_process,
+        )
+        annotated_tokens = annotate(st_text, st_analyze_results, st_entities)
+        # annotated_tokens
+        annotated_text(*annotated_tokens)
+# vertical space
+st.text("")
+
+st.subheader("Anonymized")
+
+with st.spinner("Anonymizing..."):
+    if button or st.session_state.first_load:
+        st_anonymize_results = anonymize(st_text, st_analyze_results)
+        st_anonymize_results
+
+
+# table result
+st.subheader("Detailed Findings")
+if st_analyze_results:
+    res_dicts = [r.to_dict() for r in st_analyze_results]
+    for d in res_dicts:
+        d['Value'] = st_text[d['start']:d['end']]
+    df = pd.DataFrame.from_records(res_dicts)
+    df = df[["entity_type", "Value", "score", "start", "end"]].rename(
+        {
+            "entity_type": "Entity type",
+            "start": "Start",
+            "end": "End",
+            "score": "Confidence",
+        },
+        axis=1,
+    )
+
+    st.dataframe(df, width=1000)
+else:
+    st.text("No findings")
+
+st.session_state['first_load'] = True
+
+# json result
+
+
+class ToDictListEncoder(JSONEncoder):
+    """Encode dict to json."""
+
+    def default(self, o):
+        """Encode to JSON using to_dict."""
+        if o:
+            return o.to_dict()
+        return []
+
+
+if st_return_decision_process:
+    st.json(json.dumps(st_analyze_results, cls=ToDictListEncoder))

flair_recognizer.py

@@ -0,0 +1,245 @@
+import logging
+from typing import Optional, List, Tuple, Set
+
+from presidio_analyzer import (
+    RecognizerResult,
+    EntityRecognizer,
+    AnalysisExplanation,
+)
+from presidio_analyzer.nlp_engine import NlpArtifacts
+
+try:
+    from flair.data import Sentence
+    from flair.models import SequenceTagger
+except ImportError:
+    print("Flair is not installed")
+
+
+logger = logging.getLogger("presidio-analyzer")
+
+
+class FlairRecognizer(EntityRecognizer):
+    """
+    Wrapper for a flair model, if needed to be used within Presidio Analyzer.
+    :example:
+    >from presidio_analyzer import AnalyzerEngine, RecognizerRegistry
+    >flair_recognizer = FlairRecognizer()
+    >registry = RecognizerRegistry()
+    >registry.add_recognizer(flair_recognizer)
+    >analyzer = AnalyzerEngine(registry=registry)
+    >results = analyzer.analyze(
+    >    "My name is Christopher and I live in Irbid.",
+    >    language="en",
+    >    return_decision_process=True,
+    >)
+    >for result in results:
+    >    print(result)
+    >    print(result.analysis_explanation)
+    """
+
+    ENTITIES = [
+        "LOCATION",
+        "PERSON",
+        "NRP",
+        "GPE",
+        "ORGANIZATION",
+        "MAC_ADDRESS",
+        "US_BANK_NUMBER",
+        "IMEI",
+        "TITLE",
+        "LICENSE_PLATE",
+        "US_PASSPORT",
+        "CURRENCY",
+        "ROUTING_NUMBER",
+        "US_ITIN",
+        "US_BANK_NUMBER",
+        "US_DRIVER_LICENSE",
+        "AGE",
+        "PASSWORD",
+        "SWIFT_CODE",
+    ]
+
+    DEFAULT_EXPLANATION = "Identified as {} by Flair's Named Entity Recognition"
+
+    CHECK_LABEL_GROUPS = [
+        ({"LOCATION"}, {"LOC", "LOCATION", "STREET_ADDRESS", "COORDINATE"}),
+        ({"PERSON"}, {"PER", "PERSON"}),
+        ({"NRP"}, {"NORP", "NRP"}),
+        ({"GPE"}, {"GPE"}),
+        ({"ORGANIZATION"}, {"ORG"}),
+        ({"MAC_ADDRESS"}, {"MAC_ADDRESS"}),
+        ({"US_BANK_NUMBER"}, {"US_BANK_NUMBER"}),
+        ({"IMEI"}, {"IMEI"}),
+        ({"TITLE"}, {"TITLE"}),
+        ({"LICENSE_PLATE"}, {"LICENSE_PLATE"}),
+        ({"US_PASSPORT"}, {"US_PASSPORT"}),
+        ({"CURRENCY"}, {"CURRENCY"}),
+        ({"ROUTING_NUMBER"}, {"ROUTING_NUMBER"}),
+        ({"AGE"}, {"AGE"}),
+        ({"CURRENCY"}, {"CURRENCY"}),
+        ({"SWIFT_CODE"}, {"SWIFT_CODE"}),
+        ({"US_ITIN"}, {"US_ITIN"}),
+        ({"US_BANK_NUMBER"}, {"US_BANK_NUMBER"}),
+        ({"US_DRIVER_LICENSE"}, {"US_DRIVER_LICENSE"}),
+    ]
+
+    MODEL_LANGUAGES = {
+        "en":"beki/flair-pii-english-large",
+        # "en":"flair-trf.pt",
+    }
+
+    PRESIDIO_EQUIVALENCES = {
+        "PER": "PERSON",
+        "LOC": "LOCATION",
+        "ORG": "ORGANIZATION",
+        "NROP": "NRP",
+        "URL": "URL",
+        "US_ITIN": "US_ITIN",
+        "US_PASSPORT": "US_PASSPORT",
+        "IBAN_CODE": "IBAN_CODE",
+        "IP_ADDRESS": "IP_ADDRESS",
+        "EMAIL_ADDRESS": "EMAIL",
+        "US_DRIVER_LICENSE": "US_DRIVER_LICENSE",
+        "US_BANK_NUMBER": "US_BANK_NUMBER",
+    }
+
+    def __init__(
+        self,
+        supported_language: str = "en",
+        supported_entities: Optional[List[str]] = None,
+        check_label_groups: Optional[Tuple[Set, Set]] = None,
+        model: SequenceTagger = None,
+    ):
+        self.check_label_groups = (
+            check_label_groups if check_label_groups else self.CHECK_LABEL_GROUPS
+        )
+
+        supported_entities = supported_entities if supported_entities else self.ENTITIES
+        self.model = (
+            model
+            if model
+            else SequenceTagger.load(self.MODEL_LANGUAGES.get(supported_language))
+        )
+
+        super().__init__(
+            supported_entities=supported_entities,
+            supported_language=supported_language,
+            name="Flair Analytics",
+        )
+
+    def load(self) -> None:
+        """Load the model, not used. Model is loaded during initialization."""
+        pass
+
+    def get_supported_entities(self) -> List[str]:
+        """
+        Return supported entities by this model.
+        :return: List of the supported entities.
+        """
+        return self.supported_entities
+
+    # Class to use Flair with Presidio as an external recognizer.
+    def analyze(
+        self, text: str, entities: List[str], nlp_artifacts: NlpArtifacts = None
+    ) -> List[RecognizerResult]:
+        """
+        Analyze text using Text Analytics.
+        :param text: The text for analysis.
+        :param entities: Not working properly for this recognizer.
+        :param nlp_artifacts: Not used by this recognizer.
+        :param language: Text language. Supported languages in MODEL_LANGUAGES
+        :return: The list of Presidio RecognizerResult constructed from the recognized
+            Flair detections.
+        """
+
+        results = []
+
+        sentences = Sentence(text)
+        self.model.predict(sentences)
+
+        # If there are no specific list of entities, we will look for all of it.
+        if not entities:
+            entities = self.supported_entities
+
+        for entity in entities:
+            if entity not in self.supported_entities:
+                continue
+
+            for ent in sentences.get_spans("ner"):
+                if not self.__check_label(
+                    entity, ent.labels[0].value, self.check_label_groups
+                ):
+                    continue
+                textual_explanation = self.DEFAULT_EXPLANATION.format(
+                    ent.labels[0].value
+                )
+                explanation = self.build_flair_explanation(
+                    round(ent.score, 2), textual_explanation
+                )
+                flair_result = self._convert_to_recognizer_result(ent, explanation)
+
+                results.append(flair_result)
+
+        return results
+
+    def _convert_to_recognizer_result(self, entity, explanation) -> RecognizerResult:
+
+        entity_type = self.PRESIDIO_EQUIVALENCES.get(entity.tag, entity.tag)
+        flair_score = round(entity.score, 2)
+
+        flair_results = RecognizerResult(
+            entity_type=entity_type,
+            start=entity.start_position,
+            end=entity.end_position,
+            score=flair_score,
+            analysis_explanation=explanation,
+        )
+
+        return flair_results
+
+    def build_flair_explanation(
+        self, original_score: float, explanation: str
+    ) -> AnalysisExplanation:
+        """
+        Create explanation for why this result was detected.
+        :param original_score: Score given by this recognizer
+        :param explanation: Explanation string
+        :return:
+        """
+        explanation = AnalysisExplanation(
+            recognizer=self.__class__.__name__,
+            original_score=original_score,
+            textual_explanation=explanation,
+        )
+        return explanation
+
+    @staticmethod
+    def __check_label(
+        entity: str, label: str, check_label_groups: Tuple[Set, Set]
+    ) -> bool:
+        return any(
+            [entity in egrp and label in lgrp for egrp, lgrp in check_label_groups]
+        )
+
+
+if __name__ == "__main__":
+
+    from presidio_analyzer import AnalyzerEngine, RecognizerRegistry
+
+    flair_recognizer = (
+        FlairRecognizer()
+    )  # This would download a very large (+2GB) model on the first run
+
+    registry = RecognizerRegistry()
+    registry.add_recognizer(flair_recognizer)
+
+    analyzer = AnalyzerEngine(registry=registry)
+
+    results = analyzer.analyze(
+        "{first_name: Moustafa, sale_id: 235234}",
+        language="en",
+        return_decision_process=True,
+    )
+    for result in results:
+        print(result)
+        print(result.analysis_explanation)

requirements.txt

@@ -0,0 +1,8 @@
+pandas
+streamlit
+presidio-anonymizer
+presidio-analyzer
+torch
+#flair@git+https://github.com/flairNLP/flair.git@d4ed67bf663e4066517f00397412510d90043653
+st-annotated-text
+https://huggingface.co/beki/en_spacy_pii_distilbert/resolve/main/en_spacy_pii_distilbert-any-py3-none-any.whl

spacy_recognizer.py

@@ -0,0 +1,131 @@
+import logging
+from typing import Optional, List, Tuple, Set
+
+from presidio_analyzer import (
+    RecognizerResult,
+    LocalRecognizer,
+    AnalysisExplanation,
+)
+from presidio_analyzer.nlp_engine import NlpArtifacts
+from presidio_analyzer.predefined_recognizers.spacy_recognizer import SpacyRecognizer
+
+logger = logging.getLogger("presidio-analyzer")
+
+
+class CustomSpacyRecognizer(LocalRecognizer):
+
+    ENTITIES = [
+        "LOCATION",
+        "PERSON",
+        "NRP",
+        "ORGANIZATION",
+        "DATE_TIME",
+    ]
+
+    DEFAULT_EXPLANATION = "Identified as {} by Spacy's Named Entity Recognition (Privy-trained)"
+
+    CHECK_LABEL_GROUPS = [
+        ({"LOCATION"}, {"LOC", "LOCATION", "STREET_ADDRESS", "COORDINATE"}),
+        ({"PERSON"}, {"PER", "PERSON"}),
+        ({"NRP"}, {"NORP", "NRP"}),
+        ({"ORGANIZATION"}, {"ORG"}),
+        ({"DATE_TIME"}, {"DATE_TIME"}),
+    ]
+
+    MODEL_LANGUAGES = {
+        "en": "beki/en_spacy_pii_distilbert",
+    }
+
+    PRESIDIO_EQUIVALENCES = {
+        "PER": "PERSON",
+        "LOC": "LOCATION",
+        "ORG": "ORGANIZATION",
+        "NROP": "NRP",
+        "DATE_TIME": "DATE_TIME",
+    }
+
+    def __init__(
+        self,
+        supported_language: str = "en",
+        supported_entities: Optional[List[str]] = None,
+        check_label_groups: Optional[Tuple[Set, Set]] = None,
+        context: Optional[List[str]] = None,
+        ner_strength: float = 0.85,
+    ):
+        self.ner_strength = ner_strength
+        self.check_label_groups = (
+            check_label_groups if check_label_groups else self.CHECK_LABEL_GROUPS
+        )
+        supported_entities = supported_entities if supported_entities else self.ENTITIES
+        super().__init__(
+            supported_entities=supported_entities,
+            supported_language=supported_language,
+        )
+
+    def load(self) -> None:
+        """Load the model, not used. Model is loaded during initialization."""
+        pass
+
+    def get_supported_entities(self) -> List[str]:
+        """
+        Return supported entities by this model.
+        :return: List of the supported entities.
+        """
+        return self.supported_entities
+
+    def build_spacy_explanation(
+        self, original_score: float, explanation: str
+    ) -> AnalysisExplanation:
+        """
+        Create explanation for why this result was detected.
+        :param original_score: Score given by this recognizer
+        :param explanation: Explanation string
+        :return:
+        """
+        explanation = AnalysisExplanation(
+            recognizer=self.__class__.__name__,
+            original_score=original_score,
+            textual_explanation=explanation,
+        )
+        return explanation
+
+    def analyze(self, text, entities, nlp_artifacts=None):  # noqa D102
+        results = []
+        if not nlp_artifacts:
+            logger.warning("Skipping SpaCy, nlp artifacts not provided...")
+            return results
+
+        ner_entities = nlp_artifacts.entities
+
+        for entity in entities:
+            if entity not in self.supported_entities:
+                continue
+            for ent in ner_entities:
+                if not self.__check_label(entity, ent.label_, self.check_label_groups):
+                    continue
+                textual_explanation = self.DEFAULT_EXPLANATION.format(
+                    ent.label_)
+                explanation = self.build_spacy_explanation(
+                    self.ner_strength, textual_explanation
+                )
+                spacy_result = RecognizerResult(
+                    entity_type=entity,
+                    start=ent.start_char,
+                    end=ent.end_char,
+                    score=self.ner_strength,
+                    analysis_explanation=explanation,
+                    recognition_metadata={
+                        RecognizerResult.RECOGNIZER_NAME_KEY: self.name
+                    },
+                )
+                results.append(spacy_result)
+
+        return results
+
+    @staticmethod
+    def __check_label(
+        entity: str, label: str, check_label_groups: Tuple[Set, Set]
+    ) -> bool:
+        return any(
+            [entity in egrp and label in lgrp for egrp, lgrp in check_label_groups]
+        )