klas/t6_mem0
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

[Improvement] Add support for min chunk size (#1007)

Browse Source
This commit is contained in:
Deven Patel
2023-12-15 05:59:15 +05:30
committed by GitHub
parent 9303a1bf81
commit c0ee680546
11 changed files with 59 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
tests/chunkers/test_base_chunker.py
View File

@@ -4,6 +4,7 @@ from unittest.mock import MagicMock
import pytest
from embedchain.chunkers.base_chunker import BaseChunker
from embedchain.config.add_config import ChunkerConfig
from embedchain.models.data_type import DataType
@@ -35,6 +36,18 @@ def chunker(text_splitter_mock, data_type):
return chunker
def test_create_chunks_with_config(chunker, text_splitter_mock, loader_mock, app_id, data_type):
text_splitter_mock.split_text.return_value = ["Chunk 1", "long chunk"]
loader_mock.load_data.return_value = {
"data": [{"content": "Content 1", "meta_data": {"url": "URL 1"}}],
"doc_id": "DocID",
}
config = ChunkerConfig(chunk_size=50, chunk_overlap=0, length_function=len, min_chunk_size=10)
result = chunker.create_chunks(loader_mock, "test_src", app_id, config)
assert result["documents"] == ["long chunk"]
def test_create_chunks(chunker, text_splitter_mock, loader_mock, app_id, data_type):
text_splitter_mock.split_text.return_value = ["Chunk 1", "Chunk 2"]
loader_mock.load_data.return_value = {

4
tests/chunkers/test_image_chunker.py
View File

@@ -11,7 +11,7 @@ class TestImageChunker(unittest.TestCase):
Test the chunks generated by TextChunker.
# TODO: Not a very precise test.
"""
chunker_config = ChunkerConfig(chunk_size=1, chunk_overlap=0, length_function=len)
chunker_config = ChunkerConfig(chunk_size=1, chunk_overlap=0, length_function=len, min_chunk_size=0)
chunker = ImagesChunker(config=chunker_config)
# Data type must be set manually in the test
chunker.set_data_type(DataType.IMAGES)
@@ -51,7 +51,7 @@ class TestImageChunker(unittest.TestCase):
self.assertEqual(expected_chunks, result)
def test_word_count(self):
chunker_config = ChunkerConfig(chunk_size=1, chunk_overlap=0, length_function=len)
chunker_config = ChunkerConfig(chunk_size=1, chunk_overlap=0, length_function=len, min_chunk_size=0)
chunker = ImagesChunker(config=chunker_config)
chunker.set_data_type(DataType.IMAGES)

18
tests/chunkers/test_text.py
View File

@@ -10,12 +10,12 @@ class TestTextChunker:
"""
Test the chunks generated by TextChunker.
"""
chunker_config = ChunkerConfig(chunk_size=10, chunk_overlap=0, length_function=len)
chunker_config = ChunkerConfig(chunk_size=10, chunk_overlap=0, length_function=len, min_chunk_size=0)
chunker = TextChunker(config=chunker_config)
text = "Lorem ipsum dolor sit amet, consectetur adipiscing elit."
# Data type must be set manually in the test
chunker.set_data_type(DataType.TEXT)
result = chunker.create_chunks(MockLoader(), text)
result = chunker.create_chunks(MockLoader(), text, chunker_config)
documents = result["documents"]
assert len(documents) > 5
@@ -23,11 +23,11 @@ class TestTextChunker:
"""
Test the chunks generated by TextChunker with app_id
"""
chunker_config = ChunkerConfig(chunk_size=10, chunk_overlap=0, length_function=len)
chunker_config = ChunkerConfig(chunk_size=10, chunk_overlap=0, length_function=len, min_chunk_size=0)
chunker = TextChunker(config=chunker_config)
text = "Lorem ipsum dolor sit amet, consectetur adipiscing elit."
chunker.set_data_type(DataType.TEXT)
result = chunker.create_chunks(MockLoader(), text)
result = chunker.create_chunks(MockLoader(), text, chunker_config)
documents = result["documents"]
assert len(documents) > 5
@@ -35,12 +35,12 @@ class TestTextChunker:
"""
Test that if an infinitely high chunk size is used, only one chunk is returned.
"""
chunker_config = ChunkerConfig(chunk_size=9999999999, chunk_overlap=0, length_function=len)
chunker_config = ChunkerConfig(chunk_size=9999999999, chunk_overlap=0, length_function=len, min_chunk_size=0)
chunker = TextChunker(config=chunker_config)
text = "Lorem ipsum dolor sit amet, consectetur adipiscing elit."
# Data type must be set manually in the test
chunker.set_data_type(DataType.TEXT)
result = chunker.create_chunks(MockLoader(), text)
result = chunker.create_chunks(MockLoader(), text, chunker_config)
documents = result["documents"]
assert len(documents) == 1
@@ -48,18 +48,18 @@ class TestTextChunker:
"""
Test that if a chunk size of one is used, every character is a chunk.
"""
chunker_config = ChunkerConfig(chunk_size=1, chunk_overlap=0, length_function=len)
chunker_config = ChunkerConfig(chunk_size=1, chunk_overlap=0, length_function=len, min_chunk_size=0)
chunker = TextChunker(config=chunker_config)
# We can't test with lorem ipsum because chunks are deduped, so would be recurring characters.
text = """0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ!"#$%&\'()*+,-./:;<=>?@[\\]^_`{|}~ \t\n\r\x0b\x0c"""
# Data type must be set manually in the test
chunker.set_data_type(DataType.TEXT)
result = chunker.create_chunks(MockLoader(), text)
result = chunker.create_chunks(MockLoader(), text, chunker_config)
documents = result["documents"]
assert len(documents) == len(text)
def test_word_count(self):
chunker_config = ChunkerConfig(chunk_size=1, chunk_overlap=0, length_function=len)
chunker_config = ChunkerConfig(chunk_size=1, chunk_overlap=0, length_function=len, min_chunk_size=0)
chunker = TextChunker(config=chunker_config)
chunker.set_data_type(DataType.TEXT)

2
tests/embedchain/test_add.py
View File

@@ -33,7 +33,7 @@ def test_add_forced_type(app):
def test_dry_run(app):
chunker_config = ChunkerConfig(chunk_size=1, chunk_overlap=0)
chunker_config = ChunkerConfig(chunk_size=1, chunk_overlap=0, min_chunk_size=0)
text = """0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"""
result = app.add(source=text, config=AddConfig(chunker=chunker_config), dry_run=True)