Revert OpenAI batch embedder (#9361)

GitOrigin-RevId: 41aaaaebba111607c0e9c090422b9c4200540d6f

Author: voodoo11

integration_tests/xpack/test_embedders.py

@@ -1,7 +1,6 @@
 import pytest
 
 import pathway as pw
-from pathway.internals.udfs.retries import ExponentialBackoffRetryStrategy
 from pathway.internals.udfs.utils import _coerce_sync
 from pathway.xpacks.llm import embedders
 
@@ -17,9 +16,6 @@
 )
 @pytest.mark.parametrize("strategy", ["start", "end"])
 def test_openai_embedder(text: str, model: str, 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,
@@ -32,102 +28,28 @@ def test_openai_embedder(text: str, model: str, strategy: str):
             retry_strategy=pw.udfs.ExponentialBackoffRetryStrategy(),
         )
 
-    table = table.select(embedding=embedder(pw.this.text))
+    sync_embedder = _coerce_sync(embedder.func)
 
-    result = pw.debug.table_to_pandas(table).to_dict("records")
+    embedding = sync_embedder(text)
 
-    assert len(result) == 1
-    assert isinstance(result[0]["embedding"][0], float)
-    assert len(result[0]["embedding"]) > 1500
+    assert len(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,
-        retry_strategy=ExponentialBackoffRetryStrategy(),
+        model=model, truncation_keep_strategy=truncation_keep_strategy
     )
 
     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=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=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 Any, Literal
+from typing import Literal
 
 import numpy as np
 
@@ -85,23 +85,6 @@ 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.
 
@@ -130,8 +113,6 @@ 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
@@ -176,7 +157,6 @@ 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"):
@@ -185,7 +165,8 @@ def __init__(
         _monkeypatch_openai_async()
         executor = udfs.async_executor(capacity=capacity, retry_strategy=retry_strategy)
         super().__init__(
-            executor=executor, cache_strategy=cache_strategy, max_batch_size=batch_size
+            executor=executor,
+            cache_strategy=cache_strategy,
         )
         self.truncation_keep_strategy = truncation_keep_strategy
         self.kwargs = dict(openai_kwargs)
@@ -194,64 +175,32 @@ def __init__(
         if model is not None:
             self.kwargs["model"] = model
 
-    async def __wrapped__(self, inputs: list[str], **kwargs) -> list[np.ndarray]:
+    async def __wrapped__(self, input, **kwargs) -> np.ndarray:
         """Embed the documents
 
         Args:
-            inputs: mandatory, the strings to embed.
+            input: mandatory, the string to embed.
             **kwargs: optional parameters, if unset defaults from the constructor
               will be taken.
         """
+        input = input or "."
 
+        kwargs = {**self.kwargs, **kwargs}
         kwargs = _extract_value_inside_dict(kwargs)
 
-        if kwargs.get("model") is None and self.kwargs.get("model") is None:
+        if 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:
-            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)
-                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)
-                ]
-
-            result_list = [task.result() for task in tasks]
-            return result_list
+            input = self.truncate_context(
+                kwargs["model"], input, self.truncation_keep_strategy
+            )
 
-        else:
-            ret = await self.client.embeddings.create(input=inputs, **constant_kwargs)
-            return [np.array(datum.embedding) for datum in ret.data]
+        ret = await self.client.embeddings.create(input=[input], **kwargs)
+        return np.array(ret.data[0].embedding)
 
     @staticmethod
     def truncate_context(
@@ -298,25 +247,6 @@ def truncate_context(
 
         return tokenizer.decode(tokens)
 
-    @staticmethod
-    def _count_tokens(text: str, model: str) -> int:
-        with optional_imports("xpack-llm"):
-            import tiktoken
-
-        tokenizer = tiktoken.encoding_for_model(model)
-        tokens = tokenizer.encode(text)
-        return len(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()}
-        return len(_coerce_sync(self.__wrapped__)(["."], **kwargs_as_list)[0])
-
 
 class LiteLLMEmbedder(BaseEmbedder):
     """Pathway wrapper for `litellm.embedding`.
@@ -470,7 +400,16 @@ def __wrapped__(self, input: list[str], **kwargs) -> list[np.ndarray]:
         """  # noqa: E501
 
         kwargs = _extract_value_inside_dict(kwargs)
-        constant_kwargs, per_row_kwargs = _split_batched_kwargs(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
 
         # if kwargs are not the same for every input we cannot batch them
         if per_row_kwargs: