diff --git a/mteb/models/openai_models.py b/mteb/models/openai_models.py index adf96fbe4e..aecacf549a 100644 --- a/mteb/models/openai_models.py +++ b/mteb/models/openai_models.py @@ -63,7 +63,7 @@ def encode(self, sentences: list[str], **kwargs: Any) -> np.ndarray: else: trimmed_sentences.append(sentence) - max_batch_size = 2048 + max_batch_size = kwargs.get("batch_size", 2048) sublists = [ trimmed_sentences[i : i + max_batch_size] for i in range(0, len(trimmed_sentences), max_batch_size)