Coverage for presidio_analyzer / nlp_engine / transformers_nlp_engine.py: 94%

1import logging 

2from typing import Dict, List, Optional 

3 

4import spacy 

5from spacy.tokens import Doc, Span 

6 

7try: 

8 import spacy_huggingface_pipelines 

9 import transformers 

10except ImportError: 

11 spacy_huggingface_pipelines = None 

12 transformers = None 

13 

14from presidio_analyzer.nlp_engine import ( 

15 NerModelConfiguration, 

16 SpacyNlpEngine, 

17) 

18 

19logger = logging.getLogger("presidio-analyzer") 

20 

21 

22class TransformersNlpEngine(SpacyNlpEngine): 

23 """ 

24 

25 TransformersNlpEngine is a transformers based NlpEngine. 

26 

27 It comprises a spacy pipeline used for tokenization, 

28 lemmatization, pos, and a transformers component for NER. 

29 

30 Both the underlying spacy pipeline and the transformers engine could be 

31 configured by the user. 

32 :param models: A dict holding the model's configuration. 

33 :example: 

34 [{"lang_code": "en", "model_name": { 

35 "spacy": "en_core_web_sm", 

36 "transformers": "dslim/bert-base-NER" 

37 } 

38 }] 

39 :param ner_model_configuration: Parameters for the NER model. 

40 See conf/transformers.yaml for an example 

41 

42 

43 Note that since the spaCy model is not used for NER, 

44 we recommend using a simple model, such as en_core_web_sm for English. 

45 For potential Transformers models, see a list of models here: 

46 https://huggingface.co/models?pipeline_tag=token-classification 

47 It is further recommended to fine-tune these models 

48 to the specific scenario in hand. 

49 

50 """ 

51 

52 engine_name = "transformers" 

53 is_available = bool(spacy_huggingface_pipelines) 

54 

55 def __init__( 

56 self, 

57 models: Optional[List[Dict]] = None, 

58 ner_model_configuration: Optional[NerModelConfiguration] = None, 

59 ): 

60 if not models: 

61 models = [ 

62 { 

63 "lang_code": "en", 

64 "model_name": { 

65 "spacy": "en_core_web_sm", 

66 "transformers": "obi/deid_roberta_i2b2", 

67 }, 

68 } 

69 ] 

70 super().__init__(models=models, ner_model_configuration=ner_model_configuration) 

71 self.entity_key = "bert-base-ner" 

72 

73 def load(self) -> None: 

74 """Load the spaCy and transformers models.""" 

75 

76 logger.debug(f"Loading SpaCy and transformers models: {self.models}") 

77 

78 super()._enable_gpu() 

79 

80 self.nlp = {} 

81 

82 for model in self.models: 

83 self._validate_model_params(model) 

84 spacy_model = model["model_name"]["spacy"] 

85 transformers_model = model["model_name"]["transformers"] 

86 self._download_spacy_model_if_needed(spacy_model) 

87 

88 nlp = spacy.load(spacy_model, disable=["parser", "ner"]) 

89 

90 pipe_config = { 

91 "model": transformers_model, 

92 "annotate": "spans", 

93 "stride": self.ner_model_configuration.stride, 

94 "alignment_mode": self.ner_model_configuration.alignment_mode, 

95 "aggregation_strategy": self.ner_model_configuration.aggregation_strategy, # noqa: E501 

96 "annotate_spans_key": self.entity_key, 

97 } 

98 

99 nlp.add_pipe("hf_token_pipe", config=pipe_config) 

100 self.nlp[model["lang_code"]] = nlp 

101 

102 @staticmethod 

103 def _validate_model_params(model: Dict) -> None: 

104 if "lang_code" not in model: 

105 raise ValueError("lang_code is missing from model configuration") 

106 if "model_name" not in model: 

107 raise ValueError("model_name is missing from model configuration") 

108 if not isinstance(model["model_name"], dict): 

109 raise ValueError("model_name must be a dictionary") 

110 if "spacy" not in model["model_name"]: 

111 raise ValueError("spacy model name is missing from model configuration") 

112 if "transformers" not in model["model_name"]: 

113 raise ValueError( 

114 "transformers model name is missing from model configuration" 

115 ) 

116 

117 def _get_entities(self, doc: Doc) -> List[Span]: 

118 """ 

119 Extract entities out of a spaCy pipeline, depending on the type of pipeline. 

120 

121 For spacy-huggingface-pipeline, this would be doc.spans[key] 

122 :param doc: the output spaCy doc. 

123 :return: List of entities 

124 """ 

125 

126 return doc.spans[self.entity_key] 

127 

128 def _get_scores_for_entities(self, doc: Doc) -> List[float]: 

129 """Extract scores for entities from the doc. 

130 

131 While spaCy does not provide confidence scores, 

132 the spacy-huggingface-pipeline flow adds confidence scores 

133 as SpanGroup attributes. 

134 :param doc: SpaCy doc 

135 """ 

136 

137 return [float(score) for score in doc.spans[self.entity_key].attrs["scores"]]