-
Notifications
You must be signed in to change notification settings - Fork 348
fix(semantic): split sentences that exceed chunk_size #627
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from 1 commit
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -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) | ||
| 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) | ||
| ] | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. In the current implementation, all split 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. | ||
|
|
||
|
|
@@ -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) | ||
|
|
@@ -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 [] | ||
|
|
||
|
|
||
Uh oh!
There was an error while loading. Please reload this page.