refactor: Add specific max token values to embedding models

This commit is contained in:
Igor Ilic 2025-01-23 18:18:45 +01:00
parent 6d5679f9d2
commit 7dea1d54d7
4 changed files with 4 additions and 5 deletions

View file

@ -31,12 +31,11 @@ class LiteLLMEmbeddingEngine(EmbeddingEngine):
api_key: str = None,
endpoint: str = None,
api_version: str = None,
max_tokens: int = float("inf"),
max_tokens: int = 512,
):
self.api_key = api_key
self.endpoint = endpoint
self.api_version = api_version
# TODO: Add or remove provider info
self.provider = provider
self.model = model
self.dimensions = dimensions

View file

@ -7,7 +7,7 @@ class GeminiTokenizer(TokenizerInterface):
def __init__(
self,
model: str,
max_tokens: int = float("inf"),
max_tokens: int = 3072,
):
self.model = model
self.max_tokens = max_tokens

View file

@ -9,7 +9,7 @@ class HuggingFaceTokenizer(TokenizerInterface):
def __init__(
self,
model: str,
max_tokens: int = float("inf"),
max_tokens: int = 512,
):
self.model = model
self.max_tokens = max_tokens

View file

@ -13,7 +13,7 @@ class TikTokenTokenizer(TokenizerInterface):
def __init__(
self,
model: str,
max_tokens: int = float("inf"),
max_tokens: int = 8191,
):
self.model = model
self.max_tokens = max_tokens