Update presidio_helpers.py

presidio_helpers.py

@@ -9,9 +9,7 @@ from presidio_analyzer import (
     RecognizerResult,
     RecognizerRegistry,
     PatternRecognizer,
-    Pattern,
 )
-from presidio_analyzer.nlp_engine import NlpEngine
 from presidio_anonymizer import AnonymizerEngine
 from presidio_anonymizer.entities import OperatorConfig
 
@@ -21,14 +19,13 @@ logger = logging.getLogger("presidio-streamlit")
 def nlp_engine_and_registry(
     model_family: str,
     model_path: str,
-) -> Tuple[NlpEngine, RecognizerRegistry]:
+) -> Tuple[object, RecognizerRegistry]:
     """Create the NLP Engine instance based on the requested model."""
     registry = RecognizerRegistry()
     
-    if model_family.lower() == "spacy":
-        from spacy.language import Language
-        import spacy
-        try:
+    try:
+        if model_family.lower() == "spacy":
+            import spacy
             nlp = spacy.load(model_path)
             registry.load_predefined_recognizers()
             registry.add_recognizer_from_dict({
@@ -39,13 +36,8 @@ def nlp_engine_and_registry(
                 "package": "spacy",
             })
             return nlp, registry
-        except Exception as e:
-            logger.error(f"Failed to load spaCy model {model_path}: {str(e)}")
-            raise
-    elif model_family.lower() == "flair":
-        from flair.models import SequenceTagger
-        from flair.data import Sentence
-        try:
+        elif model_family.lower() == "flair":
+            from flair.models import SequenceTagger
             tagger = SequenceTagger.load(model_path)
             registry.load_predefined_recognizers()
             registry.add_recognizer_from_dict({
@@ -56,12 +48,8 @@ def nlp_engine_and_registry(
                 "package": "flair",
             })
             return tagger, registry
-        except Exception as e:
-            logger.error(f"Failed to load Flair model {model_path}: {str(e)}")
-            raise
-    elif model_family.lower() == "huggingface":
-        from transformers import pipeline
-        try:
+        elif model_family.lower() == "huggingface":
+            from transformers import pipeline
             nlp = pipeline("ner", model=model_path, tokenizer=model_path)
             registry.load_predefined_recognizers()
             registry.add_recognizer_from_dict({
@@ -72,11 +60,11 @@ def nlp_engine_and_registry(
                 "package": "transformers",
             })
             return nlp, registry
-        except Exception as e:
-            logger.error(f"Failed to load HuggingFace model {model_path}: {str(e)}")
-            raise
-    else:
-        raise ValueError(f"Model family {model_family} not supported")
+        else:
+            raise ValueError(f"Model family {model_family} not supported")
+    except Exception as e:
+        logger.error(f"Error loading model {model_path} for {model_family}: {str(e)}")
+        raise RuntimeError(f"Failed to load model: {str(e)}. Ensure model is downloaded and accessible.")
 
 @st.cache_resource
 def analyzer_engine(
@@ -110,24 +98,28 @@ def analyze(
     deny_list: List[str],
 ) -> List[RecognizerResult]:
     """Analyze text for PHI entities."""
-    results = analyzer.analyze(
-        text=text,
-        entities=entities,
-        language=language,
-        score_threshold=score_threshold,
-        return_decision_process=return_decision_process,
-    )
-    # Apply allow and deny lists
-    filtered_results = []
-    for result in results:
-        text_snippet = text[result.start:result.end].lower()
-        if any(word.lower() in text_snippet for word in allow_list):
-            continue
-        if any(word.lower() in text_snippet for word in deny_list):
-            filtered_results.append(result)
-        elif not deny_list:
-            filtered_results.append(result)
-    return filtered_results
+    try:
+        results = analyzer.analyze(
+            text=text,
+            entities=entities,
+            language=language,
+            score_threshold=score_threshold,
+            return_decision_process=return_decision_process,
+        )
+        # Apply allow and deny lists
+        filtered_results = []
+        for result in results:
+            text_snippet = text[result.start:result.end].lower()
+            if any(word.lower() in text_snippet for word in allow_list):
+                continue
+            if any(word.lower() in text_snippet for word in deny_list):
+                filtered_results.append(result)
+            elif not deny_list:
+                filtered_results.append(result)
+        return filtered_results
+    except Exception as e:
+        logger.error(f"Analysis error: {str(e)}")
+        raise
 
 def anonymize(
     text: str,
@@ -137,20 +129,24 @@ def anonymize(
     number_of_chars: int = 15,
 ) -> dict:
     """Anonymize detected PHI entities in the text."""
-    anonymizer = AnonymizerEngine()
-    operator_config = {
-        "DEFAULT": OperatorConfig(operator, {})
-    }
-    if operator == "mask":
-        operator_config["DEFAULT"] = OperatorConfig(operator, {
-            "masking_char": mask_char,
-            "chars_to_mask": number_of_chars,
-        })
-    return anonymizer.anonymize(
-        text=text,
-        analyzer_results=analyze_results,
-        operators=operator_config,
-    )
+    try:
+        anonymizer = AnonymizerEngine()
+        operator_config = {
+            "DEFAULT": OperatorConfig(operator, {})
+        }
+        if operator == "mask":
+            operator_config["DEFAULT"] = OperatorConfig(operator, {
+                "masking_char": mask_char,
+                "chars_to_mask": number_of_chars,
+            })
+        return anonymizer.anonymize(
+            text=text,
+            analyzer_results=analyze_results,
+            operators=operator_config,
+        )
+    except Exception as e:
+        logger.error(f"Anonymization error: {str(e)}")
+        raise
 
 def create_ad_hoc_deny_list_recognizer(
     deny_list: Optional[List[str]] = None,
@@ -158,7 +154,11 @@ def create_ad_hoc_deny_list_recognizer(
     """Create a recognizer for deny list items."""
     if not deny_list:
         return None
-    deny_list_recognizer = PatternRecognizer(
-        supported_entity="GENERIC_PII", deny_list=deny_list
-    )
-    return deny_list_recognizer
+    try:
+        deny_list_recognizer = PatternRecognizer(
+            supported_entity="GENERIC_PII", deny_list=deny_list
+        )
+        return deny_list_recognizer
+    except Exception as e:
+        logger.error(f"Error creating deny list recognizer: {str(e)}")
+        raise