voyageai: new models (#26907)

Co-authored-by: fzowl <zoltan@voyageai.com>
Co-authored-by: fzowl <160063452+fzowl@users.noreply.github.com>
This commit is contained in:
Erick Friis
2024-09-26 10:07:10 -07:00
committed by GitHub
parent 2a0d9d05fb
commit 8bc12df2eb
4 changed files with 32 additions and 3 deletions

View File

@@ -16,6 +16,11 @@ from typing_extensions import Self
logger = logging.getLogger(__name__)
DEFAULT_VOYAGE_2_BATCH_SIZE = 72
DEFAULT_VOYAGE_3_LITE_BATCH_SIZE = 30
DEFAULT_VOYAGE_3_BATCH_SIZE = 10
DEFAULT_BATCH_SIZE = 7
class VoyageAIEmbeddings(BaseModel, Embeddings):
"""VoyageAIEmbeddings embedding model.
@@ -55,7 +60,19 @@ class VoyageAIEmbeddings(BaseModel, Embeddings):
model = values.get("model")
batch_size = values.get("batch_size")
if batch_size is None:
values["batch_size"] = 72 if model in ["voyage-2", "voyage-02"] else 7
values["batch_size"] = (
DEFAULT_VOYAGE_2_BATCH_SIZE
if model in ["voyage-2", "voyage-02"]
else (
DEFAULT_VOYAGE_3_LITE_BATCH_SIZE
if model == "voyage-3-lite"
else (
DEFAULT_VOYAGE_3_BATCH_SIZE
if model == "voyage-3"
else DEFAULT_BATCH_SIZE
)
)
)
return values
@model_validator(mode="after")