Coverage for presidio_analyzer / nlp_engine / stanza_nlp_engine.py: 91%
218 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-03-29 09:03 +0000
1import logging
2import warnings
3from typing import Any, Dict, Generator, List, Optional, Tuple, Union
5try:
6 import stanza
7 from stanza import Pipeline
8 from stanza.models.common.pretrain import Pretrain
9 from stanza.models.common.vocab import UNK_ID
10 from stanza.resources.common import DEFAULT_MODEL_DIR
12except ImportError:
13 stanza = None
16from spacy import Language, blank
17from spacy.tokens import Doc, Token
18from spacy.util import registry
20from presidio_analyzer.nlp_engine import (
21 NerModelConfiguration,
22 NlpArtifacts,
23 SpacyNlpEngine,
24 device_detector,
25)
27logger = logging.getLogger("presidio-analyzer")
30class StanzaNlpEngine(SpacyNlpEngine):
31 """
32 StanzaNlpEngine is an abstraction layer over the nlp module.
34 It provides processing functionality as well as other queries
35 on tokens.
36 The StanzaNlpEngine uses spacy-stanza and stanza as its NLP module
38 :param models: Dictionary with the name of the spaCy model per language.
39 For example: models = [{"lang_code": "en", "model_name": "en"}]
40 :param ner_model_configuration: Parameters for the NER model.
41 See conf/stanza.yaml for an example
43 """
45 engine_name = "stanza"
46 is_available = bool(stanza)
48 def __init__(
49 self,
50 models: Optional[List[Dict[str, str]]] = None,
51 ner_model_configuration: Optional[NerModelConfiguration] = None,
52 download_if_missing: bool = True,
53 ):
54 super().__init__(models, ner_model_configuration)
55 self.download_if_missing = download_if_missing
56 self.device = device_detector.get_device()
58 def load(self) -> None:
59 """Load the NLP model."""
61 logger.debug(f"Loading Stanza models: {self.models}")
63 # Enable GPU support using parent class method
64 super()._enable_gpu()
66 self.nlp = {}
67 for model in self.models:
68 self._validate_model_params(model)
69 self.nlp[model["lang_code"]] = load_pipeline(
70 model["model_name"],
71 processors="tokenize,pos,lemma,ner",
72 download_method="DOWNLOAD_RESOURCES"
73 if self.download_if_missing
74 else None,
75 device=self.device,
76 )
78 def process_batch(
79 self,
80 texts: Union[List[str], List[Tuple[str, object]]],
81 language: str,
82 batch_size: int = 1,
83 n_process: int = 1,
84 as_tuples: bool = False,
85 ) -> Generator[
86 Union[Tuple[Any, NlpArtifacts, Any], Tuple[Any, NlpArtifacts]], Any, None
87 ]:
88 """Execute the NLP pipeline on a batch of texts using Stanza's bulk processing.
90 This method overrides SpacyNlpEngine.process_batch to leverage Stanza's
91 efficient bulk_process method, which processes multiple documents together
92 for better GPU utilization.
94 Note: Stanza batches internally at the sentence/token level, not docs.
95 For optimal GPU performance, use larger batch sizes (e.g., 16-32 docs).
96 GPU utilization depends on total sentences/tokens across all docs in batch.
98 :param texts: A list of texts to process. if as_tuples is set to True,
99 texts should be a list of tuples (text, context).
100 :param language: The language of the texts.
101 :param batch_size: Number of documents per bulk_process call.
102 Recommended: 16-32+ for GPU, lower values acceptable for CPU.
103 :param n_process: Not used for Stanza (kept for API compatibility).
104 :param as_tuples: If set to True, inputs should be a sequence of
105 (text, context) tuples. Output will then be a sequence of
106 (text, NlpArtifacts, context) tuples. Defaults to False.
108 :return: A generator of tuples (text, NlpArtifacts, context) or
109 (text, NlpArtifacts) depending on the value of as_tuples.
110 """
112 if not self.nlp:
113 raise ValueError("NLP engine is not loaded. Consider calling .load()")
115 # Get the StanzaTokenizer (which wraps the Stanza pipeline)
116 # In spaCy, tokenizers are accessed via .tokenizer, not .get_pipe()
117 stanza_tokenizer = self.nlp[language].tokenizer
118 stanza_pipeline = stanza_tokenizer.snlp
120 text_list = list(texts) if not isinstance(texts, list) else texts
122 for batch_start in range(0, len(text_list), batch_size):
123 batch_end = min(batch_start + batch_size, len(text_list))
124 batch = text_list[batch_start:batch_end]
126 if as_tuples:
127 batch_texts = [str(text) for text, context in batch]
128 contexts = [context for text, context in batch]
129 else:
130 batch_texts = [str(text) for text in batch]
131 contexts = None
133 # Create Stanza Document objects and process via bulk_process
134 # Stanza handles internal batching at sentence/token level
135 stanza_docs = [stanza.Document([], text=text) for text in batch_texts]
136 processed_stanza_docs = stanza_pipeline.bulk_process(stanza_docs)
138 # Convert processed Stanza docs to spaCy docs using spacy-stanza's logic
139 # We call _convert_doc() which reuses StanzaTokenizer's conversion path
140 for idx, processed_stanza_doc in enumerate(processed_stanza_docs):
141 spacy_doc = stanza_tokenizer._convert_doc(processed_stanza_doc)
142 nlp_artifacts = self._doc_to_nlp_artifact(spacy_doc, language)
144 if as_tuples:
145 yield batch_texts[idx], nlp_artifacts, contexts[idx]
146 else:
147 yield batch_texts[idx], nlp_artifacts
150# Code taken from https://github.com/explosion/spacy-stanza
151# Supports Stanza > 1.7.0
152def load_pipeline(
153 name: str,
154 *,
155 lang: str = "",
156 dir: Optional[str] = None,
157 package: str = "default",
158 processors: Union[dict, str] = None,
159 logging_level: Optional[Union[int, str]] = None,
160 verbose: Optional[bool] = None,
161 device: Optional[str] = None,
162 **kwargs,
163) -> Language:
164 """Create a blank nlp object for a given language code.
166 with a stanza pipeline as part of the tokenizer.
168 To use the default stanza pipeline with
169 the same language code, leave the tokenizer config empty. Otherwise, pass
170 in the stanza pipeline settings in config["nlp"]["tokenizer"].
172 name (str): The language code, e.g. "en".
173 lang: str = "",
174 dir: Optional[str] = None,
175 package: str = "default",
176 processors: Union[dict, str] = {},
177 logging_level: Optional[Union[int, str]] = None,
178 verbose: Optional[bool] = None,
179 device: Optional[str] = None,
180 **kwargs: Options for the individual stanza processors.
181 RETURNS (Language): The nlp object.
182 """
184 if not processors:
185 processors = {}
187 # Create an empty config skeleton
188 config = {"nlp": {"tokenizer": {"kwargs": {}}}}
189 if lang == "":
190 lang = name
191 # Set the stanza tokenizer
192 config["nlp"]["tokenizer"]["@tokenizers"] = "PipelineAsTokenizer.v1"
193 # Set the stanza options
194 config["nlp"]["tokenizer"]["lang"] = lang
195 config["nlp"]["tokenizer"]["dir"] = dir
196 config["nlp"]["tokenizer"]["package"] = package
197 config["nlp"]["tokenizer"]["processors"] = processors
198 config["nlp"]["tokenizer"]["logging_level"] = logging_level
199 config["nlp"]["tokenizer"]["verbose"] = verbose
200 config["nlp"]["tokenizer"]["device"] = device
201 config["nlp"]["tokenizer"]["kwargs"].update(kwargs)
202 return blank(name, config=config)
205@registry.tokenizers("PipelineAsTokenizer.v1")
206def create_tokenizer(
207 lang: str = "",
208 dir: Optional[str] = None,
209 package: str = "default",
210 processors: Union[dict, str] = None,
211 logging_level: Optional[Union[int, str]] = None,
212 verbose: Optional[bool] = None,
213 device: Optional[str] = None,
214 kwargs: dict = None,
215):
216 """Create a tokenizer factory for a given language code.
218 :param lang: The language code, e.g. "en".
219 :param dir: The model directory.
220 :param package: The model package.
221 :param processors: The processors to use.
222 :param logging_level: The logging level.
223 :param verbose: Whether to be verbose.
224 :param device: The device to use (e.g., "cpu", "cuda").
225 Note: MPS is currently not supported.
226 :param kwargs: Additional keyword arguments.
227 """
228 if not processors:
229 processors = {}
230 if not kwargs:
231 kwargs = {}
233 def tokenizer_factory(
234 nlp,
235 lang=lang,
236 dir=dir,
237 package=package,
238 processors=processors,
239 logging_level=logging_level,
240 verbose=verbose,
241 device=device,
242 kwargs=kwargs,
243 ) -> StanzaTokenizer:
244 if dir is None:
245 dir = DEFAULT_MODEL_DIR
246 snlp = Pipeline(
247 lang=lang,
248 dir=dir,
249 package=package,
250 processors=processors,
251 logging_level=logging_level,
252 verbose=verbose,
253 device=device,
254 **kwargs,
255 )
256 return StanzaTokenizer(
257 snlp,
258 nlp.vocab,
259 )
261 return tokenizer_factory
264# Code taken from https://github.com/explosion/spacy-stanza
265# Supports Stanza > 1.7.0
266class StanzaTokenizer(object):
267 """The entire stanza pipeline in a custom tokenizer.
269 Because we're only running the Stanza pipeline once and don't split
270 it up into spaCy pipeline components, we'll set all the attributes within
271 a custom tokenizer.
272 """
274 def __init__(self, snlp, vocab):
275 """Initialize the tokenizer.
277 snlp (stanza.Pipeline): The initialized Stanza pipeline.
278 vocab (spacy.vocab.Vocab): The vocabulary to use.
279 RETURNS (Tokenizer): The custom tokenizer.
280 """
281 self.snlp = snlp
282 self.vocab = vocab
283 self.svecs = self._find_embeddings(snlp)
285 def __call__(self, text):
286 """Convert a Stanza Doc to a spaCy Doc.
288 text (Unicode): The text to process.
289 RETURNS (spacy.tokens.Doc): The spaCy Doc object.
290 """
291 if not text:
292 return Doc(self.vocab)
293 elif text.isspace():
294 return Doc(self.vocab, words=[text], spaces=[False])
296 snlp_doc = self.snlp(text)
297 return self._convert_doc(snlp_doc)
299 def _convert_doc(self, snlp_doc):
300 """Convert a processed Stanza Document to a spaCy Doc.
302 This method contains the conversion logic separated from text processing,
303 allowing it to be called with already-processed Stanza documents.
305 :param snlp_doc: Processed Stanza Document
306 :return: spaCy Doc object
307 """
308 if not snlp_doc.text:
309 return Doc(self.vocab)
310 elif snlp_doc.text.isspace():
311 return Doc(self.vocab, words=[snlp_doc.text], spaces=[False])
313 text = snlp_doc.text
314 snlp_tokens, snlp_heads = self.__get_tokens_with_heads(snlp_doc)
315 pos = []
316 tags = []
317 morphs = []
318 deps = []
319 heads = []
320 lemmas = []
321 token_texts = [t.text for t in snlp_tokens]
322 is_aligned = True
323 try:
324 words, spaces = self.__get_words_and_spaces(token_texts, text)
325 except ValueError:
326 words = token_texts
327 spaces = [True] * len(words)
328 is_aligned = False
329 warnings.warn(
330 "Due to multiword token expansion or an alignment "
331 "issue, the original text has been replaced by space-separated "
332 "expanded tokens.",
333 stacklevel=4,
334 )
335 offset = 0
336 for i, word in enumerate(words):
337 if word.isspace() and (
338 i + offset >= len(snlp_tokens) or word != snlp_tokens[i + offset].text
339 ):
340 # insert a space token
341 pos.append("SPACE")
342 tags.append("_SP")
343 morphs.append("")
344 deps.append("")
345 lemmas.append(word)
347 # increment any heads left of this position that point beyond
348 # this position to the right (already present in heads)
349 for j in range(0, len(heads)):
350 if j + heads[j] >= i:
351 heads[j] += 1
353 # decrement any heads right of this position that point beyond
354 # this position to the left (yet to be added from snlp_heads)
355 for j in range(i + offset, len(snlp_heads)):
356 if j + snlp_heads[j] < i + offset:
357 snlp_heads[j] -= 1
359 # initial space tokens are attached to the following token,
360 # otherwise attach to the preceding token
361 if i == 0:
362 heads.append(1)
363 else:
364 heads.append(-1)
366 offset -= 1
367 else:
368 token = snlp_tokens[i + offset]
369 assert word == token.text
371 pos.append(token.upos or "")
372 tags.append(token.xpos or token.upos or "")
373 morphs.append(token.feats or "")
374 deps.append(token.deprel or "")
375 heads.append(snlp_heads[i + offset])
376 lemmas.append(token.lemma or "")
378 doc = Doc(
379 self.vocab,
380 words=words,
381 spaces=spaces,
382 pos=pos,
383 tags=tags,
384 morphs=morphs,
385 lemmas=lemmas,
386 deps=deps,
387 heads=[head + i for i, head in enumerate(heads)],
388 )
389 ents = []
390 for ent in snlp_doc.entities:
391 ent_span = doc.char_span(ent.start_char, ent.end_char, ent.type)
392 ents.append(ent_span)
393 if not is_aligned or not all(ents):
394 warnings.warn(
395 f"Can't set named entities because of multi-word token "
396 f"expansion or because the character offsets don't map to "
397 f"valid tokens produced by the Stanza tokenizer:\n"
398 f"Words: {words}\n"
399 f"Entities: {[(e.text, e.type, e.start_char, e.end_char) for e in snlp_doc.entities]}", # noqa
400 stacklevel=4,
401 )
402 else:
403 doc.ents = ents
405 if self.svecs is not None:
406 doc.user_token_hooks["vector"] = self.token_vector
407 doc.user_token_hooks["has_vector"] = self.token_has_vector
408 return doc
410 def pipe(self, texts):
411 """Tokenize a stream of texts.
413 texts: A sequence of Unicode texts.
414 YIELDS (Doc): A sequence of Doc objects, in order.
415 """
416 for text in texts:
417 yield self(text)
419 @staticmethod
420 def __get_tokens_with_heads(snlp_doc):
421 """Flatten the tokens in the Stanza Doc and extract the token indices.
423 extract the token indices of the sentence start tokens to set is_sent_start.
425 snlp_doc (stanza.Document): The processed Stanza doc.
426 RETURNS (list): The tokens (words).
427 """
428 tokens = []
429 heads = []
430 offset = 0
431 for sentence in snlp_doc.sentences:
432 for token in sentence.tokens:
433 for word in token.words:
434 # Here, we're calculating the absolute token index in the doc,
435 # then the *relative* index of the head, -1 for zero-indexed
436 # and if the governor is 0 (root), we leave it at 0
437 if word.head:
438 head = word.head + offset - len(tokens) - 1
439 else:
440 head = 0
441 heads.append(head)
442 tokens.append(word)
443 offset += sum(len(token.words) for token in sentence.tokens)
444 return tokens, heads
446 @staticmethod
447 def __get_words_and_spaces(words, text):
448 if "".join("".join(words).split()) != "".join(text.split()):
449 raise ValueError("Unable to align mismatched text and words.")
450 text_words = []
451 text_spaces = []
452 text_pos = 0
453 # normalize words to remove all whitespace tokens
454 norm_words = [word for word in words if not word.isspace()]
455 # align words with text
456 for word in norm_words:
457 try:
458 word_start = text[text_pos:].index(word)
459 except ValueError:
460 raise ValueError("Unable to align mismatched text and words.")
461 if word_start > 0:
462 text_words.append(text[text_pos : text_pos + word_start])
463 text_spaces.append(False)
464 text_pos += word_start
465 text_words.append(word)
466 text_spaces.append(False)
467 text_pos += len(word)
468 if text_pos < len(text) and text[text_pos] == " ":
469 text_spaces[-1] = True
470 text_pos += 1
471 if text_pos < len(text):
472 text_words.append(text[text_pos:])
473 text_spaces.append(False)
474 return text_words, text_spaces
476 def token_vector(self, token: Token):
477 """Get Stanza's pretrained word embedding for given token.
479 :param token: The token whose embedding will be returned
480 :return (np.ndarray[ndim=1, dtype='float32']): the embedding/vector.
481 token.vector.size > 0 if Stanza pipeline contains a processor with
482 embeddings, else token.vector.size == 0.
483 A 0-vector (origin) will be returned
484 when the token doesn't exist in snlp's pretrained embeddings.
485 """
486 unit_id = self.svecs.vocab.unit2id(token.text)
487 return self.svecs.emb[unit_id]
489 def token_has_vector(self, token):
490 """Check if the token exists as a unit in snlp's pretrained embeddings."""
491 return self.svecs.vocab.unit2id(token.text) != UNK_ID
493 @staticmethod
494 def _find_embeddings(snlp):
495 """Find pretrained word embeddings in any of a SNLP's processors.
497 RETURNS (Pretrain): Or None if no embeddings were found.
498 """
499 embs = None
500 for proc in snlp.processors.values():
501 if hasattr(proc, "pretrain") and isinstance(proc.pretrain, Pretrain):
502 embs = proc.pretrain
503 break
504 return embs
506 # dummy serialization methods
507 def to_bytes(self, **kwargs): # noqa
508 return b""
510 def from_bytes(self, _bytes_data, **kwargs): # noqa
511 return self
513 def to_disk(self, _path, **kwargs): # noqa
514 return None
516 def from_disk(self, _path, **kwargs): # noqa
517 return self