scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-12 23:01:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 89

fix(utils.py): support max token adjustment for sagemaker

Browse source
This commit is contained in:
Krrish Dholakia 2024-01-31 19:09:54 -08:00
parent 93a52a2d35
commit a32639fa79
3 changed files with 41 additions and 10 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
litellm/tests/test_completion_cost.py
View file

@ -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)

28
litellm/tests/test_model_max_token_adjust.py Normal file
View file

@ -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()

21
litellm/utils.py
View file

@ -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: