ai-workspace-services/litellm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(sagemaker.py): bring back llama2 templating for sagemaker

Browse Source
This commit is contained in:
Krrish Dholakia 2023-12-05 16:42:19 -08:00
parent 54d8a9df3f
commit 88845dddb1
1 changed files with 6 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
litellm/llms/sagemaker.py
View File

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