Update presidio_nlp_engine_config.py

presidio_nlp_engine_config.py

@@ -1,5 +1,6 @@
 import logging
 from typing import Tuple
+import os
 import spacy
 from presidio_analyzer import RecognizerRegistry
 from presidio_analyzer.nlp_engine import (
@@ -8,6 +9,7 @@ from presidio_analyzer.nlp_engine import (
 )
 from transformers import AutoTokenizer, AutoModelForTokenClassification
 from presidio_analyzer.nlp_engine import TransformersNlpEngine
+from huggingface_hub import login
 
 logger = logging.getLogger("presidio-streamlit")
 
@@ -21,7 +23,7 @@ def create_nlp_engine_with_spacy(
     nlp = spacy.load(model_path)
     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",
@@ -56,37 +58,16 @@ def create_nlp_engine_with_transformers(
     """
     print(f"Loading Transformers model: {model_path} of type {type(model_path)}")
     
-    if model_path == "iiiorg/piiranha-v1-detect-personal-information":
-        # Specific configuration for your model
-        tokenizer = AutoTokenizer.from_pretrained(model_path)
-        model = AutoModelForTokenClassification.from_pretrained(model_path)
-        
-        nlp_engine = TransformersNlpEngine(tokenizer=tokenizer, model=model, device="cpu")
-        
-        # You might want to add specific entity mappings for this model
-        entity_mapping = {
-            "PERSON": "PERSON",
-            "LOCATION": "LOCATION",
-            "ORGANIZATION": "ORGANIZATION",
-            # Add more mappings as needed
-        }
-        
-        registry = RecognizerRegistry()
-        registry.load_predefined_recognizers(nlp_engine=nlp_engine)
-        
-        # You might want to add custom recognizers for this model
-        # For example:
-        # from presidio_analyzer import EntityRecognizer
-        # custom_recognizer = EntityRecognizer(supported_entities=["PERSON", "LOCATION", "ORGANIZATION"])
-        # registry.add_recognizer(custom_recognizer)
-        
-    else:
-        # Default configuration for other transformer models
-        tokenizer = AutoTokenizer.from_pretrained(model_path)
-        model = AutoModelForTokenClassification.from_pretrained(model_path)
-        
-        nlp_engine = TransformersNlpEngine(tokenizer=tokenizer, model=model, device="cpu")
-        registry = RecognizerRegistry()
-        registry.load_predefined_recognizers(nlp_engine=nlp_engine)
+    hf_token = os.getenv("HUGGING_FACE_TOKEN")
+    if hf_token:
+        login(hf_token)
+    
+    tokenizer = AutoTokenizer.from_pretrained(model_path)
+    model = AutoModelForTokenClassification.from_pretrained(model_path)
+    
+    nlp_engine = TransformersNlpEngine(tokenizer=tokenizer, model=model, device="cpu")
+    
+    registry = RecognizerRegistry()
+    registry.load_predefined_recognizers(nlp_engine=nlp_engine)
     
     return nlp_engine, registry