Batch OpenAI embedder (#9381)

GitOrigin-RevId: 4dd920c04055365d0992802b600e48c2788535b1

Author: szymondudycz

integration_tests/xpack/test_embedders.py

@@ -15,10 +15,13 @@
     "model", [None, "text-embedding-ada-002", "text-embedding-3-small"]
 )
 @pytest.mark.parametrize("strategy", ["start", "end"])
-def test_openai_embedder(text: str, model: str, strategy: str):
+def test_openai_embedder(text: str, model: str | None, strategy: str):
+    table = pw.debug.table_from_rows(
+        schema=pw.schema_from_types(text=str), rows=[(text,)]
+    )
     if model is None:
         embedder = embedders.OpenAIEmbedder(
-            truncation_keep_strategy=strategy,
+            truncation_keep_strategy=strategy,  # type: ignore
             retry_strategy=pw.udfs.ExponentialBackoffRetryStrategy(),
         )
     else:
@@ -28,28 +31,102 @@ def test_openai_embedder(text: str, model: str, strategy: str):
             retry_strategy=pw.udfs.ExponentialBackoffRetryStrategy(),
         )
 
-    sync_embedder = _coerce_sync(embedder.func)
+    table = table.select(embedding=embedder(pw.this.text))
 
-    embedding = sync_embedder(text)
+    result = pw.debug.table_to_pandas(table).to_dict("records")
 
-    assert len(embedding) > 1500
+    assert len(result) == 1
+    assert isinstance(result[0]["embedding"][0], float)
+    assert len(result[0]["embedding"]) > 1500
 
 
 @pytest.mark.parametrize("model", ["text-embedding-ada-002", "text-embedding-3-small"])
 def test_openai_embedder_fails_no_truncation(model: str):
     truncation_keep_strategy = None
     embedder = embedders.OpenAIEmbedder(
-        model=model, truncation_keep_strategy=truncation_keep_strategy
+        model=model,
+        truncation_keep_strategy=truncation_keep_strategy,
+        retry_strategy=pw.udfs.ExponentialBackoffRetryStrategy(),
     )
 
     sync_embedder = _coerce_sync(embedder.func)
 
     with pytest.raises(Exception) as exc:
-        sync_embedder(LONG_TEXT)
+        sync_embedder([LONG_TEXT])
 
     assert "maximum context length" in str(exc)
 
 
+def test_openai_embedder_with_common_parameter():
+    table = pw.debug.table_from_rows(
+        schema=pw.schema_from_types(text=str), rows=[("aaa",), ("bbb",)]
+    )
+
+    embedder = embedders.OpenAIEmbedder(
+        model="text-embedding-3-small",
+        retry_strategy=pw.udfs.ExponentialBackoffRetryStrategy(),
+    )
+
+    table = table.select(embedding=embedder(pw.this.text, dimensions=700))
+
+    result = pw.debug.table_to_pandas(table).to_dict("records")
+
+    assert len(result) == 2
+    assert isinstance(result[0]["embedding"][0], float)
+    assert len(result[0]["embedding"]) == 700
+    assert isinstance(result[1]["embedding"][0], float)
+    assert len(result[1]["embedding"]) == 700
+
+
+def test_openai_embedder_with_different_parameter():
+    table = pw.debug.table_from_rows(
+        schema=pw.schema_from_types(text=str, dimensions=int),
+        rows=[("aaa", 300), ("bbb", 800)],
+    )
+
+    embedder = embedders.OpenAIEmbedder(
+        model="text-embedding-3-small",
+        retry_strategy=pw.udfs.ExponentialBackoffRetryStrategy(),
+    )
+
+    table = table.select(
+        text=pw.this.text,
+        embedding=embedder(pw.this.text, dimensions=pw.this.dimensions),
+    )
+
+    result = pw.debug.table_to_pandas(table).to_dict("records")
+
+    assert len(result) == 2
+    assert isinstance(result[0]["embedding"][0], float)
+    assert isinstance(result[1]["embedding"][0], float)
+    if result[0]["text"] == "aaa":
+        assert len(result[0]["embedding"]) == 300
+    else:
+        assert len(result[1]["embedding"]) == 300
+    if result[0]["text"] == "bbb":
+        assert len(result[0]["embedding"]) == 800
+    else:
+        assert len(result[1]["embedding"]) == 800
+
+
+def test_openai_embedder_input_as_kwarg():
+    table = pw.debug.table_from_rows(
+        schema=pw.schema_from_types(text=str), rows=[("foo",)]
+    )
+    embedder = embedders.OpenAIEmbedder(
+        model="text-embedding-3-small",
+        retry_strategy=pw.udfs.ExponentialBackoffRetryStrategy(),
+    )
+
+    table = table.select(embedding=embedder(input=pw.this.text))
+
+    result = pw.debug.table_to_pandas(table).to_dict("records")
+
+    assert len(result) == 1
+    assert isinstance(result[0]["embedding"][0], float)
+    assert len(result[0]["embedding"]) > 1500
+
+
 def test_sentence_transformer_embedder():
     table = pw.debug.table_from_rows(
         schema=pw.schema_from_types(text=str), rows=[("aaa",), ("bbb",)]

python/pathway/xpacks/llm/embedders.py

@@ -4,7 +4,7 @@
 """
 import asyncio
 import logging
-from typing import Literal
+from typing import Any, Literal
 
 import numpy as np
 
@@ -85,6 +85,23 @@ def __call__(
         return super().__call__(input, *args, **kwargs)
 
 
+def _split_batched_kwargs(
+    kwargs: dict[str, list[Any]]
+) -> tuple[dict[str, Any], dict[str, list[Any]]]:
+    constant_kwargs = {}
+    per_row_kwargs = {}
+
+    if kwargs:
+        for key, values in kwargs.items():
+            v = values[0]
+            if all(value == v for value in values):
+                constant_kwargs[key] = v
+            else:
+                per_row_kwargs[key] = values
+
+    return constant_kwargs, per_row_kwargs
+
+
 class OpenAIEmbedder(BaseEmbedder):
     """Pathway wrapper for OpenAI Embedding services.
 
@@ -114,6 +131,8 @@ class OpenAIEmbedder(BaseEmbedder):
             Can be ``"start"``, ``"end"`` or ``None``. ``"start"`` will keep the first part of the text
             and remove the rest. ``"end"`` will keep the last part of the text.
             If `None`, no truncation will be applied to any of the documents, this may cause API exceptions.
+        batch_size: maximum size of a single batch to be sent to the embedder. Bigger
+            batches may reduce the time needed for embedding.
         encoding_format: The format to return the embeddings in. Can be either `float` or
             `base64 <https://pypi.org/project/pybase64/>`_.
         user: A unique identifier representing your end-user, which can help OpenAI to monitor
@@ -160,6 +179,7 @@ def __init__(
         cache_strategy: udfs.CacheStrategy | None = None,
         model: str | None = "text-embedding-3-small",
         truncation_keep_strategy: Literal["start", "end"] | None = "start",
+        batch_size: int = 128,
         **openai_kwargs,
     ):
         with optional_imports("xpack-llm"):
@@ -168,42 +188,82 @@ def __init__(
         _monkeypatch_openai_async()
         executor = udfs.async_executor(capacity=capacity, retry_strategy=retry_strategy)
         super().__init__(
-            executor=executor,
-            cache_strategy=cache_strategy,
+            executor=executor, cache_strategy=cache_strategy, max_batch_size=batch_size
         )
         self.truncation_keep_strategy = truncation_keep_strategy
         self.kwargs = dict(openai_kwargs)
-        api_key = self.kwargs.pop("api_key", None)
-        self.client = openai.AsyncOpenAI(api_key=api_key, max_retries=0)
+        self.api_key = self.kwargs.pop("api_key", None)
+        self.client: openai.AsyncOpenAI | None = None
+
+        # Initialization of OpenAI for the purpose of checking if api_key was provided
+        # Actual initialization of the client is delayed to __wrapped__ to avoid issues
+        # with the event loop.
+        _ = openai.AsyncOpenAI(api_key=self.api_key, max_retries=0)
         if model is not None:
             self.kwargs["model"] = model
 
-    async def __wrapped__(self, input, **kwargs) -> np.ndarray:
+    async def __wrapped__(self, inputs: list[str], **kwargs) -> list[np.ndarray]:
         """Embed the documents
 
         Args:
-            input: mandatory, the string to embed.
+            inputs: mandatory, the strings to embed.
             **kwargs: optional parameters, if unset defaults from the constructor
               will be taken.
-        """
-        input = input or "."
+        #"""
+        import openai
+
+        if self.client is None:
+            self.client = openai.AsyncOpenAI(api_key=self.api_key, max_retries=0)
 
-        kwargs = {**self.kwargs, **kwargs}
         kwargs = _extract_value_inside_dict(kwargs)
 
-        if kwargs.get("model") is None:
+        if kwargs.get("model") is None and self.kwargs.get("model") is None:
             raise ValueError(
                 "`model` parameter is missing in `OpenAIEmbedder`. "
                 "Please provide the model name either in the constructor or in the function call."
             )
 
+        constant_kwargs, per_row_kwargs = _split_batched_kwargs(kwargs)
+        constant_kwargs = {**self.kwargs, **constant_kwargs}
+
         if self.truncation_keep_strategy:
-            input = self.truncate_context(
-                kwargs["model"], input, self.truncation_keep_strategy
-            )
+            if "model" in per_row_kwargs:
+                inputs = [
+                    self.truncate_context(model, input, self.truncation_keep_strategy)
+                    for (model, input) in zip(per_row_kwargs["model"], inputs)
+                ]
+            else:
+                inputs = [
+                    self.truncate_context(
+                        constant_kwargs["model"], input, self.truncation_keep_strategy
+                    )
+                    for input in inputs
+                ]
+
+        # if kwargs are not the same for every input we cannot batch them
+        if per_row_kwargs:
+
+            async def embed_single(input, kwargs) -> np.ndarray:
+                kwargs = {**constant_kwargs, **kwargs}
+                ret = await self.client.embeddings.create(input=[input], **kwargs)  # type: ignore
+                return np.array(ret.data[0].embedding)
+
+            list_of_per_row_kwargs = [
+                dict(zip(per_row_kwargs, values))
+                for values in zip(*per_row_kwargs.values())
+            ]
+            async with asyncio.TaskGroup() as tg:
+                tasks = [
+                    tg.create_task(embed_single(input, kwargs))
+                    for input, kwargs in zip(inputs, list_of_per_row_kwargs)
+                ]
 
-        ret = await self.client.embeddings.create(input=[input], **kwargs)
-        return np.array(ret.data[0].embedding)
+            result_list = [task.result() for task in tasks]
+            return result_list
+
+        else:
+            ret = await self.client.embeddings.create(input=inputs, **constant_kwargs)
+            return [np.array(datum.embedding) for datum in ret.data]
 
     @staticmethod
     def truncate_context(
@@ -250,6 +310,18 @@ def truncate_context(
 
         return tokenizer.decode(tokens)
 
+    def get_embedding_dimension(self, **kwargs):
+        """Computes number of embedder's dimensions by asking the embedder to embed ``"."``.
+
+        Args:
+            **kwargs: parameters of the embedder, if unset defaults from the constructor
+              will be taken.
+        """
+        kwargs_as_list = {k: [v] for k, v in kwargs.items()}
+        n_dimensions = len(_coerce_sync(self.__wrapped__)(["."], **kwargs_as_list)[0])
+        self.client = None
+        return n_dimensions
+
 
 class LiteLLMEmbedder(BaseEmbedder):
     """Pathway wrapper for `litellm.embedding`.
@@ -406,16 +478,7 @@ def __wrapped__(self, input: list[str], **kwargs) -> list[np.ndarray]:
         """  # noqa: E501
 
         kwargs = _extract_value_inside_dict(kwargs)
-        constant_kwargs = {}
-        per_row_kwargs = {}
-
-        if kwargs:
-            for key, values in kwargs.items():
-                v = values[0]
-                if all(value == v for value in values):
-                    constant_kwargs[key] = v
-                else:
-                    per_row_kwargs[key] = values
+        constant_kwargs, per_row_kwargs = _split_batched_kwargs(kwargs)
 
         # if kwargs are not the same for every input we cannot batch them
         if per_row_kwargs: