(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,6 +181,7 @@ def test_function_to_dict():
def test_token_counter(): def test_token_counter():
try:
messages = [ messages = [
{ {
"role": "user", "role": "user",
@ -191,7 +192,9 @@ def test_token_counter():
model = "gpt-3.5-turbo", model = "gpt-3.5-turbo",
messages=messages messages=messages
) )
print("gpt-35-turbo")
print(tokens) print(tokens)
assert tokens > 0
tokens = token_counter( tokens = token_counter(
model = "claude-2", model = "claude-2",
@ -199,15 +202,34 @@ def test_token_counter():
) )
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()