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 

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 @component.output_types(texts=list[str]) 

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

66 """ 

67 Cleans up the given list of strings. 

68 

69 :param texts: List of strings to clean. 

70 :returns: A dictionary with the following key: 

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

72 """ 

73 

74 if self._regex: 

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

76 

77 if self._convert_to_lowercase: 

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

79 

80 if self._translator: 

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

82 

83 return {"texts": texts}