refactor: Add specific max token values to embedding models

topoteretes · Jan 23, 2025 · 7dea1d5 · 7dea1d5
1 parent 6d5679f
commit 7dea1d5
Show file tree

Hide file tree

Showing 4 changed files with 4 additions and 5 deletions.
diff --git a/cognee/infrastructure/databases/vector/embeddings/LiteLLMEmbeddingEngine.py b/cognee/infrastructure/databases/vector/embeddings/LiteLLMEmbeddingEngine.py
@@ -31,12 +31,11 @@ def __init__(
         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

diff --git a/cognee/infrastructure/llm/tokenizer/Gemini/adapter.py b/cognee/infrastructure/llm/tokenizer/Gemini/adapter.py
@@ -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

diff --git a/cognee/infrastructure/llm/tokenizer/HuggingFace/adapter.py b/cognee/infrastructure/llm/tokenizer/HuggingFace/adapter.py
@@ -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

diff --git a/cognee/infrastructure/llm/tokenizer/TikToken/adapter.py b/cognee/infrastructure/llm/tokenizer/TikToken/adapter.py
@@ -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