Fix tests for text type outputs

chonkie-ai · Jan 7, 2025 · 6ad760b · 6ad760b
1 parent 825cc03
commit 6ad760b
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/tests/chunker/test_sdpm_chunker.py b/tests/chunker/test_sdpm_chunker.py
@@ -179,8 +179,9 @@ def test_sdpm_chunker_return_type(embedding_model, sample_text):
     """Test that SDPMChunker's return type is correctly set."""
     chunker = SDPMChunker(embedding_model=embedding_model, chunk_size=512, threshold=0.5, return_type="texts")
     chunks = chunker.chunk(sample_text)
+    tokenizer = embedding_model.get_tokenizer_or_token_counter()
     assert all([type(chunk) is str for chunk in chunks])
-    assert all([len(embedding_model.encode(chunk)) <= 512 for chunk in chunks])
+    assert all([len(tokenizer.encode(chunk)) <= 512 for chunk in chunks])
 
 if __name__ == "__main__":
     pytest.main()
diff --git a/tests/chunker/test_semantic_chunker.py b/tests/chunker/test_semantic_chunker.py
@@ -291,8 +291,9 @@ def test_semantic_chunker_return_type(embedding_model, sample_text):
     """Test that SemanticChunker's return type is correctly set."""
     chunker = SemanticChunker(embedding_model=embedding_model, chunk_size=512, threshold=0.5, return_type="texts")
     chunks = chunker.chunk(sample_text)
+    tokenizer = embedding_model.get_tokenizer_or_token_counter()
     assert all([type(chunk) is str for chunk in chunks])
-    assert all([len(embedding_model.encode(chunk)) <= 512 for chunk in chunks])
+    assert all([len(tokenizer.encode(chunk)) <= 512 for chunk in chunks])
 
 if __name__ == "__main__":
     pytest.main()