Skip to content
Open
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
68 changes: 49 additions & 19 deletions src/chonkie/chunker/semantic.py
Original file line number Diff line number Diff line change
Expand Up @@ -395,6 +395,35 @@ def _group_sentences(

return groups

def _split_oversized_sentence(self, sentence: Sentence) -> list[Sentence]:
"""Split a single sentence that exceeds chunk_size into token-sized pieces.

A sentence longer than chunk_size can never fit into any group, so grouping
alone cannot enforce the limit. We split its text at the token level, the same
way TokenChunker does, so no resulting piece exceeds chunk_size.

Args:
sentence: A sentence whose token_count exceeds chunk_size

Returns:
List of sentences, each with token_count <= chunk_size

"""
tokens = self.tokenizer.encode(sentence.text)
token_groups = [
tokens[i : i + self.chunk_size] for i in range(0, len(tokens), self.chunk_size)
]
texts = self.tokenizer.decode_batch(token_groups)
Comment thread
coderabbitai[bot] marked this conversation as resolved.
Outdated
return [
Sentence(
text=text,
start_index=sentence.start_index,
end_index=sentence.start_index + len(text),
token_count=len(group),
)
for text, group in zip(texts, token_groups)
]

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

In the current implementation, all split Sentence pieces are created with the same start_index (i.e., sentence.start_index). This results in overlapping and incorrect index ranges for the split sentences. Instead, we should accumulate the character lengths of the processed pieces to correctly offset the start_index and end_index for each subsequent piece.

        sentences = []
        current_start = sentence.start_index
        for text, group in zip(texts, token_groups):
            sentences.append(
                Sentence(
                    text=text,
                    start_index=current_start,
                    end_index=current_start + len(text),
                    token_count=len(group),
                )
            )
            current_start += len(text)
        return sentences


def _split_groups(self, groups: list[list[Sentence]]) -> list[list[Sentence]]:
"""Split groups that exceed chunk_size into smaller groups.

Expand All @@ -418,14 +447,23 @@ def _split_groups(self, groups: list[list[Sentence]]) -> list[list[Sentence]]:
current_token_count = 0

for sentence in group:
if current_token_count + sentence.token_count <= self.chunk_size:
current_group.append(sentence)
current_token_count += sentence.token_count
else:
if current_group:
final_groups.append(current_group)
current_group = [sentence]
current_token_count = sentence.token_count
# A single sentence larger than chunk_size can't fit in any
# group, so split it into token-sized pieces first.
pieces = (
self._split_oversized_sentence(sentence)
if sentence.token_count > self.chunk_size
else [sentence]
)

for piece in pieces:
if current_token_count + piece.token_count <= self.chunk_size:
current_group.append(piece)
current_token_count += piece.token_count
else:
if current_group:
final_groups.append(current_group)
current_group = [piece]
current_token_count = piece.token_count

if current_group:
final_groups.append(current_group)
Expand Down Expand Up @@ -465,18 +503,10 @@ def chunk(self, text: str) -> list[Chunk]:

# Handle edge cases - too few sentences
if len(sentences) <= self.similarity_window:
# If we have any sentences, return them as a single chunk
# If we have any sentences, group them together, still enforcing
# chunk_size so a few large sentences don't produce an oversized chunk.
if sentences:
text = "".join([s.text for s in sentences])
token_count = sum([s.token_count for s in sentences])
return [
Chunk(
text=text,
start_index=0,
end_index=len(text),
token_count=token_count,
),
]
return self._create_chunks(self._split_groups([sentences]))
else:
return []

Expand Down
50 changes: 50 additions & 0 deletions tests/chunkers/test_semantic_chunker.py
Original file line number Diff line number Diff line change
Expand Up @@ -75,6 +75,56 @@ def test_semantic_chunker_chunking(embedding_model: BaseEmbeddings, sample_text:
assert all([chunk.end_index is not None for chunk in chunks])


def test_semantic_chunker_oversized_single_sentence(embedding_model: BaseEmbeddings) -> None:
"""A single sentence longer than chunk_size must still be split to respect the limit.

Regression test for #222: a sentence with no delimiters that exceeds chunk_size
was emitted as one oversized chunk because grouping only splits between sentences.
"""
chunk_size = 100
chunker = SemanticChunker(
embedding_model=embedding_model,
chunk_size=chunk_size,
threshold=0.5,
min_sentences_per_chunk=1,
)

# One long run of words with no sentence delimiters => a single "sentence"
# far larger than chunk_size, surrounded by normal sentences.
normal = "This is a normal sentence. Here is another one. And a third here. A fourth now. "
oversized = "word " * 800
text = normal + oversized + ". " + normal

chunks = chunker.chunk(text)

assert len(chunks) > 0
assert all([chunk.token_count <= chunk_size for chunk in chunks]), (
"Every chunk must respect chunk_size even when a single sentence is oversized"
)


def test_semantic_chunker_oversized_sentence_below_similarity_window(
embedding_model: BaseEmbeddings,
) -> None:
"""The few-sentences early-return path must also enforce chunk_size (#222)."""
chunk_size = 100
chunker = SemanticChunker(
embedding_model=embedding_model,
chunk_size=chunk_size,
threshold=0.5,
similarity_window=3,
)

# Fewer sentences than similarity_window, but one is far larger than chunk_size.
text = "Short intro. " + ("word " * 600)
chunks = chunker.chunk(text)

assert len(chunks) > 1
assert all([chunk.token_count <= chunk_size for chunk in chunks]), (
"The few-sentences path must also split oversized sentences"
)


def test_semantic_chunker_empty_text(embedding_model: BaseEmbeddings) -> None:
"""Test that the SemanticChunker can handle empty text input."""
chunker = SemanticChunker(
Expand Down