176c04cd6c
Shared 512-token chunking module for the OB1/vault embedding stack. Wraps tiktoken cl100k_base; consumed by vault_mcp, ob1-enricher, ob1-reflector before passing to Ollama (nomic-embed-text, 768-d). Extracted from /opt/projects/homelab/shared/python/embedding_chunking to its own repo so consumers can pin via git ref instead of editable path source, unblocking docker builds outside the homelab tree.
56 lines
1.7 KiB
Python
56 lines
1.7 KiB
Python
from embedding_chunking import CHUNK_OVERLAP, CHUNK_SIZE, ENCODER, chunk_text
|
|
|
|
|
|
def _text_of_n_tokens(n: int) -> str:
|
|
"""Build a string that encodes to exactly n cl100k tokens.
|
|
|
|
Single ASCII letters encode to one token each, so n letters → n tokens.
|
|
"""
|
|
token_ids = ENCODER.encode("a " * n)[:n]
|
|
return ENCODER.decode(token_ids)
|
|
|
|
|
|
def test_short_text_returns_single_chunk():
|
|
text = _text_of_n_tokens(100)
|
|
chunks = chunk_text(text)
|
|
assert len(chunks) == 1
|
|
assert chunks[0].index == 0
|
|
assert chunks[0].token_count == 100
|
|
assert chunks[0].text == text
|
|
|
|
|
|
def test_1500_token_text_chunks_with_overlap():
|
|
# Brief asserted 3 chunks for 1500 tokens, but with step=448 the
|
|
# sliding-window starts are [0, 448, 896, 1344] → 4 chunks (three
|
|
# full 512-token windows plus a 156-token tail). Keeping the
|
|
# canonical sliding-window algorithm; updating the expectation.
|
|
text = _text_of_n_tokens(1500)
|
|
chunks = chunk_text(text)
|
|
|
|
step = CHUNK_SIZE - CHUNK_OVERLAP
|
|
assert len(chunks) == 4
|
|
assert [c.index for c in chunks] == [0, 1, 2, 3]
|
|
assert [c.token_count for c in chunks] == [
|
|
CHUNK_SIZE,
|
|
CHUNK_SIZE,
|
|
CHUNK_SIZE,
|
|
1500 - 3 * step,
|
|
]
|
|
|
|
tokens = ENCODER.encode(text)
|
|
for i, c in enumerate(chunks):
|
|
start = i * step
|
|
assert ENCODER.encode(c.text) == tokens[start : start + CHUNK_SIZE]
|
|
|
|
|
|
def test_600_token_text_returns_two_chunks():
|
|
text = _text_of_n_tokens(600)
|
|
chunks = chunk_text(text)
|
|
|
|
assert len(chunks) == 2
|
|
assert [c.index for c in chunks] == [0, 1]
|
|
|
|
step = CHUNK_SIZE - CHUNK_OVERLAP
|
|
assert chunks[0].token_count == CHUNK_SIZE
|
|
assert chunks[1].token_count == 600 - step
|