Skip to content

Python: Fix text chunker merging the last paragraph past max_tokens #14567

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We鈥檒l occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
Franco Ceballos Rastello (xThreeh) wants to merge 1 commit into from
Closed

Python: Fix text chunker merging the last paragraph past max_tokens #14567

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 4 additions & 10 deletions python/semantic_kernel/text/text_chunker.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -137,16 +137,10 @@ def _split_text_paragraph(text: list[str], max_tokens: int, token_counter: Calla
sec_last_para = paragraphs[-2]

if token_counter(last_para) < max_tokens / 4:
last_para_tokens = last_para.split(" ")
sec_last_para_tokens = sec_last_para.split(" ")
last_para_token_count = len(last_para_tokens)
sec_last_para_token_count = len(sec_last_para_tokens)

if last_para_token_count + sec_last_para_token_count <= max_tokens:
sec_last_para = " ".join(sec_last_para_tokens) + NEWLINE
last_para = " ".join(last_para_tokens)
new_sec_last_para = sec_last_para + last_para
paragraphs[-2] = new_sec_last_para.strip()
new_sec_last_para = f"{sec_last_para}{NEWLINE}{last_para}".strip()

if token_counter(new_sec_last_para) <= max_tokens:
paragraphs[-2] = new_sec_last_para
paragraphs.pop()

return paragraphs
Expand Down
39 changes: 31 additions & 8 deletions python/tests/unit/text/test_text_chunker.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -160,12 +160,28 @@ def test_split_text_paragraph_evenly():
"This is only a test.",
"We repeat, this is only a test. A unit test.",
"A small note. And another. And once again.",
f"Seriously, this is the end. We're finished. All set. Bye.{NEWLINE}Done.",
"Seriously, this is the end. We're finished. All set. Bye.",
"Done.",
]
split = split_plaintext_paragraph(text, max_token_per_line)
assert expected == split


def test_split_text_paragraph_evenly_respects_max_tokens():
"""Test split_paragraph() does not merge the last paragraph past max_tokens"""
text = [
"This is a test of the emergency broadcast system. This is only a test.",
"We repeat, this is only a test. A unit test.",
"A small note. And another. And once again. Seriously, this is the end. " + "We're finished. All set. Bye.",
"Done.",
]

max_token_per_line = 15

split = split_plaintext_paragraph(text, max_token_per_line)
assert all(len(paragraph) // 4 <= max_token_per_line for paragraph in split)


def test_split_text_paragraph_evenly_2():
"""Test split_paragraph() with evenly split input"""
text = [
Expand All @@ -187,7 +203,8 @@ def test_split_text_paragraph_evenly_2():
"The sun set over the horizon peacefully, the beautiful star.",
f"Cats love boxes.{NEWLINE}That is something. Incredible news that is.",
f"What a beautiful day to be alive.{NEWLINE}Seriously, this is the end.",
f"We're finished once of for all. All set. Ok.{NEWLINE}Done.{NEWLINE}" + f"Or is it?{NEWLINE}Surprise!",
f"We're finished once of for all. All set. Ok.{NEWLINE}Done.{NEWLINE}" + "Or is it?",
"Surprise!",
]
split = split_plaintext_paragraph(text, max_token_per_line)
assert expected == split
Expand Down Expand Up @@ -247,7 +264,8 @@ def test_split_paragraph_semicolon():
"This is only a test",
"We repeat; this is only a test; A unit test",
"A small note; And another; And once again;",
f"Seriously, this is the end; We're finished; All set; Bye.{NEWLINE}Done.",
"Seriously, this is the end; We're finished; All set; Bye.",
"Done.",
]
max_token_per_line = 15
split = split_plaintext_paragraph(text, max_token_per_line)
Expand All @@ -267,7 +285,8 @@ def test_split_paragraph_colon():
"This is only a test",
"We repeat: this is only a test: A unit test",
"A small note: And another: And once again:",
f"Seriously, this is the end: We're finished: All set: Bye.{NEWLINE}Done.",
"Seriously, this is the end: We're finished: All set: Bye.",
"Done.",
]
max_token_per_line = 15
split = split_plaintext_paragraph(text, max_token_per_line)
Expand Down Expand Up @@ -388,7 +407,8 @@ def test_split_md_on_dot():
"This\n is only a test",
"We repeat. this is only a test. A unit test",
"A small note. And another. And once again.",
f"Seriously, this is the end. We're finished. All set. Bye.{NEWLINE}Done.",
"Seriously, this is the end. We're finished. All set. Bye.",
"Done.",
]
max_token_per_line = 15
split = split_markdown_paragraph(text, max_token_per_line)
Expand All @@ -408,7 +428,8 @@ def test_split_md_on_colon():
"This is only a test",
"We repeat: this is only a test: A unit test",
"A small note: And another: And once again:",
f"Seriously, this is the end: We're finished: All set: Bye.{NEWLINE}Done.",
"Seriously, this is the end: We're finished: All set: Bye.",
"Done.",
]
max_token_per_line = 15
split = split_markdown_paragraph(text, max_token_per_line)
Expand All @@ -428,7 +449,8 @@ def test_split_md_on_punctuation():
"This\n is only a test",
"We repeat? this is only a test! A unit test",
"A small note? And another! And once again?",
f"Seriously, this is the end! We're finished! All set! Bye.{NEWLINE}Done.",
"Seriously, this is the end! We're finished! All set! Bye.",
"Done.",
]
max_token_per_line = 15
split = split_markdown_paragraph(text, max_token_per_line)
Expand All @@ -448,7 +470,8 @@ def test_split_md_on_semicolon():
"This is only a test",
"We repeat; this is only a test; A unit test",
"A small note; And another; And once again;",
f"Seriously, this is the end; We're finished; All set; Bye.{NEWLINE}Done.",
"Seriously, this is the end; We're finished; All set; Bye.",
"Done.",
]
max_token_per_line = 15
split = split_markdown_paragraph(text, max_token_per_line)
Expand Down
Loading