feat(settings): Configurable context_window and tokenizer (#1437)

private_gpt/components/llm/llm_component.py

@@ -1,11 +1,13 @@
 import logging
 
 from injector import inject, singleton
+from llama_index import set_global_tokenizer
 from llama_index.llms import MockLLM
 from llama_index.llms.base import LLM
+from transformers import AutoTokenizer  # type: ignore
 
 from private_gpt.components.llm.prompt_helper import get_prompt_style
-from private_gpt.paths import models_path
+from private_gpt.paths import models_cache_path, models_path
 from private_gpt.settings.settings import Settings
 
 logger = logging.getLogger(__name__)
@@ -18,6 +20,14 @@ class LLMComponent:
     @inject
     def __init__(self, settings: Settings) -> None:
         llm_mode = settings.llm.mode
+        if settings.llm.tokenizer:
+            set_global_tokenizer(
+                AutoTokenizer.from_pretrained(
+                    pretrained_model_name_or_path=settings.llm.tokenizer,
+                    cache_dir=str(models_cache_path),
+                )
+            )
+
         logger.info("Initializing the LLM in mode=%s", llm_mode)
         match settings.llm.mode:
             case "local":
@@ -29,9 +39,7 @@ class LLMComponent:
                     model_path=str(models_path / settings.local.llm_hf_model_file),
                     temperature=0.1,
                     max_new_tokens=settings.llm.max_new_tokens,
-                    # llama2 has a context window of 4096 tokens,
-                    # but we set it lower to allow for some wiggle room
-                    context_window=3900,
+                    context_window=settings.llm.context_window,
                     generate_kwargs={},
                     # All to GPU
                     model_kwargs={"n_gpu_layers": -1},
@@ -46,6 +54,8 @@ class LLMComponent:
 
                 self.llm = SagemakerLLM(
                     endpoint_name=settings.sagemaker.llm_endpoint_name,
+                    max_new_tokens=settings.llm.max_new_tokens,
+                    context_window=settings.llm.context_window,
                 )
             case "openai":
                 from llama_index.llms import OpenAI