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

1import json 

2from typing import List, Optional 

3 

4from spacy.tokens import Doc, Span 

5 

6 

7class NlpArtifacts: 

8 """ 

9 NlpArtifacts is an abstraction layer over the results of an NLP pipeline. 

10 

11 processing over a given text, it holds attributes such as entities, 

12 tokens and lemmas which can be used by any recognizer 

13 

14 :param entities: Identified entities 

15 :param tokens: Tokenized text 

16 :param tokens_indices: Indices of tokens 

17 :param lemmas: List of lemmas in text 

18 :param nlp_engine: NlpEngine object 

19 :param language: Text language 

20 :param scores: Entity confidence scores 

21 """ 

22 

23 def __init__( 

24 self, 

25 entities: List[Span], 

26 tokens: Doc, 

27 tokens_indices: List[int], 

28 lemmas: List[str], 

29 nlp_engine: "NlpEngine", # noqa: F821 

30 language: str, 

31 scores: Optional[List[float]] = None, 

32 ): 

33 self.entities = entities 

34 self.tokens = tokens 

35 self.lemmas = lemmas 

36 self.tokens_indices = tokens_indices 

37 self.keywords = self.set_keywords(nlp_engine, lemmas, language) 

38 self.nlp_engine = nlp_engine 

39 self.scores = scores if scores else [0.85] * len(entities) 

40 

41 @staticmethod 

42 def set_keywords( 

43 nlp_engine, 

44 lemmas: List[str], 

45 language: str, 

46 ) -> List[str]: 

47 """ 

48 Return keywords fpr text. 

49 

50 Extracts lemmas with certain conditions as keywords. 

51 """ 

52 if not nlp_engine: 

53 return [] 

54 keywords = [ 

55 k.lower() 

56 for k in lemmas 

57 if not nlp_engine.is_stopword(k, language) 

58 and not nlp_engine.is_punct(k, language) 

59 and k != "-PRON-" 

60 and k != "be" 

61 ] 

62 

63 # best effort, try even further to break tokens into sub tokens, 

64 # this can result in reducing false negatives 

65 keywords = [i.split(":") for i in keywords] 

66 

67 # splitting the list can, if happened, will result in list of lists, 

68 # we flatten the list 

69 keywords = [item for sublist in keywords for item in sublist] 

70 return keywords 

71 

72 def to_json(self) -> str: 

73 """Convert nlp artifacts to json.""" 

74 

75 return_dict = self.__dict__.copy() 

76 

77 # Ignore NLP engine as it's not serializable currently 

78 del return_dict["nlp_engine"] 

79 

80 # Converting spaCy tokens and spans to string as they are not serializable 

81 if "tokens" in return_dict: 

82 return_dict["tokens"] = [token.text for token in self.tokens] 

83 if "entities" in return_dict: 

84 return_dict["entities"] = [entity.text for entity in self.entities] 

85 if "scores" in return_dict: 

86 return_dict["scores"] = [float(score) for score in self.scores] 

87 

88 return json.dumps(return_dict)