Coverage for presidio_analyzer / nlp_engine / ner_model_configuration.py: 100%

1import logging 

2from typing import Collection, Dict, Optional 

3 

4from pydantic import BaseModel, ConfigDict, Field, field_validator 

5 

6logger = logging.getLogger("presidio-analyzer") 

7 

8MODEL_TO_PRESIDIO_ENTITY_MAPPING = dict( 

9 PER="PERSON", 

10 PERSON="PERSON", 

11 LOC="LOCATION", 

12 LOCATION="LOCATION", 

13 GPE="LOCATION", 

14 ORG="ORGANIZATION", 

15 DATE="DATE_TIME", 

16 TIME="DATE_TIME", 

17 NORP="NRP", 

18 AGE="AGE", 

19 ID="ID", 

20 EMAIL="EMAIL", 

21 PATIENT="PERSON", 

22 STAFF="PERSON", 

23 HOSP="ORGANIZATION", 

24 PATORG="ORGANIZATION", 

25 PHONE="PHONE_NUMBER", 

26 HCW="PERSON", 

27 HOSPITAL="ORGANIZATION", 

28) 

29 

30LOW_SCORE_ENTITY_NAMES = set() 

31 

32 

33class NerModelConfiguration(BaseModel): 

34 """NER model configuration using Pydantic validation. 

35 

36 :param labels_to_ignore: List of labels to not return predictions for. 

37 :param aggregation_strategy: 

38 See https://huggingface.co/docs/transformers/main_classes/pipelines#transformers.TokenClassificationPipeline.aggregation_strategy 

39 :param stride: 

40 See https://huggingface.co/docs/transformers/main_classes/pipelines#transformers.TokenClassificationPipeline.stride 

41 :param alignment_mode: See https://spacy.io/api/doc#char_span 

42 :param default_score: Default confidence score if the model does not provide one. 

43 :param model_to_presidio_entity_mapping: 

44 Mapping between the NER model entities and Presidio entities. 

45 :param low_score_entity_names: 

46 Set of entity names that are likely to have low detection accuracy that should be adjusted. 

47 :param low_confidence_score_multiplier: A multiplier for the score given for low_score_entity_names. 

48 Multiplier to the score given for low_score_entity_names. 

49 """ # noqa: E501 

50 

51 labels_to_ignore: Optional[Collection[str]] = Field( 

52 default_factory=list, description="List of labels to ignore" 

53 ) 

54 aggregation_strategy: Optional[str] = Field( 

55 default="max", description="Token classification aggregation strategy" 

56 ) 

57 stride: Optional[int] = Field( 

58 default=14, description="Stride for token classification" 

59 ) 

60 alignment_mode: Optional[str] = Field( 

61 default="expand", description="Alignment mode for spaCy char spans" 

62 ) 

63 default_score: Optional[float] = Field( 

64 default=0.85, ge=0.0, le=1.0, description="Default confidence score" 

65 ) 

66 model_to_presidio_entity_mapping: Optional[Dict[str, str]] = Field( 

67 default_factory=lambda: MODEL_TO_PRESIDIO_ENTITY_MAPPING.copy(), 

68 description="Mapping between model entities and Presidio entities", 

69 ) 

70 low_score_entity_names: Optional[Collection[str]] = Field( 

71 default_factory=lambda: LOW_SCORE_ENTITY_NAMES.copy(), 

72 description="Entity names with likely low detection accuracy", 

73 ) 

74 low_confidence_score_multiplier: Optional[float] = Field( 

75 default=0.4, ge=0.0, description="Score multiplier for low confidence entities" 

76 ) 

77 

78 model_config = ConfigDict(arbitrary_types_allowed=True) 

79 

80 @field_validator("aggregation_strategy") 

81 @classmethod 

82 def validate_aggregation_strategy(cls, agg_strategy: str) -> str: 

83 """Validate aggregation strategy.""" 

84 valid_strategies = ["simple", "first", "average", "max"] 

85 if agg_strategy not in valid_strategies: 

86 logger.warning( 

87 f"Aggregation strategy '{agg_strategy}' might not be supported. " 

88 f"Valid options: {valid_strategies}" 

89 ) 

90 return agg_strategy 

91 

92 @field_validator("stride") 

93 @classmethod 

94 def validate_stride(cls, stride: Optional[int]) -> int: 

95 """Validate stride and handle None values.""" 

96 if stride is None: 

97 # Get the default value from the field definition 

98 return cls.model_fields["stride"].default 

99 return stride 

100 

101 @field_validator("alignment_mode") 

102 @classmethod 

103 def validate_alignment_mode(cls, alignment: Optional[str]) -> str: 

104 """Validate alignment mode and handle None values.""" 

105 if alignment is None: 

106 # Get the default value from the field definition 

107 return cls.model_fields["alignment_mode"].default 

108 valid_modes = ["strict", "contract", "expand"] 

109 if alignment not in valid_modes: 

110 logger.warning( 

111 f"Alignment mode '{alignment}' might not be supported. " 

112 f"Valid options: {valid_modes}" 

113 ) 

114 return alignment 

115 

116 @classmethod 

117 def from_dict(cls, ner_model_configuration_dict: Dict) -> "NerModelConfiguration": 

118 """ 

119 Create NerModelConfiguration from a dictionary with Pydantic validation. 

120 

121 :param ner_model_configuration_dict: Dictionary containing configuration 

122 :return: NerModelConfiguration instance 

123 """ 

124 return cls(**ner_model_configuration_dict) 

125 

126 def to_dict(self) -> Dict: 

127 """Convert to dictionary representation.""" 

128 return self.model_dump(exclude_none=True)