Coverage for haystack/components/preprocessors/sentence_tokenizer.py: 94%
83 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-21 13:53 +0000
1# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
2#
3# SPDX-License-Identifier: Apache-2.0
5import re
6from pathlib import Path
7from typing import Any, Literal
9from haystack import logging
10from haystack.lazy_imports import LazyImport
12with LazyImport("Run 'pip install nltk>=3.9.1'") as nltk_imports:
13 import nltk
15logger = logging.getLogger(__name__)
17Language = Literal[
18 "ru", "sl", "es", "sv", "tr", "cs", "da", "nl", "en", "et", "fi", "fr", "de", "el", "it", "no", "pl", "pt", "ml"
19]
21ISO639_TO_NLTK = {
22 "ru": "russian",
23 "sl": "slovene",
24 "es": "spanish",
25 "sv": "swedish",
26 "tr": "turkish",
27 "cs": "czech",
28 "da": "danish",
29 "nl": "dutch",
30 "en": "english",
31 "et": "estonian",
32 "fi": "finnish",
33 "fr": "french",
34 "de": "german",
35 "el": "greek",
36 "it": "italian",
37 "no": "norwegian",
38 "pl": "polish",
39 "pt": "portuguese",
40 "ml": "malayalam",
41}
43QUOTE_SPANS_RE = re.compile(r'"[^"]*"|\'[^\']*\'')
45if nltk_imports.is_successful():
47 def load_sentence_tokenizer(
48 language: Language, keep_white_spaces: bool = False
49 ) -> nltk.tokenize.punkt.PunktSentenceTokenizer:
50 """
51 Utility function to load the nltk sentence tokenizer.
53 :param language: The language for the tokenizer.
54 :param keep_white_spaces: If True, the tokenizer will keep white spaces between sentences.
55 :returns: nltk sentence tokenizer.
56 """
57 try:
58 nltk.data.find("tokenizers/punkt_tab")
59 except LookupError:
60 try:
61 nltk.download("punkt_tab")
62 except FileExistsError as error:
63 logger.debug("NLTK punkt tokenizer seems to be already downloaded. Error message: {error}", error=error)
65 language_name = ISO639_TO_NLTK.get(language)
67 if language_name is not None:
68 sentence_tokenizer = nltk.data.load(f"tokenizers/punkt_tab/{language_name}.pickle")
69 else:
70 logger.warning(
71 "PreProcessor couldn't find the default sentence tokenizer model for {language}. "
72 " Using English instead. You may train your own model and use the 'tokenizer_model_folder' parameter.",
73 language=language,
74 )
75 sentence_tokenizer = nltk.data.load("tokenizers/punkt_tab/english.pickle")
77 if keep_white_spaces:
78 sentence_tokenizer._lang_vars = CustomPunktLanguageVars()
80 return sentence_tokenizer
82 class CustomPunktLanguageVars(nltk.tokenize.punkt.PunktLanguageVars):
83 # The following adjustment of PunktSentenceTokenizer is inspired by:
84 # https://stackoverflow.com/questions/33139531/preserve-empty-lines-with-nltks-punkt-tokenizer
85 # It is needed for preserving whitespace while splitting text into sentences.
86 _period_context_fmt = r"""
87 %(SentEndChars)s # a potential sentence ending
88 \s* # match potential whitespace [ \t\n\x0B\f\r]
89 (?=(?P<after_tok>
90 %(NonWord)s # either other punctuation
91 |
92 (?P<next_tok>\S+) # or some other token - original version: \s+(?P<next_tok>\S+)
93 ))"""
95 def period_context_re(self) -> re.Pattern:
96 """
97 Compiles and returns a regular expression to find contexts including possible sentence boundaries.
99 :returns: A compiled regular expression pattern.
100 """
101 try:
102 return self._re_period_context # type: ignore
103 except: # noqa: E722
104 self._re_period_context = re.compile(
105 self._period_context_fmt
106 % {
107 "NonWord": self._re_non_word_chars,
108 # SentEndChars might be followed by closing brackets, so we match them here.
109 "SentEndChars": self._re_sent_end_chars + r"[\)\]}]*",
110 },
111 re.UNICODE | re.VERBOSE,
112 )
113 return self._re_period_context
116class SentenceSplitter:
117 """
118 SentenceSplitter splits a text into sentences using the nltk sentence tokenizer
119 """
121 def __init__(
122 self,
123 language: Language = "en",
124 use_split_rules: bool = True,
125 extend_abbreviations: bool = True,
126 keep_white_spaces: bool = False,
127 ) -> None:
128 """
129 Initializes the SentenceSplitter with the specified language, split rules, and abbreviation handling.
131 :param language: The language for the tokenizer. Default is "en".
132 :param use_split_rules: If True, the additional split rules are used. If False, the rules are not used.
133 :param extend_abbreviations: If True, the abbreviations used by NLTK's PunktTokenizer are extended by a list
134 of curated abbreviations if available. If False, the default abbreviations are used.
135 Currently supported languages are: en, de.
136 :param keep_white_spaces: If True, the tokenizer will keep white spaces between sentences.
137 """
138 nltk_imports.check()
139 self.language = language
140 # after checking nltk_imports, we are sure that load_sentence_tokenizer is defined
141 self.sentence_tokenizer = load_sentence_tokenizer(language, keep_white_spaces=keep_white_spaces)
142 self.use_split_rules = use_split_rules
143 if extend_abbreviations:
144 abbreviations = SentenceSplitter._read_abbreviations(language)
145 self.sentence_tokenizer._params.abbrev_types.update(abbreviations)
146 self.keep_white_spaces = keep_white_spaces
148 def split_sentences(self, text: str) -> list[dict[str, Any]]:
149 """
150 Splits a text into sentences including references to original char positions for each split.
152 :param text: The text to split.
153 :returns: list of sentences with positions.
154 """
155 sentence_spans = list(self.sentence_tokenizer.span_tokenize(text))
156 if self.use_split_rules:
157 sentence_spans = SentenceSplitter._apply_split_rules(text, sentence_spans)
159 return [{"sentence": text[start:end], "start": start, "end": end} for start, end in sentence_spans]
161 @staticmethod
162 def _apply_split_rules(text: str, sentence_spans: list[tuple[int, int]]) -> list[tuple[int, int]]:
163 """
164 Applies additional split rules to the sentence spans.
166 :param text: The text to split.
167 :param sentence_spans: The list of sentence spans to split.
168 :returns: The list of sentence spans after applying the split rules.
169 """
170 new_sentence_spans = []
171 quote_spans = [match.span() for match in QUOTE_SPANS_RE.finditer(text)]
172 while sentence_spans:
173 span = sentence_spans.pop(0)
174 next_span = sentence_spans[0] if len(sentence_spans) > 0 else None
175 while next_span and SentenceSplitter._needs_join(text, span, next_span, quote_spans):
176 sentence_spans.pop(0)
177 span = (span[0], next_span[1])
178 next_span = sentence_spans[0] if len(sentence_spans) > 0 else None
179 start, end = span
180 new_sentence_spans.append((start, end))
181 return new_sentence_spans
183 @staticmethod
184 def _needs_join(
185 text: str, span: tuple[int, int], next_span: tuple[int, int], quote_spans: list[tuple[int, int]]
186 ) -> bool:
187 """
188 Checks if the spans need to be joined as parts of one sentence.
190 This method determines whether two adjacent sentence spans should be joined back together as a single sentence.
191 It's used to prevent incorrect sentence splitting in specific cases like quotations, numbered lists,
192 and parenthetical expressions.
194 :param text: The text containing the spans.
195 :param span: Tuple of (start, end) positions for the current sentence span.
196 :param next_span: Tuple of (start, end) positions for the next sentence span.
197 :param quote_spans: All quoted spans within text.
198 :returns:
199 True if the spans needs to be joined.
200 """
201 start, end = span
202 next_start, next_end = next_span
204 # sentence. sentence"\nsentence -> no split (end << quote_end)
205 # sentence.", sentence -> no split (end < quote_end)
206 # sentence?", sentence -> no split (end < quote_end)
207 if any(quote_start < end < quote_end for quote_start, quote_end in quote_spans):
208 # sentence boundary is inside a quote
209 return True
211 # sentence." sentence -> split (end == quote_end)
212 # sentence?" sentence -> no split (end == quote_end)
213 if any(quote_start < end == quote_end and text[quote_end - 2] == "?" for quote_start, quote_end in quote_spans):
214 # question is cited
215 return True
217 if re.search(r"(^|\n)\s*\d{1,2}\.$", text[start:end]) is not None:
218 # sentence ends with a numeration
219 return True
221 # next sentence starts with a bracket or we return False
222 return re.search(r"^\s*[\(\[]", text[next_start:next_end]) is not None
224 @staticmethod
225 def _read_abbreviations(lang: Language) -> list[str]:
226 """
227 Reads the abbreviations for a given language from the abbreviations file.
229 :param lang: The language to read the abbreviations for.
230 :returns: List of abbreviations.
231 """
232 abbreviations_file = Path(__file__).parent.parent.parent / f"data/abbreviations/{lang}.txt"
233 if not abbreviations_file.exists():
234 logger.warning("No abbreviations file found for {language}. Using default abbreviations.", language=lang)
235 return []
237 return abbreviations_file.read_text().split("\n")