entity_extraction / services /pii_service.py
adit94's picture
Upload 9 files
846b4a5 verified
from pprint import pprint
import json
from presidio_analyzer import AnalyzerEngine, RecognizerRegistry
from presidio_analyzer.nlp_engine import NlpEngineProvider, NlpArtifacts
from presidio_analyzer import PatternRecognizer
from presidio_analyzer import Pattern, PatternRecognizer
from presidio_analyzer.predefined_recognizers import SpacyRecognizer
from presidio_analyzer.predefined_recognizers import IbanRecognizer, EmailRecognizer, IpRecognizer,\
EmailRecognizer, PhoneRecognizer, UrlRecognizer, DateRecognizer
from presidio_anonymizer import AnonymizerEngine
from presidio_anonymizer.entities import OperatorConfig
import logging
from typing import Optional, List, Tuple, Set
from presidio_analyzer import (
RecognizerResult,
EntityRecognizer,
AnalysisExplanation,
)
from flair.data import Sentence
from flair.models import SequenceTagger
### Creating FlairRecognizer class for NER(names, location)
class FlairRecognizer(EntityRecognizer):
ENTITIES = [
"LOCATION",
"PERSON",
"ORGANIZATION",
# "MISCELLANEOUS" # - There are no direct correlation with Presidio entities.
]
DEFAULT_EXPLANATION = "Identified as {} by Flair's Named Entity Recognition"
CHECK_LABEL_GROUPS = [
({"LOCATION"}, {"LOC", "LOCATION"}),
({"PERSON"}, {"PER", "PERSON"}),
({"ORGANIZATION"}, {"ORG"}),
# ({"MISCELLANEOUS"}, {"MISC"}), # Probably not PII
]
MODEL_LANGUAGES = {
#"en": "flair/ner-english-large",
#"es": "flair/ner-spanish-large",
"de": "flair/ner-german-large",
#"nl": "flair/ner-dutch-large",
}
PRESIDIO_EQUIVALENCES = {
"PER": "PERSON",
"LOC": "LOCATION",
"ORG": "ORGANIZATION",
# 'MISC': 'MISCELLANEOUS' # - Probably not PII
}
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",
)
print("Flair class initialized")
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]
)
class PIIService:
def __init__(self):
configuration = {
"nlp_engine_name": "spacy",
"models": [
{"lang_code": "de", "model_name": "de_core_news_sm"}
],
}
# Create NLP engine based on configuration
provider = NlpEngineProvider(nlp_configuration=configuration)
nlp_engine = provider.create_engine()
## Creating regex for PatternRecognizers - SWIFT, vehicle number, zipcode, ssn
swift_regex = r"\b[A-Z]{4}DE[A-Z0-9]{2}(?:[A-Z0-9]{3})?"
vehicle_number_with_hyphen_regex = r"\b[A-ZÄÖÜ]{1,3}-[A-ZÄÖÜ]{1,2}-[0-9]{1,4}"
vehicle_number_without_hyphen_regex = r"\b[A-ZÄÖÜ]{1,3}[A-ZÄÖÜ]{1,2}[0-9]{1,4}"
german_zipcode_regex = r"\b((?:0[1-46-9]\d{3})|(?:[1-357-9]\d{4})|(?:[4][0-24-9]\d{3})|(?:[6][013-9]\d{3}))\b(?![\d/])"
german_ssn_regex = r"\b\d{2}\s?\d{6}\s?[A-Z]\s?\d{3}\b"
# Creating Presidio pattern object
vehicle_numbers_pattern1 = Pattern(name="vehicle_pattern", regex=vehicle_number_without_hyphen_regex, score=1)
vehicle_numbers_pattern2 = Pattern(name="vehicle_pattern", regex=vehicle_number_with_hyphen_regex, score=1)
swift_pattern = Pattern(name="bank_swift_pattern", regex=swift_regex, score=1)
germanzipcode_pattern = Pattern(name="german_zip_pattern",regex=german_zipcode_regex, score=1)
german_ssn_pattern = Pattern(name="german_ssn_pattern",regex=german_ssn_regex, score=1)
# Define the recognizer
swift_recognizer = PatternRecognizer(supported_entity="SWIFT", supported_language="de",patterns=[swift_pattern])
vehicle_number_recognizer = PatternRecognizer(supported_entity="VEHICLE_NUMBER", supported_language="de",patterns=[vehicle_numbers_pattern1,vehicle_numbers_pattern2])
germanzip_recognizer = PatternRecognizer(supported_entity="GERMAN_ZIP", supported_language="de",patterns=[germanzipcode_pattern])
germanssn_recognizer = PatternRecognizer(supported_entity="GERMAN_SSN", supported_language="de",patterns=[german_ssn_pattern])
## Lading flair entity model for person, location ID
print("Loading flair")
flair_recognizer = FlairRecognizer(supported_language="de")
print("Flair loaded")
registry = RecognizerRegistry()
#registry.load_predefined_recognizers()
#registry.add_recognizer(SpacyRecognizer(supported_language="de"))
#registry.add_recognizer(SpacyRecognizer(supported_language="en"))
registry.remove_recognizer("SpacyRecognizer")
registry.add_recognizer(flair_recognizer)
registry.add_recognizer(swift_recognizer)
registry.add_recognizer(vehicle_number_recognizer)
registry.add_recognizer(germanzip_recognizer)
registry.add_recognizer(germanssn_recognizer)
## Adding predefined recognizers
registry.add_recognizer(IbanRecognizer(supported_language="de"))
registry.add_recognizer(DateRecognizer(supported_language="de"))
registry.add_recognizer(EmailRecognizer(supported_language="de"))
registry.add_recognizer(IpRecognizer(supported_language="de"))
registry.add_recognizer(PhoneRecognizer(supported_language="de"))
registry.add_recognizer(UrlRecognizer(supported_language="de"))
registry.add_recognizer(PhoneRecognizer(supported_language="de"))
print("Recognizer registry loaded")
self.analyzer = AnalyzerEngine(registry=registry, nlp_engine=nlp_engine, supported_languages=["de"])
#print(f"Type of recognizers ::\n {self.analyzer.registry.recognizers}")
print("PII initialized")
self.anonymizer = AnonymizerEngine()
def identify(self, text):
results_de = self.analyzer.analyze(
text,
language='de'
)
#anonymized_results = self.anonymize(results_de, text)
entities = []
for result in results_de:
result_dict = result.to_dict()
temp_entity = {
"start":result_dict['start'],
"end":result_dict['end'],
"entity_type":result_dict['entity_type'],
"score":result_dict['score'],
"word":text[result_dict['start']:result_dict['end']]
}
entities.append(temp_entity)
return {"entities":entities, "text":text}#, "anonymized_text":anonymized_results['text']}
"""def anonymize(self, entities, text):
anonymized_results = self.anonymizer.anonymize(
text=text,
analyzer_results=entities,
#operators={"DEFAULT": OperatorConfig("replace", {"new_value": "<ANONYMIZED>"})},
)
return ""#json.loads(anonymized_results.to_json())"""
def add_mask(self, data):
masked_data = []
entity_count = {}
for item_idx,item in enumerate(data['entities']):
entity_type = item['entity_type']
word = item['word']
suffix = entity_count.get(entity_type, 0) + 1
entity_count[entity_type] = suffix
masked_word = f"{entity_type}_{suffix}"
item['mask'] = masked_word
#data['entities'][item_idx]['mask'] = masked_word
masked_data.append(item)
return masked_data
def anonymize(self, entities, text):
print("anonymyzing")
updated_text = text
for ent_idx, ent in enumerate(entities):
#text[ent['start']:ent['end']] = ent['mask']
updated_text = updated_text[:ent['start']] + " " + ent['mask'] + " " + updated_text[ent['end']:]
return updated_text
def remove_overlapping_entities(entities):
return