Coverage for presidio_analyzer / predefined_recognizers / ner / huggingface_ner_recognizer.py: 93%
137 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-03-29 09:03 +0000
1"""HuggingFace NER Recognizer for direct NER model inference.
3This recognizer uses HuggingFace Transformers pipeline directly,
4bypassing spaCy tokenizer alignment issues that commonly occur
5with agglutinative languages (Korean, Japanese, Turkish, etc.).
7The spaCy tokenizer tends to include particles/postpositions with nouns,
8causing alignment_mode issues:
9- NER result: "김태웅" (entity only)
10- spaCy token: "김태웅이고" (entity + particle)
11- char_span alignment fails
13This recognizer solves the problem by using the HuggingFace NER model's
14own tokenizer and returning results directly without spaCy alignment.
15"""
17import logging
18from typing import Any, Dict, List, Optional, Union
20from presidio_analyzer import (
21 AnalysisExplanation,
22 LocalRecognizer,
23 RecognizerResult,
24)
25from presidio_analyzer.chunkers import BaseTextChunker, CharacterBasedTextChunker
26from presidio_analyzer.nlp_engine import NlpArtifacts, device_detector
28try:
29 from transformers import pipeline as hf_pipeline
30except ImportError:
31 hf_pipeline = None
33try:
34 import torch
35except ImportError:
36 torch = None
39logger = logging.getLogger("presidio-analyzer")
42class HuggingFaceNerRecognizer(LocalRecognizer):
43 """HuggingFace Transformers based NER Recognizer.
45 This recognizer uses HuggingFace pipeline directly for NER,
46 bypassing spaCy tokenizer alignment issues. It's particularly
47 useful for agglutinative languages (Korean, Japanese, Turkish, etc.)
48 where particles attach to nouns.
50 Unlike the standard TransformersNlpEngine approach, this recognizer:
51 - Uses HuggingFace pipeline directly (not through spaCy)
52 - Returns NER results without char_span alignment
53 - Supports any HuggingFace token-classification model
54 - Works with any language that has a HuggingFace NER model
56 Example:
57 >>> from presidio_analyzer import AnalyzerEngine
58 >>> from presidio_analyzer.predefined_recognizers.ner import (
59 ... HuggingFaceNerRecognizer
60 ... )
61 >>>
62 >>> # For Korean
63 >>> recognizer = HuggingFaceNerRecognizer(
64 ... model_name="test-owner/test-model",
65 ... supported_language="ko"
66 ... )
67 >>>
68 >>> # For English
69 >>> recognizer = HuggingFaceNerRecognizer(
70 ... model_name="dslim/bert-base-NER",
71 ... supported_language="en"
72 ... )
73 >>>
74 >>> analyzer = AnalyzerEngine()
75 >>> analyzer.registry.add_recognizer(recognizer)
76 """
78 # Default label mapping from common NER models to Presidio entities
79 DEFAULT_LABEL_MAPPING = {
80 # Standard NER labels (CoNLL format)
81 "PER": "PERSON",
82 "LOC": "LOCATION",
83 "ORG": "ORGANIZATION",
84 "MISC": "MISC",
85 # Date/Time labels
86 "DAT": "DATE_TIME",
87 "DATE": "DATE_TIME",
88 "TIME": "DATE_TIME",
89 # Korean model labels (KoELECTRA, etc.)
90 "PS": "PERSON",
91 "LC": "LOCATION",
92 "OG": "ORGANIZATION",
93 "DT": "DATE_TIME",
94 "TI": "DATE_TIME",
95 # BERT-base-NER style labels (with B-/I- prefix stripped)
96 "PERSON": "PERSON",
97 "LOCATION": "LOCATION",
98 "ORGANIZATION": "ORGANIZATION",
99 "DATE_TIME": "DATE_TIME",
100 }
101 DEFAULT_HF_TASK = "token-classification"
103 def __init__(
104 self,
105 supported_entities: Optional[List[str]] = None,
106 name: str = "HuggingFaceNerRecognizer",
107 supported_language: str = "en",
108 version: str = "0.0.1",
109 context: Optional[List[str]] = None,
110 model_name: Optional[str] = None,
111 label_mapping: Optional[Dict[str, str]] = None,
112 threshold: float = 0.3,
113 aggregation_strategy: str = "simple",
114 chunk_overlap: int = 40,
115 chunk_size: int = 400,
116 device: Optional[Union[str, int]] = None,
117 tokenizer_name: Optional[str] = None,
118 text_chunker: Optional[BaseTextChunker] = None,
119 label_prefixes: Optional[List[str]] = None,
120 **kwargs,
121 ):
122 """Initialize the HuggingFace NER Recognizer.
124 :param supported_entities: List of supported entities.
125 If None, uses entities from label_mapping.
126 :param name: Name of the recognizer
127 :param supported_language: Language code (e.g., "en", "ko", "ja")
128 :param version: Version of the recognizer
129 :param context: Context words (N/A for this recognizer)
130 :param model_name: HuggingFace model name or path.
131 If None, must be set before calling load().
132 :param label_mapping: Mapping from model labels to Presidio entities.
133 If None, uses DEFAULT_LABEL_MAPPING.
134 :param threshold: Minimum confidence score threshold (0.0 - 1.0)
135 :param aggregation_strategy: Token aggregation strategy
136 ("simple", "first", "average", "max").
137 Recommendation: Use "simple" or "first" so that entities are pre-aggregated
138 by the model, preserving performance and alignment.
139 :param device: Device to use. Accepts:
140 - "cpu" or -1 for CPU
141 - "cuda" or "cuda:N" or int N for GPU
142 - None for auto-detection (GPU if available, else CPU)
143 Defaults to None.
144 :param chunk_overlap: Number of characters to overlap between chunks.
145 :param chunk_size: Maximum number of characters per chunk.
146 :param tokenizer_name: Name of the tokenizer. Defaults to model_name.
147 :param text_chunker: Custom text chunking strategy. If None, uses
148 CharacterBasedTextChunker with provided chunk_size and chunk_overlap.
149 :param label_prefixes: List of label prefixes to strip (e.g., B-, I-).
150 :raises ImportError: If transformers or torch libraries are not installed.
151 """
152 # Early check for required dependencies
153 if hf_pipeline is None:
154 raise ImportError(
155 "transformers is not installed. Please install it "
156 "(pip install transformers torch) to use this recognizer."
157 )
158 if torch is None:
159 raise ImportError(
160 "torch is not installed. Please install it "
161 "(pip install torch) to use this recognizer."
162 )
164 self.model_name = model_name
165 self.tokenizer_name = tokenizer_name or model_name
166 self.label_mapping = (
167 label_mapping if label_mapping is not None else self.DEFAULT_LABEL_MAPPING
168 )
169 self.threshold = threshold
170 self.aggregation_strategy = aggregation_strategy
171 if self.aggregation_strategy == "none":
172 logger.warning(
173 "aggregation_strategy='none' may result in fragmented entities "
174 "(e.g., 'B-PER', 'I-PER'). Recommended: 'simple' or 'first'."
175 )
176 self.device = self._parse_device(device)
177 self.label_prefixes = label_prefixes or ["B-", "I-", "U-", "L-"]
178 self.ner_pipeline = None
180 if kwargs:
181 logger.warning(
182 "Ignoring unsupported kwargs in %s: %s",
183 name,
184 sorted(kwargs.keys()),
185 )
187 # Derive supported entities from label mapping
188 if supported_entities:
189 final_supported_entities = supported_entities
190 else:
191 # Use sorted set for deterministic order
192 final_supported_entities = sorted(list(set(self.label_mapping.values())))
194 super().__init__(
195 supported_entities=final_supported_entities,
196 name=name,
197 supported_language=supported_language,
198 version=version,
199 context=context,
200 )
202 # Initialize the text chunker
203 if text_chunker:
204 self.text_chunker = text_chunker
205 else:
206 self.text_chunker = CharacterBasedTextChunker(
207 chunk_size=chunk_size,
208 chunk_overlap=chunk_overlap,
209 )
211 logger.info(
212 f"Initialized {self.name} with model={self.model_name}, "
213 f"threshold={self.threshold}, device={self.device}"
214 )
216 def _parse_device(self, device: Optional[Union[str, int]]) -> int:
217 """Parse device string or int into a transformer-compatible int.
219 Normalize diverse device inputs (None, "cpu", "cuda", "cuda:1", 0)
220 to a standard integer format for the Transformers pipeline.
221 If None, it uses Presidio's DeviceDetector for auto-detection.
222 """
223 if device is None:
224 detected = device_detector.get_device()
225 return 0 if detected == "cuda" else -1
227 if isinstance(device, int):
228 return device
230 device_str = str(device).strip().lower()
231 if device_str == "cpu":
232 return -1
233 if device_str == "cuda":
234 return 0
235 if device_str.startswith("cuda:"):
236 try:
237 return int(device_str.split(":", 1)[1])
238 except (ValueError, IndexError):
239 pass
241 raise ValueError(
242 f"Invalid device specified: {device}. "
243 "Accepts 'cpu', 'cuda', 'cuda:N', or integer index."
244 )
246 def load(self) -> None:
247 """Load the HuggingFace NER pipeline.
249 This method handles:
250 1. Hardware acceleration setup (CUDA validation and fallback)
251 2. Lazy-loading of the heavyweight ML pipeline.
253 :raises ValueError: If model_name is not set
254 """
255 if self.ner_pipeline is not None:
256 return
258 if not self.model_name:
259 raise ValueError(
260 "model_name must be set before calling load(). "
261 "Pass it to __init__() or set it directly."
262 )
264 # Device validation and fallback
265 device = self.device
266 if device >= 0:
267 if not torch.cuda.is_available():
268 logger.warning("CUDA is not available. Falling back to CPU.")
269 device = -1
270 elif device >= torch.cuda.device_count():
271 logger.warning(
272 "Device index %d out of range (count=%d). Falling back to CPU.",
273 device,
274 torch.cuda.device_count(),
275 )
276 device = -1
278 logger.info(f"Loading HuggingFace model: {self.model_name}, device={device}")
280 try:
281 self.ner_pipeline = hf_pipeline(
282 self.DEFAULT_HF_TASK,
283 model=self.model_name,
284 tokenizer=self.tokenizer_name,
285 aggregation_strategy=self.aggregation_strategy,
286 device=device,
287 )
288 logger.info(f"Successfully loaded {self.model_name}")
289 except Exception:
290 logger.exception(f"Failed to load model {self.model_name}")
291 raise
293 def _normalize_label(self, label: str) -> str:
294 """Normalize label by removing prefixes like B-/I-/U-/L-.
296 This method strips the prefix so that label_mapping can correctly
297 map to Presidio entities.
299 :param label: The raw label from the model.
300 :return: Normalized label.
301 """
302 if isinstance(label, str):
303 for prefix in self.label_prefixes:
304 if label.startswith(prefix):
305 return label[len(prefix) :]
306 return label
308 def _predict_chunk(self, chunk_text: str) -> List[RecognizerResult]:
309 """Perform NER prediction on a single text chunk.
311 This is a callback method used by the CharacterBasedTextChunker.
313 :param chunk_text: The chunk of text to analyze.
314 :return: List of RecognizerResult objects.
315 """
316 chunk_results = []
317 # Run inference on the chunk
318 try:
319 preds = self.ner_pipeline(chunk_text)
320 except Exception as e:
321 logger.warning(f"NER prediction failed for chunk: {e}", exc_info=True)
322 return []
324 # Helper to process a single prediction dictionary
325 def process_pred(pred: Dict[str, Any]) -> None:
326 """Convert a single HuggingFace prediction dict to RecognizerResult."""
327 raw_label = pred.get("entity_group") or pred.get("entity")
328 if not raw_label:
329 return
331 model_label = self._normalize_label(raw_label)
333 presidio_entity = self.label_mapping.get(model_label)
334 if not presidio_entity:
335 # If label is not mapped, use the model's label as is. This allows
336 # discovering entities not explicitly defined in the mapping.
337 presidio_entity = model_label
339 raw_score = pred.get("score", 0.0)
340 try:
341 score = float(raw_score)
342 except (TypeError, ValueError):
343 logger.warning("Failed to convert score to float: %r", raw_score)
344 return
346 if score < self.threshold:
347 return
349 start = pred.get("start")
350 end = pred.get("end")
351 if start is None or end is None:
352 return
354 if raw_label == model_label:
355 textual_explanation = (
356 f"Identified as {presidio_entity} by {self.model_name} "
357 f"(label: {raw_label})"
358 )
359 else:
360 textual_explanation = (
361 f"Identified as {presidio_entity} by {self.model_name} "
362 f"(original label: {raw_label}, normalized: {model_label})"
363 )
365 explanation = AnalysisExplanation(
366 recognizer=self.name,
367 original_score=score,
368 textual_explanation=textual_explanation,
369 )
371 chunk_results.append(
372 RecognizerResult(
373 entity_type=presidio_entity,
374 start=start,
375 end=end,
376 score=score,
377 analysis_explanation=explanation,
378 )
379 )
381 # Validate preds is a list before iterating
382 if not isinstance(preds, list):
383 logger.warning("Unexpected pipeline output type: %s", type(preds))
384 return []
386 for pred in preds:
387 if isinstance(pred, dict):
388 process_pred(pred)
389 else:
390 logger.warning("Unexpected prediction item type: %s", type(pred))
392 return chunk_results
394 def analyze(
395 self,
396 text: str,
397 entities: List[str],
398 nlp_artifacts: Optional[NlpArtifacts] = None,
399 ) -> List[RecognizerResult]:
400 """Analyze text for NER entities using HuggingFace model.
402 This method uses the CharacterBasedTextChunker to handle long texts
403 and ignores nlp_artifacts (spaCy results) to bypass tokenizer alignment issues.
405 :param text: The text to analyze
406 :param entities: List of entity types to detect
407 :param nlp_artifacts: Ignored (spaCy artifacts not used)
408 :return: List of RecognizerResult with detected entities
409 """
410 if not text or not text.strip():
411 return []
413 # Defensive guard for entities input
414 entities = entities or []
416 if not self.ner_pipeline:
417 self.load()
419 # Use the standard predict_with_chunking method from BaseTextChunker
420 # This handles chunking, processing, and duplicating/merging results
421 results = self.text_chunker.predict_with_chunking(
422 text=text,
423 predict_func=self._predict_chunk,
424 )
426 # Filter policy:
427 # 1. If an entity is requested, it is always kept.
428 # 2. If it's a known 'supported' entity but NOT requested, it is filtered out.
429 # 3. If it's an unmapped/unexpected entity, it is kept for Discovery.
430 if entities:
431 requested = set(entities)
432 supported = set(self.supported_entities)
433 results = [
434 r
435 for r in results
436 if (r.entity_type in requested) or (r.entity_type not in supported)
437 ]
439 return results