wehub-resource-sync
308 lines
8.4 KiB
Python
308 lines
8.4 KiB
Python
# Copyright (c) Microsoft. All rights reserved.
|
|
"""A Text splitter.
|
|
|
|
Split text in chunks, attempting to leave meaning intact.
|
|
For plain text, split looking at new lines first, then periods, and so on.
|
|
For markdown, split looking at punctuation first, and so on.
|
|
"""
|
|
|
|
import os
|
|
import re
|
|
from collections.abc import Callable
|
|
|
|
NEWLINE = os.linesep
|
|
|
|
TEXT_SPLIT_OPTIONS: list[list[str] | None] = [
|
|
["\n", "\r"],
|
|
["."],
|
|
["?", "!"],
|
|
[";"],
|
|
[":"],
|
|
[","],
|
|
[")", "]", "}"],
|
|
[" "],
|
|
["-"],
|
|
None,
|
|
]
|
|
|
|
MD_SPLIT_OPTIONS: list[list[str] | None] = [
|
|
["."],
|
|
["?", "!"],
|
|
[";"],
|
|
[":"],
|
|
[","],
|
|
[")", "]", "}"],
|
|
[" "],
|
|
["-"],
|
|
["\n", "\r"],
|
|
None,
|
|
]
|
|
|
|
|
|
def _token_counter(text: str) -> int:
|
|
"""Count the number of tokens in a string.
|
|
|
|
TODO: chunking methods should be configurable to allow for different
|
|
tokenization strategies depending on the model to be called.
|
|
For now, we use an extremely rough estimate.
|
|
"""
|
|
return len(text) // 4
|
|
|
|
|
|
def split_plaintext_lines(text: str, max_token_per_line: int, token_counter: Callable = _token_counter) -> list[str]:
|
|
"""Split plain text into lines.
|
|
|
|
it will split on new lines first, and then on punctuation.
|
|
"""
|
|
return _split_text_lines(
|
|
text=text,
|
|
max_token_per_line=max_token_per_line,
|
|
trim=True,
|
|
token_counter=token_counter,
|
|
)
|
|
|
|
|
|
def split_markdown_lines(text: str, max_token_per_line: int, token_counter: Callable = _token_counter) -> list[str]:
|
|
"""Split markdown into lines.
|
|
|
|
It will split on punctuation first, and then on space and new lines.
|
|
"""
|
|
return _split_markdown_lines(
|
|
text=text,
|
|
max_token_per_line=max_token_per_line,
|
|
trim=True,
|
|
token_counter=token_counter,
|
|
)
|
|
|
|
|
|
def split_plaintext_paragraph(text: list[str], max_tokens: int, token_counter: Callable = _token_counter) -> list[str]:
|
|
"""Split plain text into paragraphs."""
|
|
split_lines = []
|
|
for line in text:
|
|
split_lines.extend(
|
|
_split_text_lines(
|
|
text=line,
|
|
max_token_per_line=max_tokens,
|
|
trim=True,
|
|
token_counter=token_counter,
|
|
)
|
|
)
|
|
|
|
return _split_text_paragraph(text=split_lines, max_tokens=max_tokens, token_counter=token_counter)
|
|
|
|
|
|
def split_markdown_paragraph(text: list[str], max_tokens: int, token_counter: Callable = _token_counter) -> list[str]:
|
|
"""Split markdown into paragraphs."""
|
|
split_lines = []
|
|
for line in text:
|
|
split_lines.extend(
|
|
_split_markdown_lines(
|
|
text=line,
|
|
max_token_per_line=max_tokens,
|
|
trim=False,
|
|
token_counter=token_counter,
|
|
)
|
|
)
|
|
|
|
return _split_text_paragraph(text=split_lines, max_tokens=max_tokens, token_counter=token_counter)
|
|
|
|
|
|
def _split_text_paragraph(text: list[str], max_tokens: int, token_counter: Callable = _token_counter) -> list[str]:
|
|
"""Split text into paragraphs."""
|
|
if not text:
|
|
return []
|
|
|
|
paragraphs: list[str] = []
|
|
current_paragraph: list[str] = []
|
|
|
|
for line in text:
|
|
num_tokens_line = token_counter(line)
|
|
num_tokens_paragraph = token_counter("".join(current_paragraph))
|
|
|
|
if num_tokens_paragraph + num_tokens_line + 1 >= max_tokens and len(current_paragraph) > 0:
|
|
paragraphs.append("".join(current_paragraph).strip())
|
|
current_paragraph = []
|
|
|
|
current_paragraph.append(f"{line}{NEWLINE}")
|
|
|
|
if len(current_paragraph) > 0:
|
|
paragraphs.append("".join(current_paragraph).strip())
|
|
current_paragraph = []
|
|
|
|
# Distribute text more evenly in the last paragraphs
|
|
# when the last paragraph is too short.
|
|
|
|
if len(paragraphs) > 1:
|
|
last_para = paragraphs[-1]
|
|
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()
|
|
paragraphs.pop()
|
|
|
|
return paragraphs
|
|
|
|
|
|
def _split_markdown_lines(
|
|
text: str,
|
|
max_token_per_line: int,
|
|
trim: bool,
|
|
token_counter: Callable = _token_counter,
|
|
) -> list[str]:
|
|
"""Split markdown into lines."""
|
|
return _split_str_lines(
|
|
text=text,
|
|
max_tokens=max_token_per_line,
|
|
separators=MD_SPLIT_OPTIONS,
|
|
trim=trim,
|
|
token_counter=token_counter,
|
|
)
|
|
|
|
|
|
def _split_text_lines(
|
|
text: str,
|
|
max_token_per_line: int,
|
|
trim: bool,
|
|
token_counter: Callable = _token_counter,
|
|
) -> list[str]:
|
|
"""Split text into lines."""
|
|
return _split_str_lines(
|
|
text=text,
|
|
max_tokens=max_token_per_line,
|
|
separators=TEXT_SPLIT_OPTIONS,
|
|
trim=trim,
|
|
token_counter=token_counter,
|
|
)
|
|
|
|
|
|
def _split_str_lines(
|
|
text: str,
|
|
max_tokens: int,
|
|
separators: list[list[str] | None],
|
|
trim: bool,
|
|
token_counter: Callable = _token_counter,
|
|
) -> list[str]:
|
|
"""Split text into lines."""
|
|
if not text:
|
|
return []
|
|
|
|
text = text.replace("\r\n", "\n")
|
|
lines: list[str] = []
|
|
was_split = False
|
|
for split_option in separators:
|
|
if not lines:
|
|
lines, was_split = _split_str(
|
|
text=text,
|
|
max_tokens=max_tokens,
|
|
separators=split_option,
|
|
trim=trim,
|
|
token_counter=token_counter,
|
|
)
|
|
else:
|
|
lines, was_split = _split_list(
|
|
text=lines,
|
|
max_tokens=max_tokens,
|
|
separators=split_option,
|
|
trim=trim,
|
|
token_counter=token_counter,
|
|
)
|
|
if was_split:
|
|
break # pragma: no cover
|
|
|
|
return lines
|
|
|
|
|
|
def _split_str(
|
|
text: str,
|
|
max_tokens: int,
|
|
separators: list[str] | None,
|
|
trim: bool,
|
|
token_counter: Callable = _token_counter,
|
|
) -> tuple[list[str], bool]:
|
|
"""Split text into lines."""
|
|
input_was_split = False
|
|
if not text:
|
|
return [], input_was_split # pragma: no cover
|
|
|
|
if trim:
|
|
text = text.strip()
|
|
|
|
text_as_is = [text]
|
|
|
|
if token_counter(text) <= max_tokens:
|
|
return text_as_is, input_was_split
|
|
|
|
half = len(text) // 2
|
|
|
|
cutpoint = -1
|
|
|
|
if not separators:
|
|
cutpoint = half
|
|
elif set(separators) & set(text) and len(text) > 2:
|
|
regex_separators = re.compile("|".join(re.escape(s) for s in separators))
|
|
min_dist = half
|
|
for match in re.finditer(regex_separators, text):
|
|
end = match.end()
|
|
dist = abs(half - end)
|
|
if dist < min_dist:
|
|
min_dist = dist
|
|
cutpoint = end
|
|
elif end > half:
|
|
# distance is increasing, so we can stop searching
|
|
break
|
|
else:
|
|
return text_as_is, input_was_split
|
|
|
|
if 0 < cutpoint < len(text):
|
|
lines = []
|
|
for text_part in [text[:cutpoint], text[cutpoint:]]:
|
|
split, has_split = _split_str(
|
|
text=text_part,
|
|
max_tokens=max_tokens,
|
|
separators=separators,
|
|
trim=trim,
|
|
token_counter=token_counter,
|
|
)
|
|
lines.extend(split)
|
|
input_was_split = input_was_split or has_split
|
|
else:
|
|
return text_as_is, input_was_split
|
|
|
|
return lines, input_was_split
|
|
|
|
|
|
def _split_list(
|
|
text: list[str],
|
|
max_tokens: int,
|
|
separators: list[str] | None,
|
|
trim: bool,
|
|
token_counter: Callable = _token_counter,
|
|
) -> tuple[list[str], bool]:
|
|
"""Split list of string into lines."""
|
|
if not text:
|
|
return [], False # pragma: no cover
|
|
|
|
lines = []
|
|
input_was_split = False
|
|
for line in text:
|
|
split_str, was_split = _split_str(
|
|
text=line,
|
|
max_tokens=max_tokens,
|
|
separators=separators,
|
|
trim=trim,
|
|
token_counter=token_counter,
|
|
)
|
|
lines.extend(split_str)
|
|
input_was_split = input_was_split or was_split
|
|
|
|
return lines, input_was_split
|