From ac474fce382b0600404a2cc08e8a94904c450665 Mon Sep 17 00:00:00 2001 From: sabaimran Date: Tue, 23 Apr 2024 21:22:23 +0530 Subject: [PATCH] Ensure that the tokenizer and max prompt size are used the wrapper method --- src/khoj/routers/helpers.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/src/khoj/routers/helpers.py b/src/khoj/routers/helpers.py index 20dbb1e4..8fdddc07 100644 --- a/src/khoj/routers/helpers.py +++ b/src/khoj/routers/helpers.py @@ -388,6 +388,7 @@ async def send_message_to_model_wrapper( chat_model = conversation_config.chat_model max_tokens = conversation_config.max_prompt_size + tokenizer = conversation_config.tokenizer if conversation_config.model_type == "offline": if state.offline_chat_processor_config is None or state.offline_chat_processor_config.loaded_model is None: @@ -395,7 +396,12 @@ async def send_message_to_model_wrapper( loaded_model = state.offline_chat_processor_config.loaded_model truncated_messages = generate_chatml_messages_with_context( - user_message=message, system_message=system_message, model_name=chat_model, loaded_model=loaded_model + user_message=message, + system_message=system_message, + model_name=chat_model, + loaded_model=loaded_model, + tokenizer_name=tokenizer, + max_prompt_size=max_tokens, ) return send_message_to_model_offline( @@ -409,7 +415,11 @@ async def send_message_to_model_wrapper( openai_chat_config = await ConversationAdapters.aget_openai_conversation_config() api_key = openai_chat_config.api_key truncated_messages = generate_chatml_messages_with_context( - user_message=message, system_message=system_message, model_name=chat_model, max_prompt_size=max_tokens + user_message=message, + system_message=system_message, + model_name=chat_model, + max_prompt_size=max_tokens, + tokenizer_name=tokenizer, ) openai_response = send_message_to_model(