Coverage for haystack/components/preprocessors/document_splitter.py: 99%

280 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 

4 

5from collections.abc import Callable 

6from copy import deepcopy 

7from typing import Any, Literal 

8 

9from more_itertools import windowed 

10 

11from haystack import Document, component, logging 

12from haystack.components.preprocessors._page_numbers import _leading_page_breaks 

13from haystack.components.preprocessors.sentence_tokenizer import Language, SentenceSplitter, nltk_imports 

14from haystack.core.serialization import default_from_dict, default_to_dict 

15from haystack.lazy_imports import LazyImport 

16from haystack.utils import deserialize_callable, serialize_callable 

17 

18with LazyImport("Run 'pip install tiktoken'") as tiktoken_imports: 

19 import tiktoken 

20 

21logger = logging.getLogger(__name__) 

22 

23# mapping of split by character, 'function' and 'sentence' don't split by character 

24_CHARACTER_SPLIT_BY_MAPPING = {"page": "\f", "passage": "\n\n", "period": ".", "word": " ", "line": "\n"} 

25 

26 

27@component 

28class DocumentSplitter: 

29 """ 

30 Splits long documents into smaller chunks. 

31 

32 This is a common preprocessing step during indexing. It helps Embedders create meaningful semantic representations 

33 and prevents exceeding language model context limits. 

34 

35 The DocumentSplitter is compatible with the following DocumentStores: 

36 - [Astra](https://docs.haystack.deepset.ai/docs/astradocumentstore) 

37 - [Chroma](https://docs.haystack.deepset.ai/docs/chromadocumentstore) limited support, overlapping information is 

38 not stored 

39 - [Elasticsearch](https://docs.haystack.deepset.ai/docs/elasticsearch-document-store) 

40 - [OpenSearch](https://docs.haystack.deepset.ai/docs/opensearch-document-store) 

41 - [Pgvector](https://docs.haystack.deepset.ai/docs/pgvectordocumentstore) 

42 - [Pinecone](https://docs.haystack.deepset.ai/docs/pinecone-document-store) limited support, overlapping 

43 information is not stored 

44 - [Qdrant](https://docs.haystack.deepset.ai/docs/qdrant-document-store) 

45 - [Weaviate](https://docs.haystack.deepset.ai/docs/weaviatedocumentstore) 

46 

47 ### Usage example 

48 

49 ```python 

50 from haystack import Document 

51 from haystack.components.preprocessors import DocumentSplitter 

52 

53 doc = Document(content="Moonlight shimmered softly, wolves howled nearby, night enveloped everything.") 

54 

55 splitter = DocumentSplitter(split_by="word", split_length=3, split_overlap=0) 

56 result = splitter.run(documents=[doc]) 

57 ``` 

58 """ 

59 

60 def __init__( 

61 self, 

62 split_by: Literal["function", "page", "passage", "period", "word", "line", "sentence", "token"] = "word", 

63 split_length: int = 200, 

64 split_overlap: int = 0, 

65 split_threshold: int = 0, 

66 splitting_function: Callable[[str], list[str]] | None = None, 

67 respect_sentence_boundary: bool = False, 

68 language: Language = "en", 

69 use_split_rules: bool = True, 

70 extend_abbreviations: bool = True, 

71 *, 

72 skip_empty_documents: bool = True, 

73 tokenizer_encoding: str = "o200k_base", 

74 ) -> None: 

75 """ 

76 Initialize DocumentSplitter. 

77 

78 :param split_by: The unit for splitting your documents. Choose from: 

79 - `word` for splitting by spaces (" ") 

80 - `period` for splitting by periods (".") 

81 - `page` for splitting by form feed ("\\f") 

82 - `passage` for splitting by double line breaks ("\\n\\n") 

83 - `line` for splitting each line ("\\n") 

84 - `sentence` for splitting by NLTK sentence tokenizer 

85 - `token` for splitting by token count using tiktoken (requires `pip install tiktoken`) 

86 

87 :param split_length: The maximum number of units in each split. 

88 :param split_overlap: The number of overlapping units for each split. 

89 :param split_threshold: The minimum number of units per split. If a split has fewer units 

90 than the threshold, it's attached to the previous split. 

91 :param splitting_function: Necessary when `split_by` is set to "function". 

92 This is a function which must accept a single `str` as input and return a `list` of `str` as output, 

93 representing the chunks after splitting. 

94 :param respect_sentence_boundary: Choose whether to respect sentence boundaries when splitting by "word". 

95 If True, uses NLTK to detect sentence boundaries, ensuring splits occur only between sentences. 

96 :param language: Choose the language for the NLTK tokenizer. The default is English ("en"). 

97 :param use_split_rules: Choose whether to use additional split rules when splitting by `sentence`. 

98 :param extend_abbreviations: Choose whether to extend NLTK's PunktTokenizer abbreviations with a list 

99 of curated abbreviations, if available. This is currently supported for English ("en") and German ("de"). 

100 :param skip_empty_documents: Choose whether to skip documents with empty content. Default is True. 

101 Set to False when downstream components in the Pipeline (like LLMDocumentContentExtractor) can extract text 

102 from non-textual documents. 

103 :param tokenizer_encoding: The tiktoken encoding to use when `split_by="token"`. Defaults to 

104 `"o200k_base"` (current OpenAI models). Only used when `split_by="token"`. 

105 Special-token strings in document content are encoded as ordinary text. 

106 """ 

107 

108 self.split_by = split_by 

109 self.split_length = split_length 

110 self.split_overlap = split_overlap 

111 self.split_threshold = split_threshold 

112 self.splitting_function = splitting_function 

113 self.respect_sentence_boundary = respect_sentence_boundary 

114 self.language = language 

115 self.use_split_rules = use_split_rules 

116 self.extend_abbreviations = extend_abbreviations 

117 self.skip_empty_documents = skip_empty_documents 

118 self.tokenizer_encoding = tokenizer_encoding 

119 

120 self._init_checks( 

121 split_by=split_by, 

122 split_length=split_length, 

123 split_overlap=split_overlap, 

124 splitting_function=splitting_function, 

125 respect_sentence_boundary=respect_sentence_boundary, 

126 ) 

127 self._use_sentence_splitter = split_by == "sentence" or (respect_sentence_boundary and split_by == "word") 

128 if self._use_sentence_splitter: 

129 nltk_imports.check() 

130 self.sentence_splitter: SentenceSplitter | None = None 

131 if split_by == "token": 

132 tiktoken_imports.check() 

133 self._tiktoken_tokenizer: "tiktoken.Encoding | None" = None 

134 

135 def _init_checks( 

136 self, 

137 *, 

138 split_by: str, 

139 split_length: int, 

140 split_overlap: int, 

141 splitting_function: Callable | None, 

142 respect_sentence_boundary: bool, 

143 ) -> None: 

144 """ 

145 Validates initialization parameters for DocumentSplitter. 

146 

147 :param split_by: The unit for splitting documents 

148 :param split_length: The maximum number of units in each split 

149 :param split_overlap: The number of overlapping units for each split 

150 :param splitting_function: Custom function for splitting when split_by="function" 

151 :param respect_sentence_boundary: Whether to respect sentence boundaries when splitting 

152 :raises ValueError: If any parameter is invalid 

153 """ 

154 valid_split_by = ["function", "page", "passage", "period", "word", "line", "sentence", "token"] 

155 if split_by not in valid_split_by: 

156 raise ValueError(f"split_by must be one of {', '.join(valid_split_by)}.") 

157 

158 if split_by == "function" and splitting_function is None: 

159 raise ValueError("When 'split_by' is set to 'function', a valid 'splitting_function' must be provided.") 

160 

161 if split_length <= 0: 

162 raise ValueError("split_length must be greater than 0.") 

163 

164 if split_overlap < 0: 

165 raise ValueError("split_overlap must be greater than or equal to 0.") 

166 

167 if split_overlap >= split_length: 

168 raise ValueError("split_overlap must be less than split_length.") 

169 

170 if respect_sentence_boundary and split_by != "word": 

171 logger.warning( 

172 "The 'respect_sentence_boundary' option is only supported for `split_by='word'`. " 

173 "The option `respect_sentence_boundary` will be set to `False`." 

174 ) 

175 self.respect_sentence_boundary = False 

176 

177 def warm_up(self) -> None: 

178 """ 

179 Warm up the DocumentSplitter by loading the sentence tokenizer or tiktoken encoder. 

180 """ 

181 if self._use_sentence_splitter and self.sentence_splitter is None: 

182 self.sentence_splitter = SentenceSplitter( 

183 language=self.language, 

184 use_split_rules=self.use_split_rules, 

185 extend_abbreviations=self.extend_abbreviations, 

186 keep_white_spaces=True, 

187 ) 

188 if self.split_by == "token" and self._tiktoken_tokenizer is None: 

189 self._tiktoken_tokenizer = tiktoken.get_encoding(self.tokenizer_encoding) 

190 

191 @component.output_types(documents=list[Document]) 

192 def run(self, documents: list[Document]) -> dict[str, list[Document]]: 

193 """ 

194 Split documents into smaller parts. 

195 

196 Splits documents by the unit expressed in `split_by`, with a length of `split_length` 

197 and an overlap of `split_overlap`. 

198 

199 :param documents: The documents to split. 

200 :returns: A dictionary with the following key: 

201 - `documents`: List of documents with the split texts. Each document includes: 

202 - A metadata field `source_id` to track the original document. 

203 - A metadata field `page_number` with the page the chunk starts on, counting form feed 

204 ("\f") characters in the original document. 

205 - All other metadata copied from the original document. 

206 

207 :raises TypeError: if the input is not a list of Documents. 

208 :raises ValueError: if the content of a document is None. 

209 """ 

210 if self._use_sentence_splitter and self.sentence_splitter is None: 

211 self.warm_up() 

212 if self.split_by == "token" and self._tiktoken_tokenizer is None: 

213 self.warm_up() 

214 

215 if not isinstance(documents, list) or (documents and not isinstance(documents[0], Document)): 

216 raise TypeError("DocumentSplitter expects a List of Documents as input.") 

217 

218 split_docs: list[Document] = [] 

219 for doc in documents: 

220 if doc.content is None: 

221 raise ValueError( 

222 f"DocumentSplitter only works with text documents but content for document ID {doc.id} is None." 

223 ) 

224 if doc.content == "" and self.skip_empty_documents: 

225 logger.warning("Document ID {doc_id} has an empty content. Skipping this document.", doc_id=doc.id) 

226 continue 

227 

228 split_docs += self._split_document(doc) 

229 return {"documents": split_docs} 

230 

231 def _split_document(self, doc: Document) -> list[Document]: 

232 if self.split_by == "sentence" or self.respect_sentence_boundary: 

233 return self._split_by_nltk_sentence(doc) 

234 

235 if self.split_by == "function" and self.splitting_function is not None: 

236 return self._split_by_function(doc) 

237 

238 if self.split_by == "token": 

239 return self._split_by_token(doc) 

240 

241 return self._split_by_character(doc) 

242 

243 def _split_by_nltk_sentence(self, doc: Document) -> list[Document]: 

244 split_docs = [] 

245 

246 result = self.sentence_splitter.split_sentences(doc.content) # type: ignore # None check is done in run() 

247 units = [sentence["sentence"] for sentence in result] 

248 

249 if self.respect_sentence_boundary: 

250 text_splits, splits_pages, splits_start_idxs = self._concatenate_sentences_based_on_word_amount( 

251 sentences=units, split_length=self.split_length, split_overlap=self.split_overlap 

252 ) 

253 else: 

254 text_splits, splits_pages, splits_start_idxs = self._concatenate_units( 

255 elements=units, 

256 split_length=self.split_length, 

257 split_overlap=self.split_overlap, 

258 split_threshold=self.split_threshold, 

259 ) 

260 metadata = deepcopy(doc.meta) 

261 metadata["source_id"] = doc.id 

262 split_docs += self._create_docs_from_splits( 

263 text_splits=text_splits, splits_pages=splits_pages, splits_start_idxs=splits_start_idxs, meta=metadata 

264 ) 

265 

266 return split_docs 

267 

268 def _split_by_token(self, doc: Document) -> list[Document]: 

269 """ 

270 Split a document by token count using tiktoken. 

271 

272 Encodes the full document text to tokens, slices into chunks of `split_length` tokens 

273 with `split_overlap` overlap, then decodes each chunk back to a string. 

274 Stops once a chunk reaches the end of the document, avoiding overlap-only trailing chunks. 

275 """ 

276 tokens = self._tiktoken_tokenizer.encode_ordinary(doc.content) # type: ignore[union-attr, arg-type] 

277 if not tokens: 

278 if self.skip_empty_documents: 

279 return [] 

280 metadata = deepcopy(doc.meta) 

281 metadata["source_id"] = doc.id 

282 return self._create_docs_from_splits( 

283 text_splits=[""], splits_pages=[1], splits_start_idxs=[0], meta=metadata 

284 ) 

285 

286 step = self.split_length - self.split_overlap 

287 full_text, offsets = self._tiktoken_tokenizer.decode_with_offsets(tokens) # type: ignore[union-attr] 

288 offsets.append(len(full_text)) 

289 

290 text_splits: list[str] = [] 

291 splits_pages: list[int] = [] 

292 splits_start_idxs: list[int] = [] 

293 cur_page = 1 

294 prev_end = 0 

295 

296 for i in range(0, len(tokens), step): 

297 chunk_token_count = min(self.split_length, len(tokens) - i) 

298 chunk_start = offsets[i] 

299 chunk_end = offsets[i + chunk_token_count] 

300 

301 if chunk_token_count < self.split_threshold and len(text_splits) > 0: 

302 overlap_start = offsets[min(i + self.split_overlap, len(tokens))] 

303 text_splits[-1] += full_text[overlap_start:chunk_end] 

304 else: 

305 text_splits.append(full_text[chunk_start:chunk_end]) 

306 splits_pages.append(cur_page) 

307 splits_start_idxs.append(chunk_start) 

308 

309 if i + chunk_token_count == len(tokens): 

310 break 

311 

312 non_overlap_end = offsets[min(i + step, len(tokens))] 

313 cur_page += full_text[prev_end:non_overlap_end].count("\f") 

314 prev_end = non_overlap_end 

315 

316 metadata = deepcopy(doc.meta) 

317 metadata["source_id"] = doc.id 

318 return self._create_docs_from_splits( 

319 text_splits=text_splits, splits_pages=splits_pages, splits_start_idxs=splits_start_idxs, meta=metadata 

320 ) 

321 

322 def _split_by_character(self, doc: Document) -> list[Document]: 

323 split_at = _CHARACTER_SPLIT_BY_MAPPING[self.split_by] 

324 units = doc.content.split(split_at) # type: ignore[union-attr] 

325 # Add the delimiter back to all units except the last one 

326 for i in range(len(units) - 1): 

327 units[i] += split_at 

328 text_splits, splits_pages, splits_start_idxs = self._concatenate_units( 

329 units, self.split_length, self.split_overlap, self.split_threshold 

330 ) 

331 metadata = deepcopy(doc.meta) 

332 metadata["source_id"] = doc.id 

333 return self._create_docs_from_splits( 

334 text_splits=text_splits, splits_pages=splits_pages, splits_start_idxs=splits_start_idxs, meta=metadata 

335 ) 

336 

337 def _split_by_function(self, doc: Document) -> list[Document]: 

338 # the check for None is done already in the run method 

339 splits = self.splitting_function(doc.content) # type: ignore 

340 content: str = doc.content # type: ignore[assignment] 

341 

342 text_splits: list[str] = [] 

343 splits_pages: list[int] = [] 

344 splits_start_idxs: list[int] = [] 

345 cur_start_idx = 0 

346 cur_page = 1 

347 prev_start_idx = 0 

348 

349 for split in splits: 

350 if self.skip_empty_documents and len(split) == 0: 

351 continue 

352 

353 # A splitting function may transform the text, so a split is not guaranteed to appear 

354 # verbatim in the source. Locate it when possible and fall back to the running offset. 

355 found_idx = content.find(split, cur_start_idx) 

356 if found_idx == -1 and text_splits: 

357 # The split may legitimately start before the end of the previous split, e.g. an 

358 # overlapping/sliding-window splitting function. Retry allowing it to be found 

359 # anywhere after the start of the previous split, so it isn't missed just because 

360 # the forward-only search above already moved past it. 

361 found_idx = content.find(split, prev_start_idx + 1) 

362 

363 if found_idx != -1: 

364 cur_start_idx = found_idx 

365 cur_page = 1 + content.count("\f", 0, found_idx) 

366 

367 text_splits.append(split) 

368 splits_start_idxs.append(cur_start_idx) 

369 splits_pages.append(cur_page) 

370 

371 prev_start_idx = cur_start_idx 

372 cur_start_idx += len(split) 

373 if found_idx == -1: 

374 cur_page += split.count("\f") 

375 

376 metadata = deepcopy(doc.meta) 

377 metadata["source_id"] = doc.id 

378 return self._create_docs_from_splits( 

379 text_splits=text_splits, splits_pages=splits_pages, splits_start_idxs=splits_start_idxs, meta=metadata 

380 ) 

381 

382 def _concatenate_units( 

383 self, elements: list[str], split_length: int, split_overlap: int, split_threshold: int 

384 ) -> tuple[list[str], list[int], list[int]]: 

385 """ 

386 Concatenates the elements into parts of split_length units. 

387 

388 Keeps track of the original page number that each element belongs. If the length of the current units is less 

389 than the pre-defined `split_threshold`, it does not create a new split. Instead, it concatenates the current 

390 units with the last split, preventing the creation of excessively small splits. 

391 A trailing segment that adds no new text beyond the already covered units (an overlap-only window, 

392 or a window holding only the empty artifact of a trailing delimiter) is skipped instead of 

393 creating a redundant chunk. 

394 """ 

395 step = split_length - split_overlap 

396 

397 text_splits: list[str] = [] 

398 splits_pages: list[int] = [] 

399 splits_start_idxs: list[int] = [] 

400 cur_start_idx = 0 

401 cur_page = 1 

402 segments = windowed(elements, n=split_length, step=step) 

403 # Number of leading units already covered by previous segments. The yielded segments advance by 

404 # `step`, so the segment at index i starts at unit i * step. 

405 covered_unit_count = 0 

406 

407 for seg_index, seg in enumerate(segments): 

408 current_units = [unit for unit in seg if unit is not None] 

409 txt = "".join(current_units) 

410 

411 # Skip a segment that adds no new text: every unit past the already covered ones is empty. 

412 # Emitting it would create a chunk fully contained in a previous chunk. This mirrors the 

413 # overlap-only trailing chunk fix of `_split_by_token` for the character-based split modes. 

414 new_units = current_units[max(0, covered_unit_count - seg_index * step) :] 

415 if split_overlap > 0 and text_splits and not "".join(new_units): 

416 continue 

417 covered_unit_count = max(covered_unit_count, seg_index * step + len(current_units)) 

418 

419 # check if length of current units is below split_threshold 

420 if len(current_units) < split_threshold and len(text_splits) > 0: 

421 # concatenate the last split with the current one. The first `split_overlap` units 

422 # of this segment are already part of the previous split (that is what the overlap 

423 # is), so only append the units beyond the overlap; otherwise the overlapping text 

424 # would be duplicated and the merged chunk would not be present in the source. 

425 text_splits[-1] += "".join(current_units[split_overlap:]) 

426 

427 # NOTE: If skip_empty_documents is True, this line skips documents that have content="" 

428 elif not self.skip_empty_documents or len(txt) > 0: 

429 text_splits.append(txt) 

430 splits_pages.append(cur_page) 

431 splits_start_idxs.append(cur_start_idx) 

432 

433 processed_units = current_units[: split_length - split_overlap] 

434 cur_start_idx += len("".join(processed_units)) 

435 

436 if self.split_by == "page": 

437 num_page_breaks = len(processed_units) 

438 else: 

439 num_page_breaks = sum(processed_unit.count("\f") for processed_unit in processed_units) 

440 

441 cur_page += num_page_breaks 

442 

443 return text_splits, splits_pages, splits_start_idxs 

444 

445 def _create_docs_from_splits( 

446 self, text_splits: list[str], splits_pages: list[int], splits_start_idxs: list[int], meta: dict[str, Any] 

447 ) -> list[Document]: 

448 """ 

449 Creates Document objects from splits enriching them with page number and the metadata of the original document. 

450 """ 

451 documents: list[Document] = [] 

452 

453 for i, (txt, split_idx) in enumerate(zip(text_splits, splits_start_idxs, strict=True)): 

454 copied_meta = deepcopy(meta) 

455 copied_meta["page_number"] = splits_pages[i] + _leading_page_breaks(txt) 

456 copied_meta["split_id"] = i 

457 copied_meta["split_idx_start"] = split_idx 

458 doc = Document(content=txt, meta=copied_meta) 

459 documents.append(doc) 

460 

461 if self.split_overlap <= 0: 

462 continue 

463 

464 doc.meta["_split_overlap"] = [] 

465 

466 if i == 0: 

467 continue 

468 

469 doc_start_idx = splits_start_idxs[i] 

470 previous_doc = documents[i - 1] 

471 previous_doc_start_idx = splits_start_idxs[i - 1] 

472 self._add_split_overlap_information(doc, doc_start_idx, previous_doc, previous_doc_start_idx) 

473 

474 return documents 

475 

476 @staticmethod 

477 def _add_split_overlap_information( 

478 current_doc: Document, current_doc_start_idx: int, previous_doc: Document, previous_doc_start_idx: int 

479 ) -> None: 

480 """ 

481 Adds split overlap information to the current and previous Document's meta. 

482 

483 :param current_doc: The Document that is being split. 

484 :param current_doc_start_idx: The starting index of the current Document. 

485 :param previous_doc: The Document that was split before the current Document. 

486 :param previous_doc_start_idx: The starting index of the previous Document. 

487 """ 

488 overlapping_range = (current_doc_start_idx - previous_doc_start_idx, len(previous_doc.content)) # type: ignore 

489 

490 if overlapping_range[0] < overlapping_range[1]: 

491 overlapping_str = previous_doc.content[overlapping_range[0] : overlapping_range[1]] # type: ignore 

492 

493 if current_doc.content.startswith(overlapping_str): # type: ignore 

494 # add split overlap information to this Document regarding the previous Document 

495 current_doc.meta["_split_overlap"].append({"doc_id": previous_doc.id, "range": overlapping_range}) 

496 

497 # add split overlap information to previous Document regarding this Document 

498 overlapping_range = (0, overlapping_range[1] - overlapping_range[0]) 

499 previous_doc.meta["_split_overlap"].append({"doc_id": current_doc.id, "range": overlapping_range}) 

500 

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

502 """ 

503 Serializes the component to a dictionary. 

504 """ 

505 serialized = default_to_dict( 

506 self, 

507 split_by=self.split_by, 

508 split_length=self.split_length, 

509 split_overlap=self.split_overlap, 

510 split_threshold=self.split_threshold, 

511 respect_sentence_boundary=self.respect_sentence_boundary, 

512 language=self.language, 

513 use_split_rules=self.use_split_rules, 

514 extend_abbreviations=self.extend_abbreviations, 

515 skip_empty_documents=self.skip_empty_documents, 

516 tokenizer_encoding=self.tokenizer_encoding, 

517 ) 

518 if self.splitting_function: 

519 serialized["init_parameters"]["splitting_function"] = serialize_callable(self.splitting_function) 

520 return serialized 

521 

522 @classmethod 

523 def from_dict(cls, data: dict[str, Any]) -> "DocumentSplitter": 

524 """ 

525 Deserializes the component from a dictionary. 

526 """ 

527 init_params = data.get("init_parameters", {}) 

528 

529 splitting_function = init_params.get("splitting_function", None) 

530 if splitting_function: 

531 init_params["splitting_function"] = deserialize_callable(splitting_function) 

532 

533 return default_from_dict(cls, data) 

534 

535 @staticmethod 

536 def _concatenate_sentences_based_on_word_amount( 

537 sentences: list[str], split_length: int, split_overlap: int 

538 ) -> tuple[list[str], list[int], list[int]]: 

539 """ 

540 Groups the sentences into chunks of `split_length` words while respecting sentence boundaries. 

541 

542 This function is only used when splitting by `word` and `respect_sentence_boundary` is set to `True`, i.e.: 

543 with NLTK sentence tokenizer. 

544 

545 :param sentences: The list of sentences to split. 

546 :param split_length: The maximum number of words in each split. 

547 :param split_overlap: The number of overlapping words in each split. 

548 :returns: A tuple containing the concatenated sentences, the start page numbers, and the start indices. 

549 """ 

550 # chunk information 

551 chunk_word_count = 0 

552 chunk_starting_page_number = 1 

553 chunk_start_idx = 0 

554 current_chunk: list[str] = [] 

555 # output lists 

556 split_start_page_numbers = [] 

557 list_of_splits: list[list[str]] = [] 

558 split_start_indices = [] 

559 

560 for sentence_idx, sentence in enumerate(sentences): 

561 current_chunk.append(sentence) 

562 chunk_word_count += len(sentence.split()) 

563 next_sentence_word_count = ( 

564 len(sentences[sentence_idx + 1].split()) if sentence_idx < len(sentences) - 1 else 0 

565 ) 

566 

567 # Number of words in the current chunk plus the next sentence is larger than the split_length, 

568 # or we reached the last sentence 

569 if (chunk_word_count + next_sentence_word_count) > split_length or sentence_idx == len(sentences) - 1: 

570 # Save current chunk and start a new one 

571 list_of_splits.append(current_chunk) 

572 split_start_page_numbers.append(chunk_starting_page_number) 

573 split_start_indices.append(chunk_start_idx) 

574 

575 # Get the number of sentences that overlap with the next chunk 

576 num_sentences_to_keep = DocumentSplitter._number_of_sentences_to_keep( 

577 sentences=current_chunk, split_length=split_length, split_overlap=split_overlap 

578 ) 

579 # Set up information for the new chunk 

580 if num_sentences_to_keep > 0: 

581 # Processed sentences are the ones that are not overlapping with the next chunk 

582 processed_sentences = current_chunk[:-num_sentences_to_keep] 

583 chunk_starting_page_number += sum(sent.count("\f") for sent in processed_sentences) 

584 chunk_start_idx += len("".join(processed_sentences)) 

585 # Next chunk starts with the sentences that were overlapping with the previous chunk 

586 current_chunk = current_chunk[-num_sentences_to_keep:] 

587 chunk_word_count = sum(len(s.split()) for s in current_chunk) 

588 else: 

589 # Here processed_sentences is the same as current_chunk since there is no overlap 

590 chunk_starting_page_number += sum(sent.count("\f") for sent in current_chunk) 

591 chunk_start_idx += len("".join(current_chunk)) 

592 current_chunk = [] 

593 chunk_word_count = 0 

594 

595 # Concatenate the sentences together within each split 

596 text_splits = [] 

597 for split in list_of_splits: 

598 text = "".join(split) 

599 if len(text) > 0: 

600 text_splits.append(text) 

601 

602 return text_splits, split_start_page_numbers, split_start_indices 

603 

604 @staticmethod 

605 def _number_of_sentences_to_keep(sentences: list[str], split_length: int, split_overlap: int) -> int: 

606 """ 

607 Returns the number of sentences to keep in the next chunk based on the `split_overlap` and `split_length`. 

608 

609 :param sentences: The list of sentences to split. 

610 :param split_length: The maximum number of words in each split. 

611 :param split_overlap: The number of overlapping words in each split. 

612 :returns: The number of sentences to keep in the next chunk. 

613 """ 

614 # If the split_overlap is 0, we don't need to keep any sentences 

615 if split_overlap == 0: 

616 return 0 

617 

618 num_sentences_to_keep = 0 

619 num_words = 0 

620 # Next overlapping Document should not start exactly the same as the previous one, so we skip the first sentence 

621 for sent in reversed(sentences[1:]): 

622 num_words += len(sent.split()) 

623 # If the number of words is larger than the split_length then don't add any more sentences 

624 if num_words > split_length: 

625 break 

626 num_sentences_to_keep += 1 

627 if num_words > split_overlap: 

628 break 

629 return num_sentences_to_keep