diff --git a/rag/llm/embedding_model.py b/rag/llm/embedding_model.py index df57f5d9f5d..1babf115280 100644 --- a/rag/llm/embedding_model.py +++ b/rag/llm/embedding_model.py @@ -217,6 +217,14 @@ def __init__(self, key, model_name="embedding-2", **kwargs): def encode(self, texts: list): arr = [] tks_num = 0 + MAX_LEN = -1 + if self.model_name.lower() == "embedding-2": + MAX_LEN = 512 + if self.model_name.lower() == "embedding-3": + MAX_LEN = 3072 + if MAX_LEN > 0: + texts = [truncate(t, MAX_LEN) for t in texts] + for txt in texts: res = self.client.embeddings.create(input=txt, model=self.model_name)