phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-25 18:54:30 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix(sagemaker.py): prompt templating fixes

Browse source
This commit is contained in:
Krrish Dholakia 2023-12-05 17:47:37 -08:00
parent 0eccc1b1f8
commit 648d41c96f
2 changed files with 24 additions and 12 deletions
Download patch file Download diff file Expand all files Collapse all files

8
litellm/llms/sagemaker.py
View file

@ -121,10 +121,10 @@ def completion(
)
else:
if hf_model_name is None:
if "llama2" in model.lower(): # llama2 model
if "chat" in model.lower():
if "llama-2" in model.lower(): # llama-2 model
if "chat" in model.lower(): # apply llama2 chat template
hf_model_name = "meta-llama/Llama-2-7b-chat-hf"
else:
else: # apply regular llama2 template
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)
@ -146,7 +146,7 @@ def completion(
logging_obj.pre_call(
input=prompt,
api_key="",
additional_args={"complete_input_dict": data, "request_str": request_str},
additional_args={"complete_input_dict": data, "request_str": request_str, "hf_model_name": hf_model_name},
)
## COMPLETION CALL
try: