Coverage for presidio_analyzer / nlp_engine / nlp_engine_provider.py: 100%

1import logging 

2from pathlib import Path 

3from typing import Dict, Optional, Tuple, Union 

4 

5import yaml 

6 

7from presidio_analyzer.input_validation import ConfigurationValidator 

8from presidio_analyzer.nlp_engine import ( 

9 NerModelConfiguration, 

10 NlpEngine, 

11 SpacyNlpEngine, 

12 StanzaNlpEngine, 

13 TransformersNlpEngine, 

14) 

15 

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

17 

18 

19class NlpEngineProvider: 

20 """Create different NLP engines from configuration. 

21 

22 :param nlp_engines: List of available NLP engines. 

23 Default: (SpacyNlpEngine, StanzaNlpEngine) 

24 :param nlp_configuration: Dict containing nlp configuration 

25 :example: configuration: 

26 { 

27 "nlp_engine_name": "spacy", 

28 "models": [{"lang_code": "en", 

29 "model_name": "en_core_web_lg" 

30 }] 

31 } 

32 Nlp engine names available by default: spacy, stanza. 

33 :param conf_file: Path to yaml file containing nlp engine configuration. 

34 """ 

35 

36 def __init__( 

37 self, 

38 nlp_engines: Optional[Tuple] = None, 

39 conf_file: Optional[Union[Path, str]] = None, 

40 nlp_configuration: Optional[Dict] = None, 

41 ): 

42 if nlp_engines is None: 

43 nlp_engines = (SpacyNlpEngine, StanzaNlpEngine, TransformersNlpEngine) 

44 

45 self.nlp_engines = { 

46 engine.engine_name: engine for engine in nlp_engines if engine.is_available 

47 } 

48 logger.debug( 

49 f"Loaded these available nlp engines: {list(self.nlp_engines.keys())}" 

50 ) 

51 

52 if conf_file and nlp_configuration: 

53 raise ValueError( 

54 "Either conf_file or nlp_configuration should be provided, not both." 

55 ) 

56 

57 if nlp_configuration: 

58 ConfigurationValidator.validate_nlp_configuration(nlp_configuration) 

59 self.nlp_configuration = nlp_configuration 

60 

61 if conf_file or conf_file == "": 

62 if conf_file == "": 

63 raise ValueError("conf_file is empty") 

64 ConfigurationValidator.validate_file_path(conf_file) 

65 self.nlp_configuration = self._read_nlp_conf(conf_file) 

66 

67 if conf_file is None and nlp_configuration is None: 

68 conf_file = self._get_full_conf_path() 

69 logger.debug(f"Reading default conf file from {conf_file}") 

70 self.nlp_configuration = self._read_nlp_conf(conf_file) 

71 ConfigurationValidator.validate_nlp_configuration(self.nlp_configuration) 

72 

73 @staticmethod 

74 def _read_nlp_conf(conf_file: Union[Path, str]) -> Dict: 

75 """Read NLP configuration from a YAML file.""" 

76 with open(conf_file) as file: 

77 return yaml.safe_load(file) 

78 

79 

80 @staticmethod 

81 def _get_full_conf_path( 

82 default_conf_file: Union[Path, str] = "default.yaml" 

83 ) -> Path: 

84 """Return a Path to the default conf file.""" 

85 return Path(Path(__file__).parent, "../conf", default_conf_file) 

86 

87 def create_engine(self) -> NlpEngine: 

88 """Create an NLP engine instance.""" 

89 # Configuration is already validated by Pydantic in __init__ 

90 nlp_engine_name = self.nlp_configuration["nlp_engine_name"] 

91 if nlp_engine_name not in self.nlp_engines: 

92 raise ValueError( 

93 f"NLP engine '{nlp_engine_name}' is not available. " 

94 "Make sure you have all required packages installed" 

95 ) 

96 

97 nlp_engine_class = self.nlp_engines[nlp_engine_name] 

98 nlp_models = self.nlp_configuration["models"] 

99 

100 ner_model_configuration = self.nlp_configuration.get("ner_model_configuration") 

101 if ner_model_configuration: 

102 ner_model_configuration = NerModelConfiguration.from_dict( 

103 ner_model_configuration 

104 ) 

105 

106 engine = nlp_engine_class( 

107 models=nlp_models, ner_model_configuration=ner_model_configuration 

108 ) 

109 engine.load() 

110 logger.info( 

111 f"Created NLP engine: {engine.engine_name}. " 

112 f"Loaded models: {list(engine.nlp.keys())}" 

113 ) 

114 return engine