From a0daac212d832243c20beb0b5be0a5c7c6667de1 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Wed, 31 Jan 2024 17:01:46 -0800 Subject: [PATCH 1/6] fix(utils.py): support checking if user defined max tokens exceeds model limit --- litellm/utils.py | 35 +++++++++++++++++++++++++++++++++-- 1 file changed, 33 insertions(+), 2 deletions(-) diff --git a/litellm/utils.py b/litellm/utils.py index d102476f315..cec43cb8ee1 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -2186,6 +2186,34 @@ def client(original_function): ) else: return cached_result + + # CHECK MAX TOKENS + if ( + kwargs("max_tokens", None) is not None + and model is not None + and litellm.drop_params + == True # user is okay with params being modified + and ( + call_type == CallTypes.acompletion.value + or call_type == CallTypes.completion.value + ) + ): + try: + max_output_tokens = get_max_tokens(model=model) + user_max_tokens = kwargs.get("max_tokens") + ## Scenario 1: User limit > model limit + if user_max_tokens > max_output_tokens: + user_max_tokens = max_output_tokens + ## Scenario 2: User limit + prompt > model limit + input_tokens = token_counter( + model=model, messages=kwargs.get("messages") + ) + if input_tokens > max_output_tokens: + pass # allow call to fail normally + elif user_max_tokens + input_tokens > max_output_tokens: + user_max_tokens = max_output_tokens - input_tokens + except Exception as e: + print_verbose(f"Error while checking max token limit: {str(e)}") # MODEL CALL result = original_function(*args, **kwargs) end_time = datetime.datetime.now() @@ -4503,7 +4531,7 @@ def get_api_key(llm_provider: str, dynamic_api_key: Optional[str]): def get_max_tokens(model: str): """ - Get the maximum number of tokens allowed for a given model. + Get the maximum number of output tokens allowed for a given model. Parameters: model (str): The name of the model. @@ -4543,7 +4571,10 @@ def get_max_tokens(model: str): try: if model in litellm.model_cost: - return litellm.model_cost[model]["max_tokens"] + if "max_output_tokens" in litellm.model_cost[model]: + return litellm.model_cost[model]["max_output_tokens"] + elif "max_tokens" in litellm.model_cost[model]: + return litellm.model_cost[model]["max_tokens"] model, custom_llm_provider, _, _ = get_llm_provider(model=model) if custom_llm_provider == "huggingface": max_tokens = _get_max_position_embeddings(model_name=model) From 9dc972de70e41f2e4b2ab72a6971c26d97596e62 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Wed, 31 Jan 2024 18:15:47 -0800 Subject: [PATCH 2/6] fix(utils.py): fix get for dict --- litellm/utils.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/litellm/utils.py b/litellm/utils.py index cec43cb8ee1..8acabd28506 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -2189,7 +2189,7 @@ def client(original_function): # CHECK MAX TOKENS if ( - kwargs("max_tokens", None) is not None + kwargs.get("max_tokens", None) is not None and model is not None and litellm.drop_params == True # user is okay with params being modified @@ -2205,9 +2205,12 @@ def client(original_function): if user_max_tokens > max_output_tokens: user_max_tokens = max_output_tokens ## Scenario 2: User limit + prompt > model limit - input_tokens = token_counter( - model=model, messages=kwargs.get("messages") - ) + messages = None + if len(args) > 1: + messages = args[1] + elif kwargs.get("messages", None): + messages = kwargs["messages"] + input_tokens = token_counter(model=model, messages=messages) if input_tokens > max_output_tokens: pass # allow call to fail normally elif user_max_tokens + input_tokens > max_output_tokens: From 93a52a2d358083ae095f7ebbd85b1db9ed1f4d83 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Wed, 31 Jan 2024 18:23:44 -0800 Subject: [PATCH 3/6] fix(utils.py): set call_type at the top of the function --- litellm/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/utils.py b/litellm/utils.py index 8acabd28506..c9fccdc6c0e 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -2094,13 +2094,13 @@ def client(original_function): logging_obj = kwargs.get("litellm_logging_obj", None) # only set litellm_call_id if its not in kwargs + call_type = original_function.__name__ if "litellm_call_id" not in kwargs: kwargs["litellm_call_id"] = str(uuid.uuid4()) try: model = args[0] if len(args) > 0 else kwargs["model"] except: model = None - call_type = original_function.__name__ if ( call_type != CallTypes.image_generation.value and call_type != CallTypes.text_completion.value From a32639fa7958c10263905c334ecc62c34bb475b8 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Wed, 31 Jan 2024 19:09:54 -0800 Subject: [PATCH 4/6] fix(utils.py): support max token adjustment for sagemaker --- litellm/tests/test_completion_cost.py | 2 +- litellm/tests/test_model_max_token_adjust.py | 28 ++++++++++++++++++++ litellm/utils.py | 21 ++++++++------- 3 files changed, 41 insertions(+), 10 deletions(-) create mode 100644 litellm/tests/test_model_max_token_adjust.py diff --git a/litellm/tests/test_completion_cost.py b/litellm/tests/test_completion_cost.py index b117223ab08..b55f9c9d6e0 100644 --- a/litellm/tests/test_completion_cost.py +++ b/litellm/tests/test_completion_cost.py @@ -13,7 +13,7 @@ import pytest def test_get_gpt3_tokens(): max_tokens = get_max_tokens("gpt-3.5-turbo") print(max_tokens) - assert max_tokens == 4097 + assert max_tokens == 4096 # print(results) diff --git a/litellm/tests/test_model_max_token_adjust.py b/litellm/tests/test_model_max_token_adjust.py new file mode 100644 index 00000000000..b026b625677 --- /dev/null +++ b/litellm/tests/test_model_max_token_adjust.py @@ -0,0 +1,28 @@ +# What this tests? +## Tests if max tokens get adjusted, if over limit + +import sys, os, time +import traceback, asyncio +import pytest + +sys.path.insert( + 0, os.path.abspath("../..") +) # Adds the parent directory to the system path +import litellm +from litellm import completion + +litellm.drop_params = True + + +def test_completion_sagemaker(): + response = completion( + model="sagemaker/berri-benchmarking-Llama-2-70b-chat-hf-4", + messages=[{"content": "Hello, how are you?", "role": "user"}], + temperature=0.2, + max_tokens=80000, + hf_model_name="meta-llama/Llama-2-70b-chat-hf", + ) + print(f"response: {response}") + + +# test_completion_sagemaker() diff --git a/litellm/utils.py b/litellm/utils.py index c9fccdc6c0e..43319246f12 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -2199,22 +2199,28 @@ def client(original_function): ) ): try: - max_output_tokens = get_max_tokens(model=model) + base_model = model + if kwargs.get("hf_model_name", None) is not None: + base_model = f"huggingface/{kwargs.get('hf_model_name')}" + max_output_tokens = ( + get_max_tokens(model=base_model) or 4096 + ) # assume min context window is 4k tokens user_max_tokens = kwargs.get("max_tokens") - ## Scenario 1: User limit > model limit - if user_max_tokens > max_output_tokens: - user_max_tokens = max_output_tokens - ## Scenario 2: User limit + prompt > model limit + ## Scenario 1: User limit + prompt > model limit messages = None if len(args) > 1: messages = args[1] elif kwargs.get("messages", None): messages = kwargs["messages"] - input_tokens = token_counter(model=model, messages=messages) + input_tokens = token_counter(model=base_model, messages=messages) + input_tokens += max( + 0.1 * input_tokens, 10 + ) # give at least a 10 token buffer. token counting can be imprecise. if input_tokens > max_output_tokens: pass # allow call to fail normally elif user_max_tokens + input_tokens > max_output_tokens: user_max_tokens = max_output_tokens - input_tokens + kwargs["max_tokens"] = user_max_tokens except Exception as e: print_verbose(f"Error while checking max token limit: {str(e)}") # MODEL CALL @@ -4553,7 +4559,6 @@ def get_max_tokens(model: str): def _get_max_position_embeddings(model_name): # Construct the URL for the config.json file config_url = f"https://huggingface.co/{model_name}/raw/main/config.json" - try: # Make the HTTP request to get the raw JSON file response = requests.get(config_url) @@ -4561,10 +4566,8 @@ def get_max_tokens(model: str): # Parse the JSON response config_json = response.json() - # Extract and return the max_position_embeddings max_position_embeddings = config_json.get("max_position_embeddings") - if max_position_embeddings is not None: return max_position_embeddings else: From 9593df23c4ada5de5249df586f4355df6da4582a Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Wed, 31 Jan 2024 19:31:07 -0800 Subject: [PATCH 5/6] test: add more logging --- litellm/tests/test_model_max_token_adjust.py | 1 + litellm/utils.py | 1 + 2 files changed, 2 insertions(+) diff --git a/litellm/tests/test_model_max_token_adjust.py b/litellm/tests/test_model_max_token_adjust.py index b026b625677..8469302a527 100644 --- a/litellm/tests/test_model_max_token_adjust.py +++ b/litellm/tests/test_model_max_token_adjust.py @@ -15,6 +15,7 @@ litellm.drop_params = True def test_completion_sagemaker(): + litellm.set_verbose = True response = completion( model="sagemaker/berri-benchmarking-Llama-2-70b-chat-hf-4", messages=[{"content": "Hello, how are you?", "role": "user"}], diff --git a/litellm/utils.py b/litellm/utils.py index 43319246f12..a954353c609 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -2220,6 +2220,7 @@ def client(original_function): pass # allow call to fail normally elif user_max_tokens + input_tokens > max_output_tokens: user_max_tokens = max_output_tokens - input_tokens + print_verbose(f"user_max_tokens: {user_max_tokens}") kwargs["max_tokens"] = user_max_tokens except Exception as e: print_verbose(f"Error while checking max token limit: {str(e)}") From a6dd3c0bf5297ec4580cf7fae3cee112fa59140b Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Wed, 31 Jan 2024 20:07:28 -0800 Subject: [PATCH 6/6] test: move drop_params inside test --- litellm/tests/test_model_max_token_adjust.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/litellm/tests/test_model_max_token_adjust.py b/litellm/tests/test_model_max_token_adjust.py index 8469302a527..b4d48b5e28e 100644 --- a/litellm/tests/test_model_max_token_adjust.py +++ b/litellm/tests/test_model_max_token_adjust.py @@ -11,11 +11,10 @@ sys.path.insert( import litellm from litellm import completion -litellm.drop_params = True - def test_completion_sagemaker(): litellm.set_verbose = True + litellm.drop_params = True response = completion( model="sagemaker/berri-benchmarking-Llama-2-70b-chat-hf-4", messages=[{"content": "Hello, how are you?", "role": "user"}],