feat: cache tokenizers (#3558)

We now load all tokenizers at the import of the llm_endpoint module, and
any initialisation of the LLMEndpoint class will use the cached
tokenizers

Closes ENT-402

Author: jacopo-chevallard

core/quivr_core/llm/llm_endpoint.py

@@ -17,17 +17,14 @@
 logger = logging.getLogger("quivr_core")
 
 
-class LLMEndpoint:
-    _cache: dict[int, "LLMEndpoint"] = {}
+class LLMTokenizer:
+    _cache: dict[int, "LLMTokenizer"] = {}
 
-    def __init__(self, llm_config: LLMEndpointConfig, llm: BaseChatModel):
-        self._config = llm_config
-        self._llm = llm
-        self._supports_func_calling = model_supports_function_calling(
-            self._config.model
-        )
+    def __init__(self, tokenizer_hub: str | None, fallback_tokenizer: str):
+        self.tokenizer_hub = tokenizer_hub
+        self.fallback_tokenizer = fallback_tokenizer
 
-        if llm_config.tokenizer_hub:
+        if self.tokenizer_hub:
             # To prevent the warning
             # huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
             os.environ["TOKENIZERS_PARALLELISM"] = (
@@ -36,34 +33,81 @@ def __init__(self, llm_config: LLMEndpointConfig, llm: BaseChatModel):
                 else os.environ["TOKENIZERS_PARALLELISM"]
             )
             try:
-                from transformers import AutoTokenizer
+                if "text-embedding-ada-002" in self.tokenizer_hub:
+                    from transformers import GPT2TokenizerFast
 
-                self.tokenizer = AutoTokenizer.from_pretrained(llm_config.tokenizer_hub)
+                    self.tokenizer = GPT2TokenizerFast.from_pretrained(
+                        self.tokenizer_hub
+                    )
+                else:
+                    from transformers import AutoTokenizer
+
+                    self.tokenizer = AutoTokenizer.from_pretrained(self.tokenizer_hub)
             except OSError:  # if we don't manage to connect to huggingface and/or no cached models are present
                 logger.warning(
-                    f"Cannot acces the configured tokenizer from {llm_config.tokenizer_hub}, using the default tokenizer {llm_config.fallback_tokenizer}"
+                    f"Cannot acces the configured tokenizer from {self.tokenizer_hub}, using the default tokenizer {self.fallback_tokenizer}"
                 )
-                self.tokenizer = tiktoken.get_encoding(llm_config.fallback_tokenizer)
+                self.tokenizer = tiktoken.get_encoding(self.fallback_tokenizer)
         else:
-            self.tokenizer = tiktoken.get_encoding(llm_config.fallback_tokenizer)
+            self.tokenizer = tiktoken.get_encoding(self.fallback_tokenizer)
+
+    @classmethod
+    def load(cls, tokenizer_hub: str, fallback_tokenizer: str):
+        cache_key = hash(str(tokenizer_hub))
+        if cache_key in cls._cache:
+            return cls._cache[cache_key]
+        instance = cls(tokenizer_hub, fallback_tokenizer)
+        cls._cache[cache_key] = instance
+        return instance
+
+    @classmethod
+    def preload_tokenizers(cls):
+        """Preload all available tokenizers from the models configuration into cache."""
+        from quivr_core.rag.entities.config import LLMModelConfig
+
+        unique_tokenizer_hubs = set()
+
+        # Collect all unique tokenizer hubs
+        for supplier_models in LLMModelConfig._model_defaults.values():
+            for config in supplier_models.values():
+                if config.tokenizer_hub:
+                    unique_tokenizer_hubs.add(config.tokenizer_hub)
+
+        # Load each unique tokenizer
+        for hub in unique_tokenizer_hubs:
+            try:
+                cls.load(hub, LLMEndpointConfig._FALLBACK_TOKENIZER)
+                logger.info(f"Successfully preloaded tokenizer: {hub}")
+            except Exception as e:
+                logger.warning(f"Failed to preload tokenizer {hub}: {str(e)}")
+
+
+# Preload tokenizers when module is imported
+LLMTokenizer.preload_tokenizers()
+
+
+class LLMEndpoint:
+    def __init__(self, llm_config: LLMEndpointConfig, llm: BaseChatModel):
+        self._config = llm_config
+        self._llm = llm
+        self._supports_func_calling = model_supports_function_calling(
+            self._config.model
+        )
+
+        self.llm_tokenizer = LLMTokenizer.load(
+            llm_config.tokenizer_hub, llm_config.fallback_tokenizer
+        )
 
     def count_tokens(self, text: str) -> int:
         # Tokenize the input text and return the token count
-        encoding = self.tokenizer.encode(text)
+        encoding = self.llm_tokenizer.tokenizer.encode(text)
         return len(encoding)
 
     def get_config(self):
         return self._config
 
     @classmethod
     def from_config(cls, config: LLMEndpointConfig = LLMEndpointConfig()):
-        # Create a cache key from the config
-        cache_key = hash(str(config.model_dump()))
-
-        # Return cached instance if it exists
-        if cache_key in cls._cache:
-            return cls._cache[cache_key]
-
         _llm: Union[AzureChatOpenAI, ChatOpenAI, ChatAnthropic, ChatMistralAI]
         try:
             if config.supplier == DefaultModelSuppliers.AZURE:
@@ -122,7 +166,6 @@ def from_config(cls, config: LLMEndpointConfig = LLMEndpointConfig()):
                     temperature=config.temperature,
                 )
             instance = cls(llm=_llm, llm_config=config)
-            cls._cache[cache_key] = instance
             return instance
 
         except ImportError as e:

core/quivr_core/rag/entities/config.py

@@ -86,73 +86,73 @@ class LLMModelConfig:
             "gpt-4o": LLMConfig(
                 max_context_tokens=128000,
                 max_output_tokens=16384,
-                tokenizer_hub="Xenova/gpt-4o",
+                tokenizer_hub="Quivr/gpt-4o",
             ),
             "gpt-4o-mini": LLMConfig(
                 max_context_tokens=128000,
                 max_output_tokens=16384,
-                tokenizer_hub="Xenova/gpt-4o",
+                tokenizer_hub="Quivr/gpt-4o",
             ),
             "gpt-4-turbo": LLMConfig(
                 max_context_tokens=128000,
                 max_output_tokens=4096,
-                tokenizer_hub="Xenova/gpt-4",
+                tokenizer_hub="Quivr/gpt-4",
             ),
             "gpt-4": LLMConfig(
                 max_context_tokens=8192,
                 max_output_tokens=8192,
-                tokenizer_hub="Xenova/gpt-4",
+                tokenizer_hub="Quivr/gpt-4",
             ),
             "gpt-3.5-turbo": LLMConfig(
                 max_context_tokens=16385,
                 max_output_tokens=4096,
-                tokenizer_hub="Xenova/gpt-3.5-turbo",
+                tokenizer_hub="Quivr/gpt-3.5-turbo",
             ),
             "text-embedding-3-large": LLMConfig(
-                max_context_tokens=8191, tokenizer_hub="Xenova/text-embedding-ada-002"
+                max_context_tokens=8191, tokenizer_hub="Quivr/text-embedding-ada-002"
             ),
             "text-embedding-3-small": LLMConfig(
-                max_context_tokens=8191, tokenizer_hub="Xenova/text-embedding-ada-002"
+                max_context_tokens=8191, tokenizer_hub="Quivr/text-embedding-ada-002"
             ),
             "text-embedding-ada-002": LLMConfig(
-                max_context_tokens=8191, tokenizer_hub="Xenova/text-embedding-ada-002"
+                max_context_tokens=8191, tokenizer_hub="Quivr/text-embedding-ada-002"
             ),
         },
         DefaultModelSuppliers.ANTHROPIC: {
             "claude-3-5-sonnet": LLMConfig(
                 max_context_tokens=200000,
                 max_output_tokens=8192,
-                tokenizer_hub="Xenova/claude-tokenizer",
+                tokenizer_hub="Quivr/claude-tokenizer",
             ),
             "claude-3-opus": LLMConfig(
                 max_context_tokens=200000,
                 max_output_tokens=4096,
-                tokenizer_hub="Xenova/claude-tokenizer",
+                tokenizer_hub="Quivr/claude-tokenizer",
             ),
             "claude-3-sonnet": LLMConfig(
                 max_context_tokens=200000,
                 max_output_tokens=4096,
-                tokenizer_hub="Xenova/claude-tokenizer",
+                tokenizer_hub="Quivr/claude-tokenizer",
             ),
             "claude-3-haiku": LLMConfig(
                 max_context_tokens=200000,
                 max_output_tokens=4096,
-                tokenizer_hub="Xenova/claude-tokenizer",
+                tokenizer_hub="Quivr/claude-tokenizer",
             ),
             "claude-2-1": LLMConfig(
                 max_context_tokens=200000,
                 max_output_tokens=4096,
-                tokenizer_hub="Xenova/claude-tokenizer",
+                tokenizer_hub="Quivr/claude-tokenizer",
             ),
             "claude-2-0": LLMConfig(
                 max_context_tokens=100000,
                 max_output_tokens=4096,
-                tokenizer_hub="Xenova/claude-tokenizer",
+                tokenizer_hub="Quivr/claude-tokenizer",
             ),
             "claude-instant-1-2": LLMConfig(
                 max_context_tokens=100000,
                 max_output_tokens=4096,
-                tokenizer_hub="Xenova/claude-tokenizer",
+                tokenizer_hub="Quivr/claude-tokenizer",
             ),
         },
         # Unclear for LLAMA models...
@@ -161,53 +161,53 @@ class LLMModelConfig:
             "llama-3.1": LLMConfig(
                 max_context_tokens=128000,
                 max_output_tokens=4096,
-                tokenizer_hub="Xenova/Meta-Llama-3.1-Tokenizer",
+                tokenizer_hub="Quivr/Meta-Llama-3.1-Tokenizer",
             ),
             "llama-3": LLMConfig(
                 max_context_tokens=8192,
                 max_output_tokens=2048,
-                tokenizer_hub="Xenova/llama3-tokenizer-new",
+                tokenizer_hub="Quivr/llama3-tokenizer-new",
             ),
             "code-llama": LLMConfig(
-                max_context_tokens=16384, tokenizer_hub="Xenova/llama-code-tokenizer"
+                max_context_tokens=16384, tokenizer_hub="Quivr/llama-code-tokenizer"
             ),
         },
         DefaultModelSuppliers.GROQ: {
             "llama-3.3-70b": LLMConfig(
                 max_context_tokens=128000,
                 max_output_tokens=32768,
-                tokenizer_hub="Xenova/Meta-Llama-3.1-Tokenizer",
+                tokenizer_hub="Quivr/Meta-Llama-3.1-Tokenizer",
             ),
             "llama-3.1-70b": LLMConfig(
                 max_context_tokens=128000,
                 max_output_tokens=32768,
-                tokenizer_hub="Xenova/Meta-Llama-3.1-Tokenizer",
+                tokenizer_hub="Quivr/Meta-Llama-3.1-Tokenizer",
             ),
             "llama-3": LLMConfig(
-                max_context_tokens=8192, tokenizer_hub="Xenova/llama3-tokenizer-new"
+                max_context_tokens=8192, tokenizer_hub="Quivr/llama3-tokenizer-new"
             ),
             "code-llama": LLMConfig(
-                max_context_tokens=16384, tokenizer_hub="Xenova/llama-code-tokenizer"
+                max_context_tokens=16384, tokenizer_hub="Quivr/llama-code-tokenizer"
             ),
         },
         DefaultModelSuppliers.MISTRAL: {
             "mistral-large": LLMConfig(
                 max_context_tokens=128000,
                 max_output_tokens=4096,
-                tokenizer_hub="Xenova/mistral-tokenizer-v3",
+                tokenizer_hub="Quivr/mistral-tokenizer-v3",
             ),
             "mistral-small": LLMConfig(
                 max_context_tokens=128000,
                 max_output_tokens=4096,
-                tokenizer_hub="Xenova/mistral-tokenizer-v3",
+                tokenizer_hub="Quivr/mistral-tokenizer-v3",
             ),
             "mistral-nemo": LLMConfig(
                 max_context_tokens=128000,
                 max_output_tokens=4096,
-                tokenizer_hub="Xenova/Mistral-Nemo-Instruct-Tokenizer",
+                tokenizer_hub="Quivr/Mistral-Nemo-Instruct-Tokenizer",
             ),
             "codestral": LLMConfig(
-                max_context_tokens=32000, tokenizer_hub="Xenova/mistral-tokenizer-v3"
+                max_context_tokens=32000, tokenizer_hub="Quivr/mistral-tokenizer-v3"
             ),
         },
     }
@@ -247,9 +247,9 @@ class LLMEndpointConfig(QuivrBaseConfig):
     llm_base_url: str | None = None
     env_variable_name: str | None = None
     llm_api_key: str | None = None
-    max_context_tokens: int = 10000
-    max_output_tokens: int = 4000
-    temperature: float = 0.7
+    max_context_tokens: int = 20000
+    max_output_tokens: int = 4096
+    temperature: float = 0.3
     streaming: bool = True
     prompt: CustomPromptsModel | None = None