Skip to content

Commit

Permalink
Fix token text splitter duplicates (#14848)
Browse files Browse the repository at this point in the history
- **Description:** 
- Add a break case to `text_splitter.py::split_text_on_tokens()` to
avoid unwanted item at the end of result.
    - Add a testcase to enforce the behavior.
  - **Issue:** 
    - #14649 
    - #5897
  - **Dependencies:** n/a,
 
---

**Quick illustration of change:**

```
text = "foo bar baz 123"

tokenizer = Tokenizer(
        chunk_overlap=3,
        tokens_per_chunk=7
)

output = split_text_on_tokens(text=text, tokenizer=tokenizer)
```
output before change: `["foo bar", "bar baz", "baz 123", "123"]`
output after change: `["foo bar", "bar baz", "baz 123"]`
  • Loading branch information
sweetkane authored Dec 19, 2023
1 parent 14d0418 commit ea331f3
Show file tree
Hide file tree
Showing 2 changed files with 19 additions and 0 deletions.
2 changes: 2 additions & 0 deletions libs/langchain/langchain/text_splitter.py
Original file line number Diff line number Diff line change
Expand Up @@ -670,6 +670,8 @@ def split_text_on_tokens(*, text: str, tokenizer: Tokenizer) -> List[str]:
chunk_ids = input_ids[start_idx:cur_idx]
while start_idx < len(input_ids):
splits.append(tokenizer.decode(chunk_ids))
if cur_idx == len(input_ids):
break
start_idx += tokenizer.tokens_per_chunk - tokenizer.chunk_overlap
cur_idx = min(start_idx + tokenizer.tokens_per_chunk, len(input_ids))
chunk_ids = input_ids[start_idx:cur_idx]
Expand Down
17 changes: 17 additions & 0 deletions libs/langchain/tests/unit_tests/test_text_splitter.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,8 @@
MarkdownHeaderTextSplitter,
PythonCodeTextSplitter,
RecursiveCharacterTextSplitter,
Tokenizer,
split_text_on_tokens,
)

FAKE_PYTHON_TEXT = """
Expand Down Expand Up @@ -1175,3 +1177,18 @@ def test_html_header_text_splitter(tmp_path: Path) -> None:
docs_from_file = splitter.split_text_from_file(tmp_path / "doc.html")

assert docs_from_file == expected


def test_split_text_on_tokens() -> None:
"""Test splitting by tokens per chunk."""
text = "foo bar baz 123"

tokenizer = Tokenizer(
chunk_overlap=3,
tokens_per_chunk=7,
decode=(lambda it: "".join(chr(i) for i in it)),
encode=(lambda it: [ord(c) for c in it]),
)
output = split_text_on_tokens(text=text, tokenizer=tokenizer)
expected_output = ["foo bar", "bar baz", "baz 123"]
assert output == expected_output

0 comments on commit ea331f3

Please sign in to comment.