Coverage for presidio_analyzer / predefined_recognizers / ner / gliner_recognizer.py: 90%

1import json 

2import logging 

3from typing import Dict, List, Optional 

4 

5from presidio_analyzer import ( 

6 AnalysisExplanation, 

7 LocalRecognizer, 

8 RecognizerResult, 

9) 

10from presidio_analyzer.chunkers import BaseTextChunker 

11from presidio_analyzer.nlp_engine import ( 

12 NerModelConfiguration, 

13 NlpArtifacts, 

14 device_detector, 

15) 

16 

17try: 

18 from gliner import GLiNER, GLiNERConfig 

19except ImportError: 

20 GLiNER = None 

21 GLiNERConfig = None 

22 

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

24 

25 

26class GLiNERRecognizer(LocalRecognizer): 

27 """GLiNER model based entity recognizer.""" 

28 

29 def __init__( 

30 self, 

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

32 name: str = "GLiNERRecognizer", 

33 supported_language: str = "en", 

34 version: str = "0.0.1", 

35 context: Optional[List[str]] = None, 

36 entity_mapping: Optional[Dict[str, str]] = None, 

37 model_name: str = "urchade/gliner_multi_pii-v1", 

38 flat_ner: bool = True, 

39 multi_label: bool = False, 

40 threshold: float = 0.30, 

41 map_location: Optional[str] = None, 

42 text_chunker: Optional[BaseTextChunker] = None, 

43 load_onnx_model: bool = False, 

44 onnx_model_file: str = "model.onnx", 

45 **model_kwargs, 

46 ): 

47 """GLiNER model based entity recognizer. 

48 

49 The model is based on the GLiNER library. 

50 

51 :param supported_entities: List of supported entities for this recognizer. 

52 If None, all entities in Presidio's default configuration will be used. 

53 see `NerModelConfiguration` 

54 :param name: Name of the recognizer 

55 :param supported_language: Language code to use for the recognizer 

56 :param version: Version of the recognizer 

57 :param context: N/A for this recognizer 

58 :param model_name: The name of the GLiNER model to load 

59 :param flat_ner: Whether to use flat NER or not (see GLiNER's documentation) 

60 :param multi_label: Whether to use multi-label classification or not 

61 (see GLiNER's documentation) 

62 :param threshold: The threshold for the model's output 

63 (see GLiNER's documentation) 

64 :param map_location: The device to use for the model. 

65 If None, will auto-detect GPU or use CPU. 

66 :param text_chunker: Custom text chunking strategy. If None, uses 

67 CharacterBasedTextChunker with default settings (chunk_size=250, 

68 chunk_overlap=50) 

69 :param load_onnx_model: Whether to load the model using ONNX Runtime. 

70 If True, uses ONNX Runtime backend which supports CPUs without AVX2. 

71 Requires onnxruntime to be installed. Default is False. 

72 :param onnx_model_file: The name of the ONNX model file to load. 

73 Only used when load_onnx_model is True. This is passed directly to 

74 GLiNER.from_pretrained(). GLiNER looks for this file in the model 

75 directory (downloaded or cached model path). Default is "model.onnx". 

76 :param model_kwargs: Additional keyword arguments to pass to 

77 GLiNER.from_pretrained(). This allows passing future parameters 

78 to the GLiNER model without explicit support in this recognizer. 

79 

80 

81 """ 

82 

83 if entity_mapping: 

84 if supported_entities: 

85 raise ValueError( 

86 "entity_mapping and supported_entities cannot be used together" 

87 ) 

88 

89 self.model_to_presidio_entity_mapping = entity_mapping 

90 else: 

91 if not supported_entities: 

92 logger.info( 

93 "No supported entities provided, " 

94 "using default entities from NerModelConfiguration" 

95 ) 

96 self.model_to_presidio_entity_mapping = ( 

97 NerModelConfiguration().model_to_presidio_entity_mapping 

98 ) 

99 else: 

100 self.model_to_presidio_entity_mapping = { 

101 entity: entity for entity in supported_entities 

102 } 

103 

104 logger.info(f"Using entity mapping {json.dumps(entity_mapping, indent=2)}") 

105 supported_entities = list(set(self.model_to_presidio_entity_mapping.values())) 

106 self.model_name = model_name 

107 

108 self.map_location = ( 

109 map_location 

110 if map_location is not None 

111 else device_detector.get_device() 

112 ) 

113 

114 self.flat_ner = flat_ner 

115 self.multi_label = multi_label 

116 self.threshold = threshold 

117 self.load_onnx_model = load_onnx_model 

118 self.onnx_model_file = onnx_model_file 

119 self.model_kwargs = model_kwargs 

120 

121 # Use provided chunker or default to in-house character-based chunker 

122 if text_chunker is not None: 

123 self.text_chunker = text_chunker 

124 else: 

125 from presidio_analyzer.chunkers import CharacterBasedTextChunker 

126 

127 self.text_chunker = CharacterBasedTextChunker( 

128 chunk_size=250, 

129 chunk_overlap=50, 

130 ) 

131 

132 self.gliner = None 

133 

134 super().__init__( 

135 supported_entities=supported_entities, 

136 name=name, 

137 supported_language=supported_language, 

138 version=version, 

139 context=context, 

140 ) 

141 

142 self.gliner_labels = list(self.model_to_presidio_entity_mapping.keys()) 

143 

144 def load(self) -> None: 

145 """Load the GLiNER model.""" 

146 if not GLiNER: 

147 raise ImportError("GLiNER is not installed. Please install it.") 

148 

149 self.gliner = GLiNER.from_pretrained( 

150 self.model_name, 

151 map_location=self.map_location, 

152 load_onnx_model=self.load_onnx_model, 

153 onnx_model_file=self.onnx_model_file, 

154 **self.model_kwargs, 

155 ) 

156 

157 def analyze( 

158 self, 

159 text: str, 

160 entities: List[str], 

161 nlp_artifacts: Optional[NlpArtifacts] = None, 

162 ) -> List[RecognizerResult]: 

163 """Analyze text to identify entities using a GLiNER model. 

164 

165 :param text: The text to be analyzed 

166 :param entities: The list of entities this recognizer is requested to return 

167 :param nlp_artifacts: N/A for this recognizer 

168 """ 

169 

170 # combine the input labels as this model allows for ad-hoc labels 

171 labels = self.__create_input_labels(entities) 

172 

173 # Process text with automatic chunking 

174 def predict_func(text: str) -> List[RecognizerResult]: 

175 # Get predictions from GLiNER (returns dicts) 

176 gliner_predictions = self.gliner.predict_entities( 

177 text=text, 

178 labels=labels, 

179 flat_ner=self.flat_ner, 

180 threshold=self.threshold, 

181 multi_label=self.multi_label, 

182 ) 

183 

184 # Convert dicts to RecognizerResult objects 

185 results = [] 

186 for pred in gliner_predictions: 

187 presidio_entity = self.model_to_presidio_entity_mapping.get( 

188 pred["label"], pred["label"] 

189 ) 

190 

191 # Filter by requested entities 

192 if entities and presidio_entity not in entities: 

193 continue 

194 

195 analysis_explanation = AnalysisExplanation( 

196 recognizer=self.name, 

197 original_score=pred["score"], 

198 textual_explanation=f"Identified as {presidio_entity} by GLiNER", 

199 ) 

200 

201 results.append( 

202 RecognizerResult( 

203 entity_type=presidio_entity, 

204 start=pred["start"], 

205 end=pred["end"], 

206 score=pred["score"], 

207 analysis_explanation=analysis_explanation, 

208 ) 

209 ) 

210 return results 

211 

212 predictions = self.text_chunker.predict_with_chunking( 

213 text=text, 

214 predict_func=predict_func, 

215 ) 

216 

217 return predictions 

218 

219 def __create_input_labels(self, entities): 

220 """Append the entities requested by the user to the list of labels if it's not there.""" # noqa: E501 

221 labels = list(self.gliner_labels) 

222 for entity in entities: 

223 if ( 

224 entity not in self.model_to_presidio_entity_mapping.values() 

225 and entity not in self.gliner_labels 

226 ): 

227 labels.append(entity) 

228 return labels