diff --git a/src/raglite/_split_chunklets.py b/src/raglite/_split_chunklets.py index 4fd92cb6..a8295ff2 100644 --- a/src/raglite/_split_chunklets.py +++ b/src/raglite/_split_chunklets.py @@ -1,5 +1,6 @@ """Split a document into chunklets.""" +import re from collections.abc import Callable import numpy as np @@ -14,8 +15,10 @@ def markdown_chunklet_boundaries(sentences: list[str]) -> FloatVector: doc = "".join(sentences) md = MarkdownIt() tokens = md.parse(doc) - # Identify the character index of each line in the document. - lines = doc.splitlines(keepends=True) + # Identify the character index of each line in the document. Split lines like markdown-it does + # (on LF, CRLF, and CR only) rather than with `str.splitlines`, which also splits on form feeds, + # U+2028, etc., so that the line numbers in `token.map` stay aligned with the character indexes. + lines = re.findall(r"[^\r\n]*(?:\r\n|\r|\n)|[^\r\n]+", doc) line_start_char = [0] for line in lines[:-1]: line_start_char.append(line_start_char[-1] + len(line)) diff --git a/src/raglite/_split_sentences.py b/src/raglite/_split_sentences.py index d383cb51..7dea00c7 100644 --- a/src/raglite/_split_sentences.py +++ b/src/raglite/_split_sentences.py @@ -1,5 +1,6 @@ """Split a document into sentences.""" +import re from collections import deque from collections.abc import Callable from functools import cache @@ -28,7 +29,10 @@ def get_markdown_heading_indexes(doc: str) -> list[tuple[int, int]]: md = MarkdownIt() tokens = md.parse(doc) headings = [] - lines = doc.splitlines(keepends=True) + # Split lines like markdown-it does (on LF, CRLF, and CR only) rather than with + # `str.splitlines`, which also splits on form feeds, U+2028, etc., so that the line numbers + # in `token.map` stay aligned with the character indexes. + lines = re.findall(r"[^\r\n]*(?:\r\n|\r|\n)|[^\r\n]+", doc) line_start_char = [0] for line in lines: line_start_char.append(line_start_char[-1] + len(line)) diff --git a/tests/test_split_chunklets.py b/tests/test_split_chunklets.py index e3a8a22a..1d74acd6 100644 --- a/tests/test_split_chunklets.py +++ b/tests/test_split_chunklets.py @@ -2,7 +2,7 @@ import pytest -from raglite._split_chunklets import split_chunklets +from raglite._split_chunklets import markdown_chunklet_boundaries, split_chunklets @pytest.mark.parametrize( @@ -102,3 +102,28 @@ def test_split_chunklets(sentences_splits: tuple[list[str], list[int]]) -> None: chunklet == expected_chunklet for chunklet, expected_chunklet in zip(chunklets, expected_chunklets, strict=True) ) + + +@pytest.mark.parametrize( + "line_separator", + [ + pytest.param("\n", id="lf"), + pytest.param("\r\n", id="crlf"), + pytest.param("\r", id="cr"), + pytest.param("\x0b", id="vertical_tab"), + pytest.param("\x0c", id="form_feed"), + pytest.param("\x1c", id="file_separator"), + pytest.param("\x1d", id="group_separator"), + pytest.param("\x1e", id="record_separator"), + pytest.param("\x85", id="next_line"), + pytest.param("\u2028", id="line_separator"), + pytest.param("\u2029", id="paragraph_separator"), + ], +) +def test_markdown_chunklet_boundaries_line_separators(line_separator: str) -> None: + """Test that boundaries stay aligned when a document has exotic line separators.""" + # Only LF, CRLF, and CR are line endings to markdown-it, so the other separators must not shift + # the line numbers that markdown-it reports for the paragraphs and the heading. + sentences = [f"Intro one.{line_separator}", "Intro two.\n\n", "# Real Heading\n\n", "Body.\n"] + boundary_probas = markdown_chunklet_boundaries(sentences) + assert boundary_probas.tolist() == [0.5, 0.0, 1.0, 0.0] diff --git a/tests/test_split_sentences.py b/tests/test_split_sentences.py index 3e8e31b7..7aa02bc9 100644 --- a/tests/test_split_sentences.py +++ b/tests/test_split_sentences.py @@ -2,10 +2,11 @@ from pathlib import Path +import numpy as np import pytest from raglite._markdown import document_to_markdown -from raglite._split_sentences import split_sentences +from raglite._split_sentences import markdown_sentence_boundaries, split_sentences def test_split_sentences() -> None: @@ -77,3 +78,32 @@ def test_split_sentences_edge_cases(case: tuple[str, list[str], tuple[int, int | sentences[: len(expected_sentences)], expected_sentences, strict=True ) ) + + +@pytest.mark.parametrize( + "line_separator", + [ + pytest.param("\n", id="lf"), + pytest.param("\r\n", id="crlf"), + pytest.param("\r", id="cr"), + pytest.param("\x0b", id="vertical_tab"), + pytest.param("\x0c", id="form_feed"), + pytest.param("\x1c", id="file_separator"), + pytest.param("\x1d", id="group_separator"), + pytest.param("\x1e", id="record_separator"), + pytest.param("\x85", id="next_line"), + pytest.param("\u2028", id="line_separator"), + pytest.param("\u2029", id="paragraph_separator"), + ], +) +def test_markdown_sentence_boundaries_line_separators(line_separator: str) -> None: + """Test that heading boundaries stay aligned when a document has exotic line separators.""" + # Only LF, CRLF, and CR are line endings to markdown-it, so the other separators must not shift + # the line numbers that markdown-it reports for the heading. + heading = "# Real Heading\n" + doc = f"Intro one.{line_separator}Intro two.\n\n{heading}\nBody.\n" + heading_start = doc.index(heading) + boundary_probas = markdown_sentence_boundaries(doc) + marked_indexes = np.flatnonzero(~np.isnan(boundary_probas)) + assert marked_indexes.min() == heading_start - 1 + assert marked_indexes.max() == heading_start + len(heading)