Coverage for presidio_analyzer / predefined_recognizers / ner / medical_ner_recognizer.py: 100%

1"""Medical NER recognizer using HuggingFace pipeline directly.""" 

2 

3from typing import Dict, List, Optional, Union 

4 

5from presidio_analyzer.chunkers import BaseTextChunker 

6from presidio_analyzer.predefined_recognizers.ner.huggingface_ner_recognizer import ( 

7 HuggingFaceNerRecognizer, 

8) 

9 

10DEFAULT_MEDICAL_ENTITY_MAPPING: Dict[str, str] = { 

11 "DISEASE_DISORDER": "MEDICAL_DISEASE_DISORDER", 

12 "MEDICATION": "MEDICAL_MEDICATION", 

13 "THERAPEUTIC_PROCEDURE": "MEDICAL_THERAPEUTIC_PROCEDURE", 

14 "CLINICAL_EVENT": "MEDICAL_CLINICAL_EVENT", 

15 "BIOLOGICAL_ATTRIBUTE": "MEDICAL_BIOLOGICAL_ATTRIBUTE", 

16 "BIOLOGICAL_STRUCTURE": "MEDICAL_BIOLOGICAL_STRUCTURE", 

17 "FAMILY_HISTORY": "MEDICAL_FAMILY_HISTORY", 

18 "HISTORY": "MEDICAL_HISTORY", 

19} 

20 

21 

22class MedicalNERRecognizer(HuggingFaceNerRecognizer): 

23 """Recognize medical/clinical entities using blaze999/Medical-NER. 

24 

25 Thin subclass of :class:`HuggingFaceNerRecognizer` that sets 

26 medical-specific defaults. Uses HuggingFace ``transformers.pipeline`` 

27 directly (no spaCy dependency). 

28 """ 

29 

30 ENTITIES = list(DEFAULT_MEDICAL_ENTITY_MAPPING.values()) 

31 

32 def __init__( 

33 self, 

34 model_name: str = "blaze999/Medical-NER", 

35 label_mapping: Optional[Dict[str, str]] = None, 

36 supported_entities: Optional[List[str]] = None, 

37 name: str = "MedicalNERRecognizer", 

38 supported_language: str = "en", 

39 aggregation_strategy: str = "simple", 

40 threshold: float = 0.3, 

41 device: Optional[Union[str, int]] = None, 

42 text_chunker: Optional[BaseTextChunker] = None, 

43 ): 

44 """Initialize the Medical NER recognizer. 

45 

46 :param model_name: HuggingFace model name/path. 

47 Default: ``blaze999/Medical-NER`` 

48 :param label_mapping: Model label -> Presidio entity mapping. 

49 Default: :data:`DEFAULT_MEDICAL_ENTITY_MAPPING` 

50 :param supported_entities: Entity types to return (None = all mapped). 

51 :param name: Recognizer name 

52 :param supported_language: Language code 

53 :param aggregation_strategy: Pipeline aggregation strategy 

54 :param threshold: Minimum confidence score (0.0 - 1.0) 

55 :param device: Device string/int (None = auto-detect) 

56 :param text_chunker: Custom text chunker (None = default) 

57 """ 

58 super().__init__( 

59 model_name=model_name, 

60 label_mapping=label_mapping or DEFAULT_MEDICAL_ENTITY_MAPPING, 

61 supported_entities=supported_entities, 

62 name=name, 

63 supported_language=supported_language, 

64 aggregation_strategy=aggregation_strategy, 

65 threshold=threshold, 

66 device=device, 

67 text_chunker=text_chunker, 

68 )