Update presidio_nlp_engine_config.py

presidio_nlp_engine_config.py

@@ -10,6 +10,12 @@ from presidio_analyzer.nlp_engine import (
 
 logger = logging.getLogger("presidio-streamlit")
 
+CZECH_ENTITY_MAPPING = {
+    "PER": "OSOBA",
+    "LOC": "ADRESA",
+    "ORG": "ORGANIZACE",
+    "MISC": "RŮZNÉ",
+}
 
 def create_nlp_engine_with_spacy(
     model_path: str,
@@ -20,23 +26,20 @@ def create_nlp_engine_with_spacy(
     """
     nlp_configuration = {
         "nlp_engine_name": "spacy",
-        "models": [{"lang_code": "en", "model_name": model_path}],
+        "models": [{"lang_code": "cs", "model_name": model_path}],
         "ner_model_configuration": {
             "model_to_presidio_entity_mapping": {
-                "PER": "PERSON",
-                "PERSON": "PERSON",
-                "NORP": "NRP",
-                "FAC": "FACILITY",
-                "LOC": "LOCATION",
-                "GPE": "LOCATION",
-                "LOCATION": "LOCATION",
-                "ORG": "ORGANIZATION",
-                "ORGANIZATION": "ORGANIZATION",
-                "DATE": "DATE_TIME",
-                "TIME": "DATE_TIME",
+                **CZECH_ENTITY_MAPPING,
+                "PERSON": "OSOBA",
+                "GPE": "ADRESA",
+                "LOCATION": "ADRESA",
+                "ORGANIZATION": "ORGANIZACE",
+                "DATE": "DATUM_NAROZENÍ",
+                "CARDINAL": "RODNÉ_ČÍSLO",
+                "ORG": "IČO",
             },
             "low_confidence_score_multiplier": 0.4,
-            "low_score_entity_names": ["ORG", "ORGANIZATION"],
+            "low_score_entity_names": ["ORG", "ORGANIZACE"],
         },
     }
 
@@ -47,7 +50,6 @@ def create_nlp_engine_with_spacy(
 
     return nlp_engine, registry
 
-
 def create_nlp_engine_with_stanza(
     model_path: str,
 ) -> Tuple[NlpEngine, RecognizerRegistry]:
@@ -57,20 +59,15 @@ def create_nlp_engine_with_stanza(
     """
     nlp_configuration = {
         "nlp_engine_name": "stanza",
-        "models": [{"lang_code": "en", "model_name": model_path}],
+        "models": [{"lang_code": "cs", "model_name": model_path}],
         "ner_model_configuration": {
             "model_to_presidio_entity_mapping": {
-                "PER": "PERSON",
-                "PERSON": "PERSON",
-                "NORP": "NRP",
-                "FAC": "FACILITY",
-                "LOC": "LOCATION",
-                "GPE": "LOCATION",
-                "LOCATION": "LOCATION",
-                "ORG": "ORGANIZATION",
-                "ORGANIZATION": "ORGANIZATION",
-                "DATE": "DATE_TIME",
-                "TIME": "DATE_TIME",
+                **CZECH_ENTITY_MAPPING,
+                "PERSON": "OSOBA",
+                "GPE": "ADRESA",
+                "LOCATION": "ADRESA",
+                "ORGANIZATION": "ORGANIZACE",
+                "DATE": "DATUM_NAROZENÍ",
             }
         },
     }
@@ -82,7 +79,6 @@ def create_nlp_engine_with_stanza(
 
     return nlp_engine, registry
 
-
 def create_nlp_engine_with_transformers(
     model_path: str,
 ) -> Tuple[NlpEngine, RecognizerRegistry]:
@@ -92,44 +88,32 @@ def create_nlp_engine_with_transformers(
     would return NlpArtifacts such as POS and lemmas.
     :param model_path: HuggingFace model path.
     """
-    print(f"Loading Transformers model: {model_path} of type {type(model_path)}")
+    logger.info(f"Loading Transformers model: {model_path}")
 
     nlp_configuration = {
         "nlp_engine_name": "transformers",
         "models": [
             {
-                "lang_code": "en",
-                "model_name": {"spacy": "en_core_web_sm", "transformers": model_path},
+                "lang_code": "cs",
+                "model_name": {"spacy": "cs_core_news_sm", "transformers": model_path},
             }
         ],
         "ner_model_configuration": {
             "model_to_presidio_entity_mapping": {
-                "PER": "PERSON",
-                "PERSON": "PERSON",
-                "LOC": "LOCATION",
-                "LOCATION": "LOCATION",
-                "GPE": "LOCATION",
-                "ORG": "ORGANIZATION",
-                "ORGANIZATION": "ORGANIZATION",
-                "NORP": "NRP",
-                "AGE": "AGE",
-                "ID": "ID",
+                **CZECH_ENTITY_MAPPING,
+                "PERSON": "OSOBA",
+                "LOC": "ADRESA",
+                "GPE": "ADRESA",
+                "ORG": "ORGANIZACE",
+                "DATE": "DATUM_NAROZENÍ",
+                "CARDINAL": "RODNÉ_ČÍSLO",
+                "ID": "IČO",
                 "EMAIL": "EMAIL",
-                "PATIENT": "PERSON",
-                "STAFF": "PERSON",
-                "HOSP": "ORGANIZATION",
-                "PATORG": "ORGANIZATION",
-                "DATE": "DATE_TIME",
-                "TIME": "DATE_TIME",
-                "PHONE": "PHONE_NUMBER",
-                "HCW": "PERSON",
-                "HOSPITAL": "ORGANIZATION",
-                "FACILITY": "LOCATION",
+                "PHONE": "TELEFON",
             },
             "low_confidence_score_multiplier": 0.4,
             "low_score_entity_names": ["ID"],
             "labels_to_ignore": [
-                "CARDINAL",
                 "EVENT",
                 "LANGUAGE",
                 "LAW",
@@ -150,7 +134,6 @@ def create_nlp_engine_with_transformers(
 
     return nlp_engine, registry
 
-
 def create_nlp_engine_with_flair(
     model_path: str,
 ) -> Tuple[NlpEngine, RecognizerRegistry]:
@@ -165,15 +148,13 @@ def create_nlp_engine_with_flair(
     registry = RecognizerRegistry()
     registry.load_predefined_recognizers()
 
-    # there is no official Flair NlpEngine, hence we load it as an additional recognizer
-
-    if not spacy.util.is_package("en_core_web_sm"):
-        spacy.cli.download("en_core_web_sm")
-    # Using a small spaCy model + a Flair NER model
+    if not spacy.util.is_package("cs_core_news_sm"):
+        spacy.cli.download("cs_core_news_sm")
+    
     flair_recognizer = FlairRecognizer(model_path=model_path)
     nlp_configuration = {
         "nlp_engine_name": "spacy",
-        "models": [{"lang_code": "en", "model_name": "en_core_web_sm"}],
+        "models": [{"lang_code": "cs", "model_name": "cs_core_news_sm"}],
     }
     registry.add_recognizer(flair_recognizer)
     registry.remove_recognizer("SpacyRecognizer")
@@ -182,7 +163,6 @@ def create_nlp_engine_with_flair(
 
     return nlp_engine, registry
 
-
 def create_nlp_engine_with_azure_ai_language(ta_key: str, ta_endpoint: str):
     """
     Instantiate an NlpEngine with a TextAnalyticsWrapper and a small spaCy model.
@@ -204,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": "en", "model_name": "en_core_web_sm"}],
+        "models": [{"lang_code": "cs", "model_name": "cs_core_news_sm"}],
     }
 
     nlp_engine = NlpEngineProvider(nlp_configuration=nlp_configuration).create_engine()
@@ -212,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