diff --git a/litellm/llms/sagemaker.py b/litellm/llms/sagemaker.py index 8ddc35b6693..b1eed4e9976 100644 --- a/litellm/llms/sagemaker.py +++ b/litellm/llms/sagemaker.py @@ -58,8 +58,7 @@ def completion( prompt += f"{message['content']}" data = { "inputs": prompt, - # "instruction": prompt, # some baseten models require the prompt to be passed in via the 'instruction' kwarg - **optional_params, + "parameters": optional_params } ## LOGGING diff --git a/litellm/tests/test_completion.py b/litellm/tests/test_completion.py index 1cd6e570a26..e582eee59c7 100644 --- a/litellm/tests/test_completion.py +++ b/litellm/tests/test_completion.py @@ -392,16 +392,19 @@ def test_completion_together_ai(): pytest.fail(f"Error occurred: {e}") -# def test_completion_sagemaker(): -# try: -# response = completion( -# model="sagemaker/jumpstart-dft-meta-textgeneration-llama-2-7b", -# messages=messages -# ) -# # Add any assertions here to check the response -# print(response) -# except Exception as e: -# pytest.fail(f"Error occurred: {e}") +def test_completion_sagemaker(): + try: + response = completion( + model="sagemaker/jumpstart-dft-meta-textgeneration-llama-2-7b", + messages=messages, + temperature=0.2, + max_tokens=80, + ) + # Add any assertions here to check the response + print(response) + except Exception as e: + pytest.fail(f"Error occurred: {e}") + # def test_vertex_ai(): # model_name = "chat-bison" diff --git a/litellm/utils.py b/litellm/utils.py index 4b3a70418fc..2ca755b94b0 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -781,6 +781,21 @@ def get_optional_params( # use the openai defaults if presence_penalty != 0: optional_params["repetition_penalty"] = presence_penalty optional_params["details"] = True + elif custom_llm_provider == "sagemaker": + if "llama-2" in model: + # llama-2 models on sagemaker support the following args + """ + max_new_tokens: Model generates text until the output length (excluding the input context length) reaches max_new_tokens. If specified, it must be a positive integer. + temperature: Controls the randomness in the output. Higher temperature results in output sequence with low-probability words and lower temperature results in output sequence with high-probability words. If temperature -> 0, it results in greedy decoding. If specified, it must be a positive float. + top_p: In each step of text generation, sample from the smallest possible set of words with cumulative probability top_p. If specified, it must be a float between 0 and 1. + return_full_text: If True, input text will be part of the output generated text. If specified, it must be boolean. The default value for it is False. + """ + if max_tokens != float("inf"): + optional_params["max_new_tokens"] = max_tokens + if temperature != 1: + optional_params["temperature"] = temperature + if top_p != 1: + optional_params["top_p"] = top_p elif model in litellm.aleph_alpha_models: if max_tokens != float("inf"): optional_params["maximum_tokens"] = max_tokens