Skip to content

Commit 35e9ff2

Browse files
authored
fix: adjust max token size for openai ADA-v2 embeddings (#3793)
* Adjust max token size for openai ADA-v2 embeddings * Added requested changes and corrected old seq len Apparently the limit for the older models is 2046 and not 2048, I included this change directly. See (https://beta.openai.com/docs/guides/embeddings/what-are-embeddings) to check.
1 parent a2c160e commit 35e9ff2

File tree

1 file changed

+4
-4
lines changed

1 file changed

+4
-4
lines changed

haystack/nodes/retriever/_embedding_encoder.py

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -391,30 +391,30 @@ def save(self, save_dir: Union[Path, str]):
391391
class _OpenAIEmbeddingEncoder(_BaseEmbeddingEncoder):
392392
def __init__(self, retriever: "EmbeddingRetriever"):
393393
# See https://beta.openai.com/docs/guides/embeddings for more details
394-
# OpenAI has a max seq length of 2048 tokens and unknown max batch size
395-
self.max_seq_len = min(2048, retriever.max_seq_len)
396394
self.url = "https://api.openai.com/v1/embeddings"
397395
self.api_key = retriever.api_key
398396
self.batch_size = min(64, retriever.batch_size)
399397
self.progress_bar = retriever.progress_bar
400398
model_class: str = next(
401399
(m for m in ["ada", "babbage", "davinci", "curie"] if m in retriever.embedding_model), "babbage"
402400
)
403-
self._setup_encoding_models(model_class, retriever.embedding_model)
401+
self._setup_encoding_models(model_class, retriever.embedding_model, retriever.max_seq_len)
404402

405403
self.tokenizer = AutoTokenizer.from_pretrained("gpt2")
406404

407-
def _setup_encoding_models(self, model_class: str, model_name: str):
405+
def _setup_encoding_models(self, model_class: str, model_name: str, max_seq_len: int):
408406
"""
409407
Setup the encoding models for the retriever.
410408
"""
411409
# new generation of embedding models (December 2022), we need to specify the full name
412410
if "text-embedding" in model_name:
413411
self.query_encoder_model = model_name
414412
self.doc_encoder_model = model_name
413+
self.max_seq_len = min(8191, max_seq_len)
415414
else:
416415
self.query_encoder_model = f"text-search-{model_class}-query-001"
417416
self.doc_encoder_model = f"text-search-{model_class}-doc-001"
417+
self.max_seq_len = min(2046, max_seq_len)
418418

419419
def _ensure_text_limit(self, text: str) -> str:
420420
"""

0 commit comments

Comments
 (0)