Added batch embedding capability.

Author: travis-bauer

CHANGELOG.md

@@ -1,6 +1,10 @@
 # Changelog
 
 ## Unreleased
+- Added batch embedding support: embedding adapters expose `execute_one` and `execute_batch`;
+  `execute` remains as a deprecated alias for `execute_one` (removed in 1.0). `llmEmbed` uses
+  `execute_one` / `execute_batch` only, accepts list-shaped items (e.g. after `makeLists`), and
+  optional `batch_size` for built-in buffering.
 - Added model2vec embedding support via `Model2VecEmbeddingAdapter` and `llmEmbed`
   source `model2vec`, using in-process static embeddings with offline-ready HF cache
   precaching. Install with `pip install talkpipe[model2vec]` or `talkpipe[all]`. Added

docs/guides/model-and-source-configuration.md

@@ -223,7 +223,7 @@ If a CLI flag is omitted and the matching `DEFAULT_*` key is unset, the value pa
 
 ## Segment parameters
 
-Those configuration keys provide the fallback values; the per-segment parameters below override them and take final precedence at construction time.
+For `llmPrompt`, `llmVisionPrompt`, and `llmEmbed`, only **`model`** and **`source`** fall back to `default_*` config keys when omitted. Every other segment parameter must be set on the segment (ChatterLang or Python); it is not read from `~/.talkpipe.toml` or `TALKPIPE_*` unless noted below for a specific higher-level segment.
 
 ### `llmPrompt` / `LLMPrompt`
 
@@ -269,7 +269,26 @@ segment = LLMVisionPrompt(
 
 ### `llmEmbed` / `LLMEmbed`
 
-Required (directly or via config): `model`, `source`. Optional: `field` (text field to embed), `set_as` (field to store the vector on the item).
+| Parameter | From config? | Notes |
+|-----------|--------------|--------|
+| `model` | Yes — `default_embedding_model_name` | Required if not passed on the segment |
+| `source` | Yes — `default_embedding_model_source` | Required if not passed on the segment |
+| `field` | No | Text field to embed on structured items |
+| `set_as` | No | Field on the item where the vector is stored |
+| `batch_size` | No | Scalar items per provider call (default `1`) |
+| `fail_on_error` | No | Default `true` |
+
+**Batching (two patterns):**
+
+1. **Built-in buffering** — set `batch_size` greater than `1` on `llmEmbed` to amortize API round-trips without changing upstream segments.
+2. **Composable buffering** — group items with `makeLists`, then embed the batch in one call:
+
+```chatterlang
+| makeLists[num_items=100, field="_"]
+| llmEmbed[model="mxbai-embed-large", source="ollama", field="content", set_as="vector"]
+```
+
+List-shaped items are expanded back to one output per document (with `set_as`, each dict is updated and yielded).
 
 ```chatterlang
 INPUT FROM echo[data="Hello world"]

docs/guides/model2vec-embeddings.md

@@ -113,6 +113,9 @@ precache_model("minishlab/potion-base-8M")
 embedder = Model2VecEmbedder()
 vector = embedder.embed_one("Paragraph text.")
 batch = embedder.embed(["first", "second"])
+
+# ChatterLang / llmEmbed use Model2VecEmbeddingAdapter, which batches via the same encode path:
+# adapter.execute_batch(["first", "second"]) or adapter(["first", "second"])
 ```
 
 Pin to a specific HF commit by passing `revision="<commit-sha>"` to either

src/talkpipe/llm/embedding.py

@@ -1,8 +1,7 @@
 """Module for embedding text using different models"""
 
-from typing import Optional, Annotated
+from typing import Optional, Annotated, Iterator, Any, List
 import logging
-from talkpipe.llm.embedding_adapters import OllamaEmbedderAdapter
 from talkpipe.pipe.core import AbstractSegment
 from talkpipe.chatterlang.registry import register_segment
 from talkpipe.util.data_manipulation import extract_property, assign_property
@@ -30,7 +29,8 @@ def __init__(
             source: Annotated[Optional[str], "The source of the embedding model (e.g., 'ollama')"] = None,
             field: Annotated[Optional[str], "If provided, extract text from this field in the input items"] = None,
             set_as: Annotated[Optional[str], "If provided, append embeddings to input items under this field name"] = None,
-            fail_on_error: Annotated[bool, "Whether to raise an error on failure or to silently ignore it"] = True
+            fail_on_error: Annotated[bool, "Whether to raise an error on failure or to silently ignore it"] = True,
+            batch_size: Annotated[int, "Number of texts to embed per provider call when items are scalars"] = 1,
             ):
         """Initialize the embedding segment with the specified parameters.
         
@@ -44,10 +44,69 @@ def __init__(
         if source not in getEmbeddingSources():
             logger.error(f"Source '{source}' is not supported. Supported sources are: {getEmbeddingSources()}")
             raise ValueError(f"Source '{source}' is not supported. Supported sources are: {getEmbeddingSources()}")
+        if batch_size < 1:
+            raise ValueError("batch_size must be a positive integer")
         self.embedder = getEmbeddingAdapter(source)(model=model)
         self.field = field
         self.set_as = set_as
         self.fail_on_error = fail_on_error
+        self.batch_size = batch_size
+
+    def _text_from_item(self, item: Any) -> str:
+        if self.field is not None:
+            return str(extract_property(item, self.field))
+        return str(item)
+
+    def _yield_embedded(
+        self, items: List[Any], vectors: List[List[float]]
+    ) -> Iterator[Any]:
+        for item, ans in zip(items, vectors):
+            logger.debug(f"Received embedding: {ans}")
+            if self.set_as is not None:
+                logger.debug(f"Appending embedding to field {self.set_as}")
+                assign_property(item, self.set_as, ans)
+                yield item
+            else:
+                logger.debug("Yielding embedding directly")
+                yield ans
+
+    def _vectors_for_texts(self, texts: List[str]) -> List[List[float]]:
+        if not texts:
+            return []
+        if len(texts) == 1:
+            return [self.embedder.execute_one(texts[0])]
+        return self.embedder.execute_batch(texts)
+
+    def _embed_and_emit(self, items: List[Any], texts: List[str]) -> Iterator[Any]:
+        if not items or not texts:
+            return
+        logger.debug(f"Embedding batch of {len(texts)} texts")
+        try:
+            vectors = self._vectors_for_texts(texts)
+            yield from self._yield_embedded(items, vectors)
+        except Exception as e:
+            logger.error(f"Error during batch embedding: {e}")
+            if self.fail_on_error:
+                raise
+            if len(texts) == 1:
+                return
+            logger.warning(
+                "Batch embedding failed; falling back to per-item embedding"
+            )
+            for item, text in zip(items, texts):
+                try:
+                    ans = self.embedder.execute_one(text)
+                except Exception as item_error:
+                    logger.error(f"Error during embedding: {item_error}")
+                    continue
+                yield from self._yield_embedded([item], [ans])
+
+    def _embed_list_item(self, list_item: list) -> Iterator[Any]:
+        if not list_item:
+            return
+        items = list(list_item)
+        texts = [self._text_from_item(item) for item in items]
+        yield from self._embed_and_emit(items, texts)
 
     def transform(self, input_iter):
         """Transform input items by creating embeddings.
@@ -59,30 +118,32 @@ def transform(self, input_iter):
             If set_as is specified, yields the original items with embeddings added.
             Otherwise, yields the embeddings directly.
         """
+        buffer_items: List[Any] = []
+        buffer_texts: List[str] = []
+
+        def flush_buffer() -> Iterator[Any]:
+            if not buffer_items:
+                return
+            yield from self._embed_and_emit(buffer_items, buffer_texts)
+            buffer_items.clear()
+            buffer_texts.clear()
+
         for item in input_iter:
             logging.debug(f"Processing input item: {item}")
-            if self.field is not None:
-                text = extract_property(item, self.field)
-                logging.debug(f"Extracted text from field {self.field}: {text}")
-            else:
-                text = item
-                logging.debug(f"Using item as text: {text}")
+            if isinstance(item, list):
+                yield from flush_buffer()
+                yield from self._embed_list_item(item)
+                continue
 
-            logger.debug(f"Embedding text: {text}")
-            try:
-                ans = self.embedder.execute(str(text))
-            except Exception as e:
-                logger.error(f"Error during embedding: {e}")
-                if self.fail_on_error:
-                    raise e
-                else:
-                    continue
-            logger.debug(f"Received embedding: {ans}")
+            text = self._text_from_item(item)
+            logging.debug(f"Embedding text: {text}")
 
-            if self.set_as is not None:
-                logger.debug(f"Appending embedding to field {self.set_as}")
-                assign_property(item, self.set_as, ans)
-                yield item
+            if self.batch_size <= 1:
+                yield from self._embed_and_emit([item], [text])
             else:
-                logger.debug("Yielding embedding directly")
-                yield ans
+                buffer_items.append(item)
+                buffer_texts.append(text)
+                if len(buffer_items) >= self.batch_size:
+                    yield from flush_buffer()
+
+        yield from flush_buffer()

src/talkpipe/llm/embedding_adapters.py

@@ -1,15 +1,35 @@
-from typing import List
+from __future__ import annotations
+
+import warnings
+from typing import List, overload, Sequence, Union
+
 import numpy as np
+
 from talkpipe.util.config import get_config
 from talkpipe.util.constants import OLLAMA_SERVER_URL
 
+
+def _vector_to_list(vec) -> List[float]:
+    return np.asarray(vec, dtype=float).tolist()
+
+
+def _vectors_to_lists(arr) -> List[List[float]]:
+    a = np.asarray(arr, dtype=float)
+    if a.size == 0:
+        return []
+    if a.ndim == 1:
+        return [_vector_to_list(a)]
+    return [_vector_to_list(row) for row in a]
+
+
 class AbstractEmbeddingAdapter:
     """Abstract class for embedding text.
 
     This class represents an abstract adapter to embedding models.
     It defines the API and a common way to interact with different embedding models.  The
     specifics for embedding the text themselves are implemented in subclasses.
     """
+
     _model_name: str
     _source: str
 
@@ -35,11 +55,41 @@ def __str__(self):
     def __repr__(self):
         return self.__str__()
 
+    def execute_one(self, text: str) -> List[float]:
+        raise NotImplementedError("Subclasses must implement execute_one.")
+
+    def execute_batch(self, texts: Sequence[str]) -> List[List[float]]:
+        if not texts:
+            return []
+        return [self.execute_one(t) for t in texts]
+
     def execute(self, text: str) -> List[float]:
-        raise NotImplementedError("This method must be implemented in a subclass.")
+        """Embed a single string (deprecated).
+
+        .. deprecated::
+            Use :meth:`execute_one` or :meth:`execute_batch` instead.
+            ``execute`` will be removed in TalkPipe 1.0.
+        """
+        warnings.warn(
+            "EmbeddingAdapter.execute() is deprecated and will be removed in "
+            "TalkPipe 1.0. Use execute_one() or execute_batch() instead.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        return self.execute_one(text)
+
+    @overload
+    def __call__(self, text: str) -> List[float]: ...
+
+    @overload
+    def __call__(self, text: Sequence[str]) -> List[List[float]]: ...
 
-    def __call__(self, text: str) -> List[float]:
-        return self.execute(text)
+    def __call__(
+        self, text: Union[str, Sequence[str]]
+    ) -> Union[List[float], List[List[float]]]:
+        if isinstance(text, str):
+            return self.execute_one(text)
+        return self.execute_batch(list(text))
 
 
 class OllamaEmbedderAdapter(AbstractEmbeddingAdapter):
@@ -49,21 +99,24 @@ def __init__(self, model: str, server_url: str = None):
         super().__init__(model, "ollama")
         self._server_url = server_url
 
-    def execute(self, text: str) -> List[float]:
+    def _client(self):
         try:
             import ollama
         except ImportError:
             raise ImportError(
                 "Ollama is not installed. Please install it with: pip install talkpipe[ollama]"
             )
-
         server_url = self._server_url
         if not server_url:
             server_url = get_config().get(OLLAMA_SERVER_URL, None)
-        client = ollama.Client(server_url) if server_url else ollama
-        response = client.embed(
-            model=self.model_name,
-            input=text
-        )
-        result = response["embeddings"][0]
-        return np.array(result)
+        return ollama.Client(server_url) if server_url else ollama
+
+    def execute_batch(self, texts: Sequence[str]) -> List[List[float]]:
+        if not texts:
+            return []
+        client = self._client()
+        response = client.embed(model=self.model_name, input=list(texts))
+        return _vectors_to_lists(response["embeddings"])
+
+    def execute_one(self, text: str) -> List[float]:
+        return self.execute_batch([text])[0]

src/talkpipe/llm/embedding_adapters_model2vec.py

@@ -1,9 +1,9 @@
-from typing import List, Optional
+from typing import List, Optional, Sequence
 
 from talkpipe.util.config import get_config
 from talkpipe.util.constants import MODEL2VEC_CACHE_DIR, MODEL2VEC_REVISION
 
-from .embedding_adapters import AbstractEmbeddingAdapter
+from .embedding_adapters import AbstractEmbeddingAdapter, _vectors_to_lists
 from .model2vec_embeddings import DEFAULT_MODEL, Model2VecEmbedder
 
 
@@ -20,5 +20,10 @@ def __init__(self, model: Optional[str] = None):
             cache_folder=cfg.get(MODEL2VEC_CACHE_DIR),
         )
 
-    def execute(self, text: str) -> List[float]:
-        return self._embedder.embed_one(text)
+    def execute_batch(self, texts: Sequence[str]) -> List[List[float]]:
+        if not texts:
+            return []
+        return _vectors_to_lists(self._embedder.embed(list(texts)))
+
+    def execute_one(self, text: str) -> List[float]:
+        return self.execute_batch([text])[0]

src/talkpipe/llm/embedding_adapters_openai.py

@@ -1,4 +1,4 @@
-from typing import List
+from typing import List, Sequence
 
 from .embedding_adapters import AbstractEmbeddingAdapter
 
@@ -21,7 +21,16 @@ def __init__(self, model: str):
         openai = _require_openai()
         self.client = openai.OpenAI()
 
-    def execute(self, text: str) -> List[float]:
+    def execute_batch(self, texts: Sequence[str]) -> List[List[float]]:
+        if not texts:
+            return []
+        response = self.client.embeddings.create(
+            model=self.model_name,
+            input=list(texts),
+        )
+        return [list(d.embedding) for d in response.data]
+
+    def execute_one(self, text: str) -> List[float]:
         response = self.client.embeddings.create(
             model=self.model_name,
             input=text,

tests/talkpipe/llm/test_embedding.py

@@ -101,7 +101,7 @@ def test_fail_on_error_parameter():
     """Test that fail_on_error parameter works correctly to prevent duplicate execute() calls."""
     # Create a mock embedder that always fails
     mock_embedder = Mock()
-    mock_embedder.execute = Mock(side_effect=RuntimeError("Embedding failed"))
+    mock_embedder.execute_one = Mock(side_effect=RuntimeError("Embedding failed"))
 
     # Test with fail_on_error=True (should raise exception)
     embedder_true = LLMEmbed(model="test-model", source="ollama", fail_on_error=True)
@@ -111,7 +111,7 @@ def test_fail_on_error_parameter():
         list(embedder_true(["test input"]))
 
     # Verify execute was called only once (not twice due to duplicate line bug)
-    assert mock_embedder.execute.call_count == 1
+    assert mock_embedder.execute_one.call_count == 1
 
     # Reset mock
     mock_embedder.reset_mock()
@@ -126,5 +126,5 @@ def test_fail_on_error_parameter():
     # Result should be empty since the embedding failed and was skipped
     assert result == []
 
-    # Verify execute was called only once (not twice due to duplicate line bug)
-    assert mock_embedder.execute.call_count == 1
+    # Verify execute_one was called only once (not twice due to duplicate line bug)
+    assert mock_embedder.execute_one.call_count == 1