Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,8 @@
from .yaml_recognizer_models import (
BaseRecognizerConfig,
CustomRecognizerConfig,
GLiNERRecognizerConfig,
HuggingFaceRecognizerConfig,
LanguageContextConfig,
PredefinedRecognizerConfig,
RecognizerRegistryConfig,
Expand All @@ -15,6 +17,8 @@
"ConfigurationValidator",
"BaseRecognizerConfig",
"CustomRecognizerConfig",
"GLiNERRecognizerConfig",
"HuggingFaceRecognizerConfig",
"LanguageContextConfig",
"PredefinedRecognizerConfig",
"RecognizerRegistryConfig",
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -176,6 +176,23 @@ class HuggingFaceRecognizerConfig(PredefinedRecognizerConfig):
)


class GLiNERRecognizerConfig(PredefinedRecognizerConfig):
"""Configuration specifically for GLiNER models."""

model_config = ConfigDict(extra="forbid")

model_name: Optional[str] = Field(None, description="GLiNER model name")
flat_ner: Optional[bool] = Field(None, description="Use flat NER")
multi_label: Optional[bool] = Field(
None, description="Use multi-label classification"
)
threshold: Optional[float] = Field(None, description="Confidence threshold")
map_location: Optional[str] = Field(None, description="Device (cpu/gpu/etc.)")
load_onnx_model: Optional[bool] = Field(None, description="Load ONNX model")
onnx_model_file: Optional[str] = Field(None, description="ONNX model file name")
entity_mapping: Optional[Dict[str, str]] = Field(None, description="Entity mapping")


class CustomRecognizerConfig(BaseRecognizerConfig):
"""Configuration for custom pattern-based recognizers."""

Expand Down Expand Up @@ -463,4 +480,5 @@ def validate_language_presence(self):
# This allows for modular expansion without polluting the base config
CONFIG_MODEL_MAP: Dict[str, Type[BaseModel]] = {
"HuggingFaceNerRecognizer": HuggingFaceRecognizerConfig,
"GLiNERRecognizer": GLiNERRecognizerConfig,
}
31 changes: 31 additions & 0 deletions presidio-analyzer/tests/test_yaml_recognizer_models.py
Original file line number Diff line number Diff line change
Expand Up @@ -627,6 +627,37 @@ def test_recognizer_registry_config_custom_name_with_hf_class():
assert recognizer.model_name == "TestModel/Ner"


def test_gliner_recognizer_config_model_name():
"""Test that GLiNERRecognizer model_name is preserved through config validation."""
from presidio_analyzer.input_validation.yaml_recognizer_models import (
GLiNERRecognizerConfig,
RecognizerRegistryConfig,
)

registry_config = {
"recognizers": [
{
"name": "GLiNERRecognizer",
"type": "predefined",
"supported_language": "en",
"model_name": "custom/gliner-model",
"threshold": 0.5,
"flat_ner": False,
"multi_label": True,
}
]
}

config = RecognizerRegistryConfig(**registry_config)
recognizer = config.recognizers[0]

assert isinstance(recognizer, GLiNERRecognizerConfig)
assert recognizer.model_name == "custom/gliner-model"
assert recognizer.threshold == 0.5
assert recognizer.flat_ner is False
assert recognizer.multi_label is True


def test_config_model_map_fallback_to_predefined():
"""Test CONFIG_MODEL_MAP falls back to Predefined for unknown class_name."""
from presidio_analyzer.input_validation.yaml_recognizer_models import (
Expand Down