Julian Carlson
57 lines
1.9 KiB
Python
57 lines
1.9 KiB
Python
import tiktoken
|
|
from .config import CONFIG
|
|
|
|
_enc = tiktoken.get_encoding("cl100k_base")
|
|
|
|
|
|
def count_tokens(text):
|
|
return len(_enc.encode(text))
|
|
|
|
|
|
def chunk_text(text, chunk_size=None, overlap=None):
|
|
chunk_size = chunk_size or CONFIG.chunk_size_tokens
|
|
overlap = overlap or CONFIG.chunk_overlap_tokens
|
|
|
|
paragraphs = [p.strip() for p in text.split("\n\n") if p.strip()]
|
|
|
|
chunks = []
|
|
current_chunk = []
|
|
current_tokens = 0
|
|
|
|
for para in paragraphs:
|
|
para_tokens = count_tokens(para)
|
|
|
|
if para_tokens > chunk_size:
|
|
sentences = para.replace("\n", " ").split(". ")
|
|
for sent in sentences:
|
|
if not sent.strip():
|
|
continue
|
|
sent_tokens = count_tokens(sent)
|
|
if current_tokens + sent_tokens > chunk_size and current_chunk:
|
|
chunks.append("\n\n".join(current_chunk))
|
|
current_chunk = current_chunk[-1:] if current_chunk else []
|
|
current_tokens = count_tokens("\n\n".join(current_chunk)) if current_chunk else 0
|
|
current_chunk.append(sent)
|
|
current_tokens += sent_tokens
|
|
else:
|
|
if current_tokens + para_tokens > chunk_size and current_chunk:
|
|
chunks.append("\n\n".join(current_chunk))
|
|
overlap_tokens = 0
|
|
overlap_paras = []
|
|
for p in reversed(current_chunk):
|
|
pt = count_tokens(p)
|
|
if overlap_tokens + pt <= overlap:
|
|
overlap_paras.insert(0, p)
|
|
overlap_tokens += pt
|
|
else:
|
|
break
|
|
current_chunk = overlap_paras
|
|
current_tokens = overlap_tokens
|
|
current_chunk.append(para)
|
|
current_tokens += para_tokens
|
|
|
|
if current_chunk:
|
|
chunks.append("\n\n".join(current_chunk))
|
|
|
|
return chunks
|