(test) token_counter

This commit is contained in:
ishaan-jaff 2023-11-13 08:02:46 -08:00
parent c91abc8ad1
commit 7a3607e00c

View file

@ -181,33 +181,55 @@ def test_function_to_dict():
def test_token_counter(): def test_token_counter():
messages = [ try:
{ messages = [
"role": "user", {
"content": "hi how are you what time is it" "role": "user",
} "content": "hi how are you what time is it"
] }
tokens = token_counter( ]
model = "gpt-3.5-turbo", tokens = token_counter(
messages=messages model = "gpt-3.5-turbo",
) messages=messages
print(tokens) )
print("gpt-35-turbo")
print(tokens)
assert tokens > 0
tokens = token_counter( tokens = token_counter(
model = "claude-2", model = "claude-2",
messages=messages messages=messages
) )
print("claude-2") print("claude-2")
print(tokens) print(tokens)
assert tokens > 0
tokens = token_counter( tokens = token_counter(
model = "command-nightly", model = "palm/chat-bison",
messages=messages messages=messages
) )
print("command-nightly") print("palm/chat-bison")
print(tokens) print(tokens)
assert tokens > 0
# test_token_counter() tokens = token_counter(
model = "ollama/llama2",
messages=messages
)
print("ollama/llama2")
print(tokens)
assert tokens > 0
tokens = token_counter(
model = "anthropic.claude-instant-v1",
messages=messages
)
print("anthropic.claude-instant-v1")
print(tokens)
assert tokens > 0
except Exception as e:
pytest.fail(f"Error occurred: {e}")
test_token_counter()