Raise ValueError for non-positive batch_size in iter_batch

fastembed/common/utils.py

@@ -43,16 +43,23 @@ def last_token_pooling(input_array: NumpyArray, attention_mask: NDArray[np.int64
 
 
 def iter_batch(iterable: Iterable[T], size: int) -> Iterable[list[T]]:
-    """
+    """Validate the batch size immediately and consume the iterable lazily.
+
     >>> list(iter_batch([1,2,3,4,5], 3))
     [[1, 2, 3], [4, 5]]
     """
-    source_iter = iter(iterable)
-    while source_iter:
-        b = list(islice(source_iter, size))
-        if len(b) == 0:
-            break
-        yield b
+    if size < 1:
+        raise ValueError(f"batch_size must be >= 1, got {size}")
+
+    def batches() -> Iterable[list[T]]:
+        source_iter = iter(iterable)
+        while source_iter:
+            b = list(islice(source_iter, size))
+            if len(b) == 0:
+                break
+            yield b
+
+    return batches()
 
 
 def define_cache_dir(cache_dir: str | None = None) -> Path:

tests/test_common.py

@@ -1,4 +1,5 @@
 import numpy as np
+import pytest
 
 from fastembed import (
     TextEmbedding,
@@ -7,7 +8,7 @@
     LateInteractionMultimodalEmbedding,
     LateInteractionTextEmbedding,
 )
-from fastembed.common.utils import last_token_pooling
+from fastembed.common.utils import iter_batch, last_token_pooling
 
 
 def test_text_list_supported_models():
@@ -59,3 +60,15 @@ def test_last_token_pooling_with_left_padding():
     pooled = last_token_pooling(token_embeddings, attention_mask)
 
     assert np.allclose(pooled, [[2.0, 2.0], [6.0, 6.0]])
+
+
+def test_iter_batch_accepts_positive_size():
+    assert list(iter_batch([1, 2, 3, 4, 5], 3)) == [[1, 2, 3], [4, 5]]
+
+
+def test_iter_batch_rejects_non_positive_size():
+    with pytest.raises(ValueError, match="batch_size must be >= 1, got 0"):
+        iter_batch([1, 2, 3], 0)
+
+    with pytest.raises(ValueError, match="batch_size must be >= 1, got -1"):
+        iter_batch([1, 2, 3], -1)