From 88845dddb1c25ebeda244ae2526fd9ba8df21944 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Tue, 5 Dec 2023 16:42:19 -0800 Subject: [PATCH] fix(sagemaker.py): bring back llama2 templating for sagemaker --- litellm/llms/sagemaker.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/litellm/llms/sagemaker.py b/litellm/llms/sagemaker.py index e1c9ccdc8c7..ca71461cfcd 100644 --- a/litellm/llms/sagemaker.py +++ b/litellm/llms/sagemaker.py @@ -120,6 +120,12 @@ def completion( messages=messages ) else: + if hf_model_name is None: + if "llama2" in model.lower(): # llama2 model + if "chat" in model.lower(): + hf_model_name = "meta-llama/Llama-2-7b-chat-hf" + else: + hf_model_name = "meta-llama/Llama-2-7b" hf_model_name = hf_model_name or model # pass in hf model name for pulling it's prompt template - (e.g. `hf_model_name="meta-llama/Llama-2-7b-chat-hf` applies the llama2 chat template to the prompt) prompt = prompt_factory(model=hf_model_name, messages=messages)