Coverage for presidio_analyzer / nlp_engine / spacy_nlp_engine.py: 92%
119 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-03-29 09:03 +0000
1import logging
2from pathlib import Path
3from typing import Any, Dict, Generator, List, Optional, Tuple, Union
5import spacy
6from spacy.language import Language
7from spacy.tokens import Doc, Span
9from presidio_analyzer.nlp_engine import (
10 NerModelConfiguration,
11 NlpArtifacts,
12 NlpEngine,
13 device_detector,
14)
16logger = logging.getLogger("presidio-analyzer")
19class SpacyNlpEngine(NlpEngine):
20 """
21 SpacyNlpEngine is an abstraction layer over the nlp module.
23 It provides processing functionality as well as other queries
24 on tokens.
25 The SpacyNlpEngine uses SpaCy as its NLP module
26 """
28 engine_name = "spacy"
29 is_available = bool(spacy)
31 def __init__(
32 self,
33 models: Optional[List[Dict[str, str]]] = None,
34 ner_model_configuration: Optional[NerModelConfiguration] = None,
35 ):
36 """
37 Initialize a wrapper on spaCy functionality.
39 :param models: Dictionary with the name of the spaCy model per language.
40 For example: models = [{"lang_code": "en", "model_name": "en_core_web_lg"}]
41 :param ner_model_configuration: Parameters for the NER model.
42 See conf/spacy.yaml for an example
43 """
44 if not models:
45 models = [{"lang_code": "en", "model_name": "en_core_web_lg"}]
46 self.models = models
48 if not ner_model_configuration:
49 ner_model_configuration = NerModelConfiguration()
50 self.ner_model_configuration = ner_model_configuration
52 self.nlp = None
54 def _enable_gpu(self) -> None:
55 """Enable GPU support for spaCy/transformers if available."""
56 device = device_detector.get_device()
57 if device != "cpu":
58 try:
59 spacy.require_gpu()
60 except Exception as e:
61 logger.warning(
62 f"Failed to enable GPU ({device}), falling back to CPU: {e}"
63 )
65 def load(self) -> None:
66 """Load the spaCy NLP model."""
67 logger.debug(f"Loading SpaCy models: {self.models}")
69 self._enable_gpu()
71 self.nlp = {}
72 for model in self.models:
73 self._validate_model_params(model)
74 self._download_spacy_model_if_needed(model["model_name"])
75 self.nlp[model["lang_code"]] = spacy.load(model["model_name"])
77 @staticmethod
78 def _download_spacy_model_if_needed(model_name: str) -> None:
79 if not (spacy.util.is_package(model_name) or Path(model_name).exists()):
80 logger.warning(f"Model {model_name} is not installed. Downloading...")
81 spacy.cli.download(model_name)
82 logger.info(f"Finished downloading model {model_name}")
84 @staticmethod
85 def _validate_model_params(model: Dict) -> None:
86 if "lang_code" not in model:
87 raise ValueError("lang_code is missing from model configuration")
88 if "model_name" not in model:
89 raise ValueError("model_name is missing from model configuration")
90 if not isinstance(model["model_name"], str):
91 raise ValueError("model_name must be a string")
93 def get_supported_entities(self) -> List[str]:
94 """Return the supported entities for this NLP engine."""
95 if not self.ner_model_configuration.model_to_presidio_entity_mapping:
96 raise ValueError(
97 "model_to_presidio_entity_mapping is missing from model configuration"
98 )
99 entities_from_mapping = list(
100 set(self.ner_model_configuration.model_to_presidio_entity_mapping.values())
101 )
102 entities = [
103 ent
104 for ent in entities_from_mapping
105 if ent not in self.ner_model_configuration.labels_to_ignore
106 ]
107 return entities
109 def get_supported_languages(self) -> List[str]:
110 """Return the supported languages for this NLP engine."""
111 if not self.nlp:
112 raise ValueError("NLP engine is not loaded. Consider calling .load()")
113 return list(self.nlp.keys())
115 def is_loaded(self) -> bool:
116 """Return True if the model is already loaded."""
117 return self.nlp is not None
119 def process_text(self, text: str, language: str) -> NlpArtifacts:
120 """Execute the SpaCy NLP pipeline on the given text and language."""
121 if not self.nlp:
122 raise ValueError("NLP engine is not loaded. Consider calling .load()")
124 doc = self.nlp[language](text)
125 return self._doc_to_nlp_artifact(doc, language)
127 def process_batch(
128 self,
129 texts: Union[List[str], List[Tuple[str, object]]],
130 language: str,
131 batch_size: int = 1,
132 n_process: int = 1,
133 as_tuples: bool = False,
134 ) -> Generator[
135 Union[Tuple[Any, NlpArtifacts, Any], Tuple[Any, NlpArtifacts]], Any, None
136 ]:
137 """Execute the NLP pipeline on a batch of texts using spacy pipe.
139 :param texts: A list of texts to process. if as_tuples is set to True,
140 texts should be a list of tuples (text, context).
141 :param language: The language of the texts.
142 :param batch_size: Default batch size for pipe and evaluate.
143 :param n_process: Number of processors to process texts.
144 :param as_tuples: If set to True, inputs should be a sequence of
145 (text, context) tuples. Output will then be a sequence of
146 (doc, context) tuples. Defaults to False.
148 :return: A generator of tuples (text, NlpArtifacts, context) or
149 (text, NlpArtifacts) depending on the value of as_tuples.
150 """
152 if not self.nlp:
153 raise ValueError("NLP engine is not loaded. Consider calling .load()")
155 if as_tuples:
156 if not all(isinstance(item, tuple) and len(item) == 2 for item in texts):
157 raise ValueError(
158 "When 'as_tuples' is True, "
159 "'texts' must be a list of tuples (text, context)."
160 )
161 texts = ((str(text), context) for text, context in texts)
162 else:
163 texts = (str(text) for text in texts)
164 batch_output = self.nlp[language].pipe(
165 texts, as_tuples=as_tuples, batch_size=batch_size, n_process=n_process
166 )
167 for output in batch_output:
168 if as_tuples:
169 doc, context = output
170 yield doc.text, self._doc_to_nlp_artifact(doc, language), context
171 else:
172 doc = output
173 yield doc.text, self._doc_to_nlp_artifact(doc, language)
175 def is_stopword(self, word: str, language: str) -> bool:
176 """
177 Return true if the given word is a stop word.
179 (within the given language)
180 """
181 return self.nlp[language].vocab[word].is_stop
183 def is_punct(self, word: str, language: str) -> bool:
184 """
185 Return true if the given word is a punctuation word.
187 (within the given language).
188 """
189 return self.nlp[language].vocab[word].is_punct
191 def get_nlp(self, language: str) -> Language:
192 """
193 Return the language model loaded for a language.
195 :param language: Language
196 :return: Model from spaCy
197 """
198 return self.nlp[language]
200 def _doc_to_nlp_artifact(self, doc: Doc, language: str) -> NlpArtifacts:
201 lemmas = [token.lemma_ for token in doc]
202 tokens_indices = [token.idx for token in doc]
204 entities = self._get_entities(doc)
205 scores = self._get_scores_for_entities(doc)
207 entities, scores = self._get_updated_entities(entities, scores)
209 return NlpArtifacts(
210 entities=entities,
211 tokens=doc,
212 tokens_indices=tokens_indices,
213 lemmas=lemmas,
214 nlp_engine=self,
215 language=language,
216 scores=scores,
217 )
219 def _get_entities(self, doc: Doc) -> List[Span]:
220 """
221 Extract entities out of a spaCy pipeline, depending on the type of pipeline.
223 For normal spaCy, this would be doc.ents
224 :param doc: the output spaCy doc.
225 :return: List of entities
226 """
228 return doc.ents
230 def _get_scores_for_entities(self, doc: Doc) -> List[float]:
231 """Extract scores for entities from the doc.
233 Since spaCy does not provide confidence scores for entities by default,
234 we use the default score from the ner model configuration.
235 :param doc: SpaCy doc
236 """
238 entities = doc.ents
239 scores = [self.ner_model_configuration.default_score] * len(entities)
240 return scores
242 def _get_updated_entities(
243 self, entities: List[Span], scores: List[float]
244 ) -> Tuple[List[Span], List[float]]:
245 """
246 Get an updated list of entities based on the ner model configuration.
248 Remove entities that are in labels_to_ignore,
249 update entity names based on model_to_presidio_entity_mapping
251 :param entities: Entities that were extracted from a spaCy pipeline
252 :param scores: Original confidence scores for the entities extracted
253 :return: Tuple holding the entities and confidence scores
254 """
255 if len(entities) != len(scores):
256 raise ValueError("Entities and scores must be the same length")
258 new_entities = []
259 new_scores = []
261 mapping = self.ner_model_configuration.model_to_presidio_entity_mapping
262 to_ignore = self.ner_model_configuration.labels_to_ignore
263 for ent, score in zip(entities, scores):
264 # Remove model labels in the ignore list
265 if ent.label_ in to_ignore:
266 continue
268 # Update entity label based on mapping
269 if ent.label_ in mapping:
270 ent.label_ = mapping[ent.label_]
271 else:
272 logger.warning(
273 f"Entity {ent.label_} is not mapped to a Presidio entity, "
274 f"but keeping anyway. "
275 f"Add to `NerModelConfiguration.labels_to_ignore` to remove."
276 )
278 # Remove presidio entities in the ignore list
279 if ent.label_ in to_ignore:
280 continue
282 new_entities.append(ent)
284 # Update score if entity is in low score entity names
285 if ent.label_ in self.ner_model_configuration.low_score_entity_names:
286 score *= self.ner_model_configuration.low_confidence_score_multiplier
288 new_scores.append(score)
290 return new_entities, new_scores