Coverage for haystack/components/preprocessors/sentence_tokenizer.py: 95%
87 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-12 09:42 +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'"[^"]*"|\'[^\']*\'')
45# Closing brackets and quotes that can trail a sentence ending, as in `He said "Hi." Bye.` or `(Hi.) Bye.`
46SENTENCE_CLOSING_CHARS_RE = r"[\)\]}\"'’”»]*"
48if nltk_imports.is_successful():
50 def load_sentence_tokenizer(
51 language: Language, keep_white_spaces: bool = False
52 ) -> nltk.tokenize.punkt.PunktSentenceTokenizer:
53 """
54 Utility function to load the nltk sentence tokenizer.
56 :param language: The language for the tokenizer.
57 :param keep_white_spaces: If True, the tokenizer will keep white spaces between sentences.
58 :returns: nltk sentence tokenizer.
59 """
60 try:
61 nltk.data.find("tokenizers/punkt_tab")
62 except LookupError:
63 try:
64 nltk.download("punkt_tab")
65 except FileExistsError as error:
66 logger.debug("NLTK punkt tokenizer seems to be already downloaded. Error message: {error}", error=error)
68 language_name = ISO639_TO_NLTK.get(language)
70 if language_name is not None:
71 sentence_tokenizer = nltk.data.load(f"tokenizers/punkt_tab/{language_name}.pickle")
72 else:
73 logger.warning(
74 "PreProcessor couldn't find the default sentence tokenizer model for {language}. "
75 " Using English instead. You may train your own model and use the 'tokenizer_model_folder' parameter.",
76 language=language,
77 )
78 sentence_tokenizer = nltk.data.load("tokenizers/punkt_tab/english.pickle")
80 if keep_white_spaces:
81 sentence_tokenizer._lang_vars = CustomPunktLanguageVars()
83 return sentence_tokenizer
85 class CustomPunktLanguageVars(nltk.tokenize.punkt.PunktLanguageVars):
86 # The following adjustment of PunktSentenceTokenizer is inspired by:
87 # https://stackoverflow.com/questions/33139531/preserve-empty-lines-with-nltks-punkt-tokenizer
88 # It is needed for preserving whitespace while splitting text into sentences.
89 _period_context_fmt = r"""
90 %(SentEndChars)s # a potential sentence ending
91 \s* # match potential whitespace [ \t\n\x0B\f\r]
92 (?=(?P<after_tok>
93 %(NonWord)s # either other punctuation
94 |
95 (?P<next_tok>\S+) # or some other token - original version: \s+(?P<next_tok>\S+)
96 ))"""
98 def period_context_re(self) -> re.Pattern:
99 """
100 Compiles and returns a regular expression to find contexts including possible sentence boundaries.
102 :returns: A compiled regular expression pattern.
103 """
104 try:
105 return self._re_period_context # type: ignore
106 except AttributeError:
107 self._re_period_context = re.compile(
108 self._period_context_fmt
109 % {
110 "NonWord": self._re_non_word_chars,
111 # SentEndChars might be followed by closing brackets or quotes, so we match them here.
112 # If we don't, the whitespace after e.g. `."` is behind the closing quote, the pattern
113 # above can't reach it and it ends up in none of the sentences we return.
114 "SentEndChars": self._re_sent_end_chars + SENTENCE_CLOSING_CHARS_RE,
115 },
116 re.UNICODE | re.VERBOSE,
117 )
118 return self._re_period_context
121class SentenceSplitter:
122 """
123 SentenceSplitter splits a text into sentences using the nltk sentence tokenizer
124 """
126 def __init__(
127 self,
128 language: Language = "en",
129 use_split_rules: bool = True,
130 extend_abbreviations: bool = True,
131 keep_white_spaces: bool = False,
132 ) -> None:
133 """
134 Initializes the SentenceSplitter with the specified language, split rules, and abbreviation handling.
136 :param language: The language for the tokenizer. Default is "en".
137 :param use_split_rules: If True, the additional split rules are used. If False, the rules are not used.
138 :param extend_abbreviations: If True, the abbreviations used by NLTK's PunktTokenizer are extended by a list
139 of curated abbreviations if available. If False, the default abbreviations are used.
140 Currently supported languages are: en, de.
141 :param keep_white_spaces: If True, the tokenizer will keep white spaces between sentences.
142 """
143 nltk_imports.check()
144 self.language = language
145 # after checking nltk_imports, we are sure that load_sentence_tokenizer is defined
146 self.sentence_tokenizer = load_sentence_tokenizer(language, keep_white_spaces=keep_white_spaces)
147 self.use_split_rules = use_split_rules
148 if extend_abbreviations:
149 abbreviations = SentenceSplitter._read_abbreviations(language)
150 self.sentence_tokenizer._params.abbrev_types.update(abbreviations)
151 self.keep_white_spaces = keep_white_spaces
153 def split_sentences(self, text: str) -> list[dict[str, Any]]:
154 """
155 Splits a text into sentences including references to original char positions for each split.
157 :param text: The text to split.
158 :returns: list of sentences with positions.
159 """
160 sentence_spans = list(self.sentence_tokenizer.span_tokenize(text))
161 if self.use_split_rules:
162 sentence_spans = SentenceSplitter._apply_split_rules(text, sentence_spans)
164 return [{"sentence": text[start:end], "start": start, "end": end} for start, end in sentence_spans]
166 @staticmethod
167 def _apply_split_rules(text: str, sentence_spans: list[tuple[int, int]]) -> list[tuple[int, int]]:
168 """
169 Applies additional split rules to the sentence spans.
171 :param text: The text to split.
172 :param sentence_spans: The list of sentence spans to split.
173 :returns: The list of sentence spans after applying the split rules.
174 """
175 new_sentence_spans = []
176 quote_spans = [match.span() for match in QUOTE_SPANS_RE.finditer(text)]
177 while sentence_spans:
178 span = sentence_spans.pop(0)
179 next_span = sentence_spans[0] if len(sentence_spans) > 0 else None
180 while next_span and SentenceSplitter._needs_join(text, span, next_span, quote_spans):
181 sentence_spans.pop(0)
182 span = (span[0], next_span[1])
183 next_span = sentence_spans[0] if len(sentence_spans) > 0 else None
184 start, end = span
185 new_sentence_spans.append((start, end))
186 return new_sentence_spans
188 @staticmethod
189 def _needs_join(
190 text: str, span: tuple[int, int], next_span: tuple[int, int], quote_spans: list[tuple[int, int]]
191 ) -> bool:
192 """
193 Checks if the spans need to be joined as parts of one sentence.
195 This method determines whether two adjacent sentence spans should be joined back together as a single sentence.
196 It's used to prevent incorrect sentence splitting in specific cases like quotations, numbered lists,
197 and parenthetical expressions.
199 :param text: The text containing the spans.
200 :param span: Tuple of (start, end) positions for the current sentence span.
201 :param next_span: Tuple of (start, end) positions for the next sentence span.
202 :param quote_spans: All quoted spans within text.
203 :returns:
204 True if the spans needs to be joined.
205 """
206 start, end = span
207 next_start, next_end = next_span
209 # with keep_white_spaces=True a span also covers the whitespace up to the next sentence, the rules below
210 # look at where the sentence itself ends
211 end = start + len(text[start:end].rstrip())
213 # sentence. sentence"\nsentence -> no split (end << quote_end)
214 # sentence.", sentence -> no split (end < quote_end)
215 # sentence?", sentence -> no split (end < quote_end)
216 if any(quote_start < end < quote_end for quote_start, quote_end in quote_spans):
217 # sentence boundary is inside a quote
218 return True
220 # sentence." sentence -> split (end == quote_end)
221 # sentence?" sentence -> no split (end == quote_end)
222 if any(quote_start < end == quote_end and text[quote_end - 2] == "?" for quote_start, quote_end in quote_spans):
223 # question is cited
224 return True
226 # sentence.", sentence -> no split (end == quote_end): widening the closing-char class (see
227 # period_context_re) moves the boundary just past the closing quote instead of leaving it inside,
228 # so anything other than whitespace directly after the quote continues the sentence.
229 if any(
230 quote_start < end == quote_end and quote_end < len(text) and not text[quote_end].isspace()
231 for quote_start, quote_end in quote_spans
232 ):
233 # e.g. `He said "Hi.", then left.` -> the punctuation continues the sentence
234 return True
236 if re.search(r"(^|\n)\s*\d{1,2}\.$", text[start:end]) is not None:
237 # sentence ends with a numeration
238 return True
240 # next sentence starts with a bracket or we return False
241 return re.search(r"^\s*[\(\[]", text[next_start:next_end]) is not None
243 @staticmethod
244 def _read_abbreviations(lang: Language) -> list[str]:
245 """
246 Reads the abbreviations for a given language from the abbreviations file.
248 :param lang: The language to read the abbreviations for.
249 :returns: List of abbreviations.
250 """
251 abbreviations_file = Path(__file__).parent.parent.parent / f"data/abbreviations/{lang}.txt"
252 if not abbreviations_file.exists():
253 logger.warning("No abbreviations file found for {language}. Using default abbreviations.", language=lang)
254 return []
256 return abbreviations_file.read_text(encoding="utf-8").split("\n")