Spaces:
Running
Running
File size: 2,383 Bytes
d8c1ecb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 | import re
from pilotcore.chunking.base import BaseChunker
from pilotcore.retrieval.embeddings import get_embedding_tokenizer
class TokenChunker(BaseChunker):
"""
Token-based chunker using the embedding model tokenizer.
Attempts to preserve natural sentence boundaries while
respecting the desired token budget.
"""
def chunk(
self,
text: str,
chunk_size: int = 256,
overlap: int = 40,
) -> list[dict]:
tokenizer = get_embedding_tokenizer()
token_ids = tokenizer.encode(
text,
add_special_tokens=False,
)
chunks = []
start = 0
while start < len(token_ids):
end = min(start + chunk_size, len(token_ids))
candidate_ids = token_ids[start:end]
candidate = tokenizer.decode(
candidate_ids,
skip_special_tokens=True,
).strip()
# Last chunk
if end == len(token_ids):
if candidate:
chunks.append(candidate)
break
# Try to end at a natural boundary
boundary = self._find_boundary(candidate)
if boundary != -1:
candidate = candidate[: boundary + 1].strip()
actual_ids = tokenizer.encode(
candidate,
add_special_tokens=False,
)
end = start + len(actual_ids)
if candidate:
chunks.append(candidate)
next_start = end - overlap
if next_start <= start:
next_start = end
start = next_start
return [
{
"text": chunk,
"metadata": {},
}
for chunk in chunks
]
@staticmethod
def _find_boundary(text: str) -> int:
"""
Find the last 'good' split location.
"""
boundaries = []
for match in re.finditer(r"\n\n", text):
boundaries.append(match.end() - 1)
for match in re.finditer(r"[.!?]\s", text):
boundaries.append(match.end() - 2)
for match in re.finditer(r"\n", text):
boundaries.append(match.end() - 1)
if not boundaries:
return -1
return max(boundaries)
|