From a869bd46a60848a998607b90e6d64ba3588920be Mon Sep 17 00:00:00 2001 From: David Kristek Date: Fri, 26 Apr 2024 19:29:46 +0200 Subject: [PATCH] fix(llamaindex): avoid batching in embeddings --- src/genai/extensions/llama_index/embeddings.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/genai/extensions/llama_index/embeddings.py b/src/genai/extensions/llama_index/embeddings.py index 195829fd..3cb7dd85 100644 --- a/src/genai/extensions/llama_index/embeddings.py +++ b/src/genai/extensions/llama_index/embeddings.py @@ -1,6 +1,8 @@ import asyncio from typing import Optional +from pydantic import Field + from genai._types import ModelLike from genai.client import Client from genai.schema import TextEmbeddingParameters @@ -17,6 +19,9 @@ class IBMGenAILlamaIndexEmbedding(BaseEmbedding): model_id: str parameters: Optional[ModelLike[TextEmbeddingParameters]] = None execution_options: Optional[ModelLike[CreateExecutionOptions]] = None + # Batch size is set to 100000 to avoid batching in + # LlamaIndex as it is handled by the SDK itself + embed_batch_size: int = Field(default=10000, description="The batch size for embedding calls.", gt=0) @classmethod def class_name(cls) -> str: