Update presidio_nlp_engine_config.py

presidio_nlp_engine_config.py

@@ -10,11 +10,11 @@ from presidio_analyzer.nlp_engine import (
 
 logger = logging.getLogger("presidio-streamlit")
 
-CZECH_ENTITY_MAPPING = {
-    "PER": "OSOBA",
-    "LOC": "ADRESA",
-    "ORG": "ORGANIZACE",
-    "MISC": "RŮZNÉ",
+ENGLISH_ENTITY_MAPPING = {
+    "PER": "PERSON",
+    "LOC": "ADDRESS",
+    "ORG": "ORGANIZATION",
+    "MISC": "MISC",
 }
 
 def create_nlp_engine_with_spacy(
@@ -26,20 +26,20 @@ def create_nlp_engine_with_spacy(
     """
     nlp_configuration = {
         "nlp_engine_name": "spacy",
-        "models": [{"lang_code": "cs", "model_name": model_path}],
+        "models": [{"lang_code": "en", "model_name": model_path}],
         "ner_model_configuration": {
             "model_to_presidio_entity_mapping": {
-                **CZECH_ENTITY_MAPPING,
-                "PERSON": "OSOBA",
-                "GPE": "ADRESA",
-                "LOCATION": "ADRESA",
-                "ORGANIZATION": "ORGANIZACE",
-                "DATE": "DATUM_NAROZENÍ",
-                "CARDINAL": "RODNÉ_ČÍSLO",
-                "ORG": "IČO",
+                **ENGLISH_ENTITY_MAPPING,
+                "PERSON": "PERSON",
+                "GPE": "ADDRESS",
+                "LOCATION": "ADDRESS",
+                "ORGANIZATION": "ORGANIZATION",
+                "DATE": "DATE_OF_BIRTH",
+                "CARDINAL": "NATIONAL_ID_NUMBER",
+                "ORG": "COMPANY_ID",
             },
             "low_confidence_score_multiplier": 0.4,
-            "low_score_entity_names": ["ORG", "ORGANIZACE"],
+            "low_score_entity_names": ["ORG", "ORGANIZATION"],
         },
     }
 
@@ -59,15 +59,15 @@ def create_nlp_engine_with_stanza(
     """
     nlp_configuration = {
         "nlp_engine_name": "stanza",
-        "models": [{"lang_code": "cs", "model_name": model_path}],
+        "models": [{"lang_code": "en", "model_name": model_path}],
         "ner_model_configuration": {
             "model_to_presidio_entity_mapping": {
-                **CZECH_ENTITY_MAPPING,
-                "PERSON": "OSOBA",
-                "GPE": "ADRESA",
-                "LOCATION": "ADRESA",
-                "ORGANIZATION": "ORGANIZACE",
-                "DATE": "DATUM_NAROZENÍ",
+                **ENGLISH_ENTITY_MAPPING,
+                "PERSON": "PERSON",
+                "GPE": "ADDRESS",
+                "LOCATION": "ADDRESS",
+                "ORGANIZATION": "ORGANIZATION",
+                "DATE": "DATE_OF_BIRTH",
             }
         },
     }
@@ -94,22 +94,22 @@ def create_nlp_engine_with_transformers(
         "nlp_engine_name": "transformers",
         "models": [
             {
-                "lang_code": "cs",
-                "model_name": {"spacy": "cs_core_news_sm", "transformers": model_path},
+                "lang_code": "en",
+                "model_name": {"spacy": "en_core_web_sm", "transformers": model_path},
             }
         ],
         "ner_model_configuration": {
             "model_to_presidio_entity_mapping": {
-                **CZECH_ENTITY_MAPPING,
-                "PERSON": "OSOBA",
-                "LOC": "ADRESA",
-                "GPE": "ADRESA",
-                "ORG": "ORGANIZACE",
-                "DATE": "DATUM_NAROZENÍ",
-                "CARDINAL": "RODNÉ_ČÍSLO",
-                "ID": "IČO",
+                **ENGLISH_ENTITY_MAPPING,
+                "PERSON": "PERSON",
+                "LOC": "ADDRESS",
+                "GPE": "ADDRESS",
+                "ORG": "ORGANIZATION",
+                "DATE": "DATE_OF_BIRTH",
+                "CARDINAL": "NATIONAL_ID_NUMBER",
+                "ID": "COMPANY_ID",
                 "EMAIL": "EMAIL",
-                "PHONE": "TELEFON",
+                "PHONE": "PHONE",
             },
             "low_confidence_score_multiplier": 0.4,
             "low_score_entity_names": ["ID"],
@@ -148,13 +148,13 @@ def create_nlp_engine_with_flair(
     registry = RecognizerRegistry()
     registry.load_predefined_recognizers()
 
-    if not spacy.util.is_package("cs_core_news_sm"):
-        spacy.cli.download("cs_core_news_sm")
+    if not spacy.util.is_package("en_core_web_sm"):
+        spacy.cli.download("en_core_web_sm")
     
     flair_recognizer = FlairRecognizer(model_path=model_path)
     nlp_configuration = {
         "nlp_engine_name": "spacy",
-        "models": [{"lang_code": "cs", "model_name": "cs_core_news_sm"}],
+        "models": [{"lang_code": "en", "model_name": "en_core_web_sm"}],
     }
     registry.add_recognizer(flair_recognizer)
     registry.remove_recognizer("SpacyRecognizer")
@@ -184,7 +184,7 @@ def create_nlp_engine_with_azure_ai_language(ta_key: str, ta_endpoint: str):
     )
     nlp_configuration = {
         "nlp_engine_name": "spacy",
-        "models": [{"lang_code": "cs", "model_name": "cs_core_news_sm"}],
+        "models": [{"lang_code": "en", "model_name": "en_core_web_sm"}],
     }
 
     nlp_engine = NlpEngineProvider(nlp_configuration=nlp_configuration).create_engine()
@@ -192,4 +192,4 @@ def create_nlp_engine_with_azure_ai_language(ta_key: str, ta_endpoint: str):
     registry.add_recognizer(azure_ai_language_recognizer)
     registry.remove_recognizer("SpacyRecognizer")
 
-    return nlp_engine, registry
+    return nlp_engine, registry