31 lines
1.1 KiB
Python
31 lines
1.1 KiB
Python
"""Lossless character boundaries constrained by an actual token counter."""
|
|
from collections.abc import Callable
|
|
|
|
|
|
def split_to_fit(text: str, fits: Callable[[str], bool]) -> list[str]:
|
|
"""Keep every character; use natural boundaries when they fit near the end."""
|
|
pieces = []
|
|
while text:
|
|
if fits(text):
|
|
pieces.append(text)
|
|
break
|
|
low, high = 0, len(text)
|
|
while low < high:
|
|
middle = (low + high + 1) // 2
|
|
if fits(text[:middle]):
|
|
low = middle
|
|
else:
|
|
high = middle - 1
|
|
if low == 0:
|
|
raise ValueError('Token budget cannot contain one Unicode character.')
|
|
# Avoid breaking a word unless its length makes that unavoidable.
|
|
boundary = max(text.rfind('\n', 0, low), text.rfind(' ', 0, low)) + 1
|
|
cut = boundary if boundary >= low * 0.7 else low
|
|
while cut and not fits(text[:cut]):
|
|
cut -= 1
|
|
if not cut:
|
|
raise ValueError('No valid token boundary.')
|
|
pieces.append(text[:cut])
|
|
text = text[cut:]
|
|
return pieces
|