Coverage for haystack/components/preprocessors/text_cleaner.py: 100%

1# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai> 

2# 

3# SPDX-License-Identifier: Apache-2.0 

4 

5import re 

6import string 

7from typing import Any 

8 

9from haystack import component, default_to_dict 

10 

11 

12@component 

13class TextCleaner: 

14 """ 

15 Cleans text strings. 

16 

17 It can remove substrings matching a list of regular expressions, convert text to lowercase, 

18 remove punctuation, and remove numbers. 

19 Use it to clean up text data before evaluation. 

20 

21 ### Usage example 

22 

23 ```python 

24 from haystack.components.preprocessors import TextCleaner 

25 

26 text_to_clean = "1Moonlight shimmered softly, 300 Wolves howled nearby, Night enveloped everything." 

27 

28 cleaner = TextCleaner(convert_to_lowercase=True, remove_punctuation=False, remove_numbers=True) 

29 result = cleaner.run(texts=[text_to_clean]) 

30 ``` 

31 """ 

32 

33 def __init__( 

34 self, 

35 remove_regexps: list[str] | None = None, 

36 convert_to_lowercase: bool = False, 

37 remove_punctuation: bool = False, 

38 remove_numbers: bool = False, 

39 ) -> None: 

40 """ 

41 Initializes the TextCleaner component. 

42 

43 :param remove_regexps: A list of regex patterns to remove matching substrings from the text. 

44 :param convert_to_lowercase: If `True`, converts all characters to lowercase. 

45 :param remove_punctuation: If `True`, removes punctuation from the text. 

46 :param remove_numbers: If `True`, removes numerical digits from the text. 

47 """ 

48 self._remove_regexps = remove_regexps 

49 self._convert_to_lowercase = convert_to_lowercase 

50 self._remove_punctuation = remove_punctuation 

51 self._remove_numbers = remove_numbers 

52 

53 self._regex = None 

54 if remove_regexps: 

55 self._regex = re.compile("|".join(remove_regexps), flags=re.IGNORECASE) 

56 to_remove = "" 

57 if remove_punctuation: 

58 to_remove = string.punctuation 

59 if remove_numbers: 

60 to_remove += string.digits 

61 

62 self._translator = str.maketrans("", "", to_remove) if to_remove else None 

63 

64 def to_dict(self) -> dict[str, Any]: 

65 """ 

66 Serializes the component to a dictionary. 

67 

68 :returns: 

69 Dictionary with serialized data. 

70 """ 

71 return default_to_dict( 

72 self, 

73 remove_regexps=self._remove_regexps, 

74 convert_to_lowercase=self._convert_to_lowercase, 

75 remove_punctuation=self._remove_punctuation, 

76 remove_numbers=self._remove_numbers, 

77 ) 

78 

79 @component.output_types(texts=list[str]) 

80 def run(self, texts: list[str]) -> dict[str, Any]: 

81 """ 

82 Cleans up the given list of strings. 

83 

84 :param texts: List of strings to clean. 

85 :returns: A dictionary with the following key: 

86 - `texts`: the cleaned list of strings. 

87 :raises TypeError: 

88 If `texts` is not a list or any element is not a string. 

89 """ 

90 if not isinstance(texts, list) or not all(isinstance(text, str) for text in texts): 

91 raise TypeError("`texts` must be a list of strings.") 

92 

93 if self._regex: 

94 texts = [self._regex.sub("", text) for text in texts] 

95 

96 if self._convert_to_lowercase: 

97 texts = [text.lower() for text in texts] 

98 

99 if self._translator: 

100 texts = [text.translate(self._translator) for text in texts] 

101 

102 return {"texts": texts}