phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

fix(sagemaker.py): add support for amazon neuron llama models

Browse source
This commit is contained in:
Krrish Dholakia 2023-12-05 16:18:13 -08:00
parent d2dab362df
commit 01fc7f1931
1 changed files with 2 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

4
litellm/llms/sagemaker.py
View file

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