From 94fd681efab37285467674dd6e082c347fe0795a Mon Sep 17 00:00:00 2001 From: Nick Sullivan Date: Fri, 24 Nov 2023 13:46:11 -0800 Subject: [PATCH] fixed bug with token size --- aicodebot/lm.py | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/aicodebot/lm.py b/aicodebot/lm.py index 5120b46..c0a16cf 100644 --- a/aicodebot/lm.py +++ b/aicodebot/lm.py @@ -119,6 +119,7 @@ def get_ollama_model( return Ollama( model=model_name, # NOTE: num_ctx is similar, but not exactly the same as response_token_size + # num_ctx is the number of tokens in the context, whereas response_token_size is the number of tokens num_ctx=response_token_size, temperature=temperature, callbacks=callbacks, @@ -205,7 +206,15 @@ def get_model_token_limit(self, model_name): def get_token_size(self, text): """Get the number of tokens in a string using the tiktoken library.""" - encoding = tiktoken.encoding_for_model(self.tiktoken_model_name) + + try: + encoding = tiktoken.encoding_for_model(self.tiktoken_model_name) + except KeyError: + # NOTE: This sets the token size to default value for Ollama. + # TikToken doesn't seem to support Ollama-based models. + # Local models will not work as intended without dynamically setting the token size. + # TODO: Dynamically set token size for Ollama-based models. + return 2048 tokens = encoding.encode(text) return len(tokens) @@ -301,12 +310,6 @@ def use_appropriate_sized_model(self, chain, token_size): def token_size(text): # Shortcut - # NOTE: This sets the token size to default value for Ollama. - # TikToken doesn't seem to support Ollama-based models. - # Local models will not work as intended without dynamically setting the token size. - # TODO: Dynamically set token size for Ollama-based models. - if LanguageModelManager.DEFAULT_PROVIDER == LanguageModelManager.OLLAMA: - return 2048 return LanguageModelManager().get_token_size(text)