Coverage for presidio_analyzer / chunkers / base_chunker.py: 94%

1"""Abstract base class for text chunking strategies.""" 

2from abc import ABC, abstractmethod 

3from dataclasses import dataclass 

4from typing import TYPE_CHECKING, Callable, List 

5 

6if TYPE_CHECKING: 

7 from presidio_analyzer import RecognizerResult 

8 

9 

10@dataclass 

11class TextChunk: 

12 """Represents a chunk of text with its position in the original text. 

13 

14 :param text: The chunk content 

15 :param start: Start position in the original text (inclusive) 

16 :param end: End position in the original text (exclusive) 

17 """ 

18 

19 text: str 

20 start: int 

21 end: int 

22 

23 

24class BaseTextChunker(ABC): 

25 """Abstract base class for text chunking strategies. 

26 

27 Subclasses must implement the chunk() method to split text into 

28 TextChunk objects that include both content and position information. 

29 

30 Provides methods for processing predictions across chunks and 

31 deduplicating overlapping entities. 

32 """ 

33 

34 @abstractmethod 

35 def chunk(self, text: str) -> List[TextChunk]: 

36 """Split text into chunks with position information. 

37 

38 :param text: The input text to split 

39 :return: List of TextChunk objects with text and position data 

40 """ 

41 pass 

42 

43 def predict_with_chunking( 

44 self, 

45 text: str, 

46 predict_func: Callable[[str], List["RecognizerResult"]], 

47 ) -> List["RecognizerResult"]: 

48 """Process text with automatic chunking for long texts. 

49 

50 For short text, calls predict_func directly. 

51 For long text, chunks it and merges predictions with deduplication. 

52 

53 :param text: Input text to process 

54 :param predict_func: Function that takes text and returns 

55 RecognizerResult objects 

56 :return: List of RecognizerResult with correct offsets 

57 """ 

58 chunks = self.chunk(text) 

59 if not chunks: 

60 return [] 

61 if len(chunks) == 1: 

62 return predict_func(text) 

63 

64 predictions = self._process_chunks(chunks, predict_func) 

65 return self.deduplicate_overlapping_entities(predictions) 

66 

67 def _process_chunks( 

68 self, 

69 chunks: List[TextChunk], 

70 process_func: Callable[[str], List["RecognizerResult"]], 

71 ) -> List["RecognizerResult"]: 

72 """Process text chunks and adjust entity offsets. 

73 

74 :param chunks: List of TextChunk objects with text and position information 

75 :param process_func: Function that takes chunk text and returns 

76 RecognizerResult objects 

77 :return: List of RecognizerResult with adjusted offsets 

78 """ 

79 from presidio_analyzer import RecognizerResult 

80 

81 all_predictions = [] 

82 

83 for chunk in chunks: 

84 chunk_predictions = process_func(chunk.text) 

85 

86 # Create new RecognizerResult objects with adjusted offsets 

87 # to avoid mutating the original predictions 

88 for pred in chunk_predictions: 

89 adjusted_pred = RecognizerResult( 

90 entity_type=pred.entity_type, 

91 start=pred.start + chunk.start, 

92 end=pred.end + chunk.start, 

93 score=pred.score, 

94 analysis_explanation=pred.analysis_explanation, 

95 recognition_metadata=pred.recognition_metadata, 

96 ) 

97 all_predictions.append(adjusted_pred) 

98 

99 return all_predictions 

100 

101 def deduplicate_overlapping_entities( 

102 self, 

103 predictions: List["RecognizerResult"], 

104 overlap_threshold: float = 0.5, 

105 ) -> List["RecognizerResult"]: 

106 """Remove duplicate entities from overlapping chunks. 

107 

108 :param predictions: List of RecognizerResult objects 

109 :param overlap_threshold: Overlap ratio threshold to consider duplicates 

110 (default: 0.5) 

111 :return: Deduplicated list of RecognizerResult sorted by position 

112 """ 

113 if not predictions: 

114 return predictions 

115 

116 # Sort by score descending to keep highest scoring entities 

117 sorted_preds = sorted(predictions, key=lambda p: p.score, reverse=True) 

118 unique = [] 

119 

120 for pred in sorted_preds: 

121 is_duplicate = False 

122 for kept in unique: 

123 # Check if same entity type and overlapping positions 

124 if pred.entity_type == kept.entity_type: 

125 overlap_start = max(pred.start, kept.start) 

126 overlap_end = min(pred.end, kept.end) 

127 

128 if overlap_start < overlap_end: 

129 # Calculate overlap ratio 

130 overlap_len = overlap_end - overlap_start 

131 pred_len = pred.end - pred.start 

132 kept_len = kept.end - kept.start 

133 

134 if pred_len <= 0 or kept_len <= 0: 

135 continue 

136 

137 # Check if overlap exceeds threshold 

138 if overlap_len / min(pred_len, kept_len) > overlap_threshold: 

139 is_duplicate = True 

140 break 

141 

142 if not is_duplicate: 

143 unique.append(pred) 

144 

145 # Sort by position for consistent output 

146 return sorted(unique, key=lambda p: p.start)