fix(factory.py): fix tgai rendering template

This commit is contained in:
Krrish Dholakia 2023-12-13 12:27:31 -08:00
parent 69c29f8f86
commit 75bcb37cb2
2 changed files with 32 additions and 5 deletions

View file

@ -161,6 +161,8 @@ def phind_codellama_pt(messages):
def hf_chat_template(model: str, messages: list, chat_template: Optional[Any]=None):
## get the tokenizer config from huggingface
bos_token = ""
eos_token = ""
if chat_template is None:
def _get_tokenizer_config(hf_model_name):
url = f"https://huggingface.co/{hf_model_name}/raw/main/tokenizer_config.json"
@ -187,7 +189,10 @@ def hf_chat_template(model: str, messages: list, chat_template: Optional[Any]=No
# Create a template object from the template text
env = Environment()
env.globals['raise_exception'] = raise_exception
template = env.from_string(chat_template)
try:
template = env.from_string(chat_template)
except Exception as e:
raise e
def _is_system_in_template():
try:
@ -227,8 +232,8 @@ def hf_chat_template(model: str, messages: list, chat_template: Optional[Any]=No
new_messages.append(reformatted_messages[-1])
rendered_text = template.render(bos_token=bos_token, eos_token=eos_token, messages=new_messages)
return rendered_text
except:
raise Exception("Error rendering template")
except Exception as e:
raise Exception(f"Error rendering template - {str(e)}")
# Anthropic template
def claude_2_1_pt(messages: list): # format - https://docs.anthropic.com/claude/docs/how-to-use-system-prompts
@ -283,6 +288,9 @@ def get_model_info(token, model):
return None, None
def format_prompt_togetherai(messages, prompt_format, chat_template):
if prompt_format is None:
return default_pt(messages)
human_prompt, assistant_prompt = prompt_format.split('{prompt}')
if chat_template is not None:

View file

@ -294,7 +294,7 @@ def hf_test_completion_tgi():
print(response)
except Exception as e:
pytest.fail(f"Error occurred: {e}")
hf_test_completion_tgi()
# hf_test_completion_tgi()
# ################### Hugging Face Conversational models ########################
# def hf_test_completion_conv():
@ -708,7 +708,7 @@ def test_completion_azure():
except Exception as e:
pytest.fail(f"Error occurred: {e}")
test_completion_azure()
# test_completion_azure()
def test_azure_openai_ad_token():
# this tests if the azure ad token is set in the request header
@ -1026,6 +1026,25 @@ def test_completion_together_ai():
except Exception as e:
pytest.fail(f"Error occurred: {e}")
def test_completion_together_ai_mixtral():
model_name = "together_ai/DiscoResearch/DiscoLM-mixtral-8x7b-v2"
try:
messages =[
{"role": "user", "content": "Who are you"},
{"role": "assistant", "content": "I am your helpful assistant."},
{"role": "user", "content": "Tell me a joke"},
]
response = completion(model=model_name, messages=messages, max_tokens=256, n=1, logger_fn=logger_fn)
# Add any assertions here to check the response
print(response)
cost = completion_cost(completion_response=response)
assert cost > 0.0
print("Cost for completion call together-computer/llama-2-70b: ", f"${float(cost):.10f}")
except Exception as e:
pytest.fail(f"Error occurred: {e}")
test_completion_together_ai_mixtral()
def test_completion_together_ai_yi_chat():
model_name = "together_ai/zero-one-ai/Yi-34B-Chat"
try: