From 01fc7f1931c2297d79d9f12654e58623a196c627 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Tue, 5 Dec 2023 16:18:13 -0800 Subject: [PATCH] fix(sagemaker.py): add support for amazon neuron llama models --- litellm/llms/sagemaker.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/llms/sagemaker.py b/litellm/llms/sagemaker.py index 1ee43ec2e..a575bf9d1 100644 --- a/litellm/llms/sagemaker.py +++ b/litellm/llms/sagemaker.py @@ -120,8 +120,8 @@ def completion( ) else: hf_model_name = model - if "jumpstart-dft-meta-textgeneration-llama" in model: # llama2 model - if model.endswith("-f") or "-f-" in model: # sagemaker default for a chat model + if "jumpstart-dft-meta-textgeneration-llama" in model or "meta-textgenerationneuron-llama-2-7b" in model: # llama2 model + if model.endswith("-f") or "-f-" in model or "chat" in model: # sagemaker default for a chat model hf_model_name = "meta-llama/Llama-2-7b-chat" # apply the prompt template for a llama2 chat model else: hf_model_name = "meta-llama/Llama-2-7b" # apply the normal prompt template