Coverage for presidio_analyzer / context_aware_enhancers / lemma_context_aware_enhancer.py: 92%
103 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-03-29 09:03 +0000
« prev ^ index » next coverage.py v7.13.1, created at 2026-03-29 09:03 +0000
1import copy
2import logging
3from typing import List, Optional
5from presidio_analyzer import EntityRecognizer, RecognizerResult
6from presidio_analyzer.context_aware_enhancers import ContextAwareEnhancer
7from presidio_analyzer.nlp_engine import NlpArtifacts
9logger = logging.getLogger("presidio-analyzer")
12class LemmaContextAwareEnhancer(ContextAwareEnhancer):
13 """
14 A class representing a lemma based context aware enhancer logic.
16 Context words might enhance confidence score of a recognized entity,
17 LemmaContextAwareEnhancer is an implementation of Lemma based context aware logic,
18 it compares spacy lemmas of each word in context of the matched entity to given
19 context and the recognizer context words,
20 if matched it enhance the recognized entity confidence score by a given factor.
22 :param context_similarity_factor: How much to enhance confidence of match entity
23 :param min_score_with_context_similarity: Minimum confidence score
24 :param context_prefix_count: how many words before the entity to match context
25 :param context_suffix_count: how many words after the entity to match context
26 :param context_matching_mode: Matching mode for context words. Options:
27 - "substring" (default): Match context words as substrings
28 (e.g., 'card' matches 'creditcard', 'lic' matches 'duplicate').
29 Maintains backward compatibility.
30 - "whole_word": Match context words only as whole words
31 (e.g., 'lic' matches 'lic' but not 'duplicate').
32 Prevents false positives.
33 """
35 def __init__(
36 self,
37 context_similarity_factor: float = 0.35,
38 min_score_with_context_similarity: float = 0.4,
39 context_prefix_count: int = 5,
40 context_suffix_count: int = 0,
41 context_matching_mode: str = "substring",
42 ):
43 super().__init__(
44 context_similarity_factor=context_similarity_factor,
45 min_score_with_context_similarity=min_score_with_context_similarity,
46 context_prefix_count=context_prefix_count,
47 context_suffix_count=context_suffix_count,
48 )
49 if context_matching_mode not in ["whole_word", "substring"]:
50 raise ValueError(
51 f"context_matching_mode must be one of: 'whole_word', 'substring'. "
52 f"Got: {context_matching_mode}"
53 )
54 self.context_matching_mode = context_matching_mode
56 def enhance_using_context(
57 self,
58 text: str,
59 raw_results: List[RecognizerResult],
60 nlp_artifacts: NlpArtifacts,
61 recognizers: List[EntityRecognizer],
62 context: Optional[List[str]] = None,
63 ) -> List[RecognizerResult]:
64 """
65 Update results in case the lemmas of surrounding words or input context
66 words are identical to the context words.
68 Using the surrounding words of the actual word matches, look
69 for specific strings that if found contribute to the score
70 of the result, improving the confidence that the match is
71 indeed of that PII entity type
73 :param text: The actual text that was analyzed
74 :param raw_results: Recognizer results which didn't take
75 context into consideration
76 :param nlp_artifacts: The nlp artifacts contains elements
77 such as lemmatized tokens for better
78 accuracy of the context enhancement process
79 :param recognizers: the list of recognizers
80 :param context: list of context words
81 """ # noqa: D205,D400
83 # create a deep copy of the results object, so we can manipulate it
84 results = copy.deepcopy(raw_results)
86 # create recognizer context dictionary
87 recognizers_dict = {recognizer.id: recognizer for recognizer in recognizers}
89 # Create empty list in None or lowercase all context words in the list
90 if not context:
91 context = []
92 else:
93 context = [word.lower() for word in context]
95 # Sanity
96 if nlp_artifacts is None:
97 logger.warning("NLP artifacts were not provided")
98 return results
100 for result in results:
101 recognizer = None
102 # get recognizer matching the result, if found.
103 if (
104 result.recognition_metadata
105 and RecognizerResult.RECOGNIZER_IDENTIFIER_KEY
106 in result.recognition_metadata.keys()
107 ):
108 recognizer = recognizers_dict.get(
109 result.recognition_metadata[
110 RecognizerResult.RECOGNIZER_IDENTIFIER_KEY
111 ]
112 )
114 if not recognizer:
115 logger.debug(
116 "Recognizer name not found as part of the "
117 "recognition_metadata dict in the RecognizerResult. "
118 )
119 continue
121 # skip recognizer result if the recognizer doesn't support
122 # context enhancement
123 if not recognizer.context:
124 logger.debug(
125 "recognizer '%s' does not support context enhancement",
126 recognizer.name,
127 )
128 continue
130 # skip context enhancement if already boosted by recognizer level
131 if result.recognition_metadata.get(
132 RecognizerResult.IS_SCORE_ENHANCED_BY_CONTEXT_KEY
133 ):
134 logger.debug("result score already boosted, skipping")
135 continue
137 # extract lemmatized context from the surrounding of the match
138 word = text[result.start : result.end]
140 surrounding_words = self._extract_surrounding_words(
141 nlp_artifacts=nlp_artifacts, word=word, start=result.start
142 )
144 # combine other sources of context with surrounding words
145 surrounding_words.extend(context)
147 supportive_context_word = self._find_supportive_word_in_context(
148 surrounding_words, recognizer.context, self.context_matching_mode
149 )
150 if supportive_context_word != "":
151 result.score += self.context_similarity_factor
152 result.score = max(result.score, self.min_score_with_context_similarity)
153 result.score = min(result.score, ContextAwareEnhancer.MAX_SCORE)
155 # Update the explainability object with context information
156 # helped to improve the score
157 result.analysis_explanation.set_supportive_context_word(
158 supportive_context_word
159 )
160 result.analysis_explanation.set_improved_score(result.score)
161 return results
163 @staticmethod
164 def _find_supportive_word_in_context(
165 context_list: List[str],
166 recognizer_context_list: List[str],
167 matching_mode: str = "substring",
168 ) -> str:
169 """
170 Find words in the text which are relevant for context evaluation.
172 A word is considered a supportive context word based on the matching mode:
173 - "substring" (default): Substring match
174 (e.g., 'card' matches 'creditcard', 'lic' matches 'duplicate')
175 - "whole_word": Exact whole-word match (case-insensitive)
176 (e.g., 'lic' matches 'lic' but not 'duplicate')
178 :param context_list words before and after the matched entity within
179 a specified window size
180 :param recognizer_context_list a list of words considered as
181 context keywords manually specified by the recognizer's author
182 :param matching_mode: Matching mode ('whole_word' or 'substring').
183 Defaults to 'substring'.
184 """
185 word = ""
186 # If the context list is empty, no need to continue
187 if context_list is None or recognizer_context_list is None:
188 return word
190 for predefined_context_word in recognizer_context_list:
191 result = False
193 if matching_mode == "substring":
194 # Substring match (case-insensitive) - default behavior
195 # for backward compatibility
196 result = next(
197 (
198 True
199 for keyword in context_list
200 if predefined_context_word.lower() in keyword.lower()
201 ),
202 False,
203 )
204 elif matching_mode == "whole_word":
205 # Exact whole-word match (case-insensitive)
206 result = next(
207 (
208 True
209 for keyword in context_list
210 if predefined_context_word.lower() == keyword.lower()
211 ),
212 False,
213 )
215 if result:
216 logger.debug("Found context keyword '%s'", predefined_context_word)
217 word = predefined_context_word
218 break
220 return word
222 def _extract_surrounding_words(
223 self, nlp_artifacts: NlpArtifacts, word: str, start: int
224 ) -> List[str]:
225 """Extract words surrounding another given word.
227 The text from which the context is extracted is given in the nlp
228 doc.
230 :param nlp_artifacts: An abstraction layer which holds different
231 items which are the result of a NLP pipeline
232 execution on a given text
233 :param word: The word to look for context around
234 :param start: The start index of the word in the original text
235 """
236 if not nlp_artifacts.tokens:
237 logger.info("Skipping context extraction due to lack of NLP artifacts")
238 # if there are no nlp artifacts, this is ok, we can
239 # extract context and we return a valid, yet empty
240 # context
241 return [""]
243 # Get the already prepared words in the given text, in their
244 # LEMMATIZED version
245 lemmatized_keywords = nlp_artifacts.keywords
247 # since the list of tokens is not necessarily aligned
248 # with the actual index of the match, we look for the
249 # token index which corresponds to the match
250 token_index = self._find_index_of_match_token(
251 word, start, nlp_artifacts.tokens, nlp_artifacts.tokens_indices
252 )
254 # index i belongs to the PII entity, take the preceding n words
255 # and the successing m words into a context list
257 backward_context = self._add_n_words_backward(
258 token_index,
259 self.context_prefix_count,
260 nlp_artifacts.lemmas,
261 lemmatized_keywords,
262 )
263 forward_context = self._add_n_words_forward(
264 token_index,
265 self.context_suffix_count,
266 nlp_artifacts.lemmas,
267 lemmatized_keywords,
268 )
270 context_list = []
271 context_list.extend(backward_context)
272 context_list.extend(forward_context)
273 context_list = list(set(context_list))
274 logger.debug("Context list is: %s", " ".join(context_list))
275 return context_list
277 @staticmethod
278 def _find_index_of_match_token(
279 word: str,
280 start: int,
281 tokens,
282 tokens_indices: List[int],
283 ) -> int:
284 found = False
285 # we use the known start index of the original word to find the actual
286 # token at that index, we are not checking for equivalence since the
287 # token might be just a substring of that word (e.g. for phone number
288 # 555-124564 the first token might be just '555' or for a match like '
289 # rocket' the actual token will just be 'rocket' hence the misalignment
290 # of indices)
291 # Note: we are iterating over the original tokens (not the lemmatized)
292 i = -1
293 for i, token in enumerate(tokens, 0):
294 # Either we found a token with the exact location, or
295 # we take a token which its characters indices covers
296 # the index we are looking for.
297 if (tokens_indices[i] == start) or (start < tokens_indices[i] + len(token)):
298 # found the interesting token, the one that around it
299 # we take n words, we save the matching lemma
300 found = True
301 break
303 if not found:
304 raise ValueError(
305 "Did not find word '" + word + "' "
306 "in the list of tokens although it "
307 "is expected to be found"
308 )
309 return i
311 @staticmethod
312 def _add_n_words(
313 index: int,
314 n_words: int,
315 lemmas: List[str],
316 lemmatized_filtered_keywords: List[str],
317 is_backward: bool,
318 ) -> List[str]:
319 """
320 Prepare a string of context words.
322 Return a list of words which surrounds a lemma at a given index.
323 The words will be collected only if exist in the filtered array
325 :param index: index of the lemma that its surrounding words we want
326 :param n_words: number of words to take
327 :param lemmas: array of lemmas
328 :param lemmatized_filtered_keywords: the array of filtered
329 lemmas from the original sentence,
330 :param is_backward: if true take the preceeding words, if false,
331 take the successing words
332 """
333 i = index
334 context_words = []
335 # The entity itself is no interest to us...however we want to
336 # consider it anyway for cases were it is attached with no spaces
337 # to an interesting context word, so we allow it and add 1 to
338 # the number of collected words
340 # collect at most n words (in lower case)
341 remaining = n_words + 1
342 while 0 <= i < len(lemmas) and remaining > 0:
343 lower_lemma = lemmas[i].lower()
344 if lower_lemma in lemmatized_filtered_keywords:
345 context_words.append(lower_lemma)
346 remaining -= 1
347 i = i - 1 if is_backward else i + 1
348 return context_words
350 def _add_n_words_forward(
351 self,
352 index: int,
353 n_words: int,
354 lemmas: List[str],
355 lemmatized_filtered_keywords: List[str],
356 ) -> List[str]:
357 return self._add_n_words(
358 index, n_words, lemmas, lemmatized_filtered_keywords, False
359 )
361 def _add_n_words_backward(
362 self,
363 index: int,
364 n_words: int,
365 lemmas: List[str],
366 lemmatized_filtered_keywords: List[str],
367 ) -> List[str]:
368 return self._add_n_words(
369 index, n_words, lemmas, lemmatized_filtered_keywords, True
370 )