PromptQuest / tests /test_token_count.py
NexAddo's picture
Upload 26 files
e4f691b
raw
history blame
976 Bytes
from hackaprompt.completers import get_completer
import os
openai_api_key = os.getenv("OPENAI_API_KEY")
def test_count_tokens_gpt_35_turbo():
completer = get_completer(model="gpt-3.5-turbo", openai_api_key=openai_api_key)
prompt = "hello"
expected_token_count = 1 # taken on https://platform.openai.com/tokenizer
assert completer.get_token_count(prompt) == expected_token_count
prompt = "This is a test to see how many tokens we have"
expected_token_count = 11 # taken on https://platform.openai.com/tokenizer
assert completer.get_token_count(prompt) == expected_token_count
def test_count_tokens_flanT5():
completer = get_completer(model="FlanT5-XXL")
prompt = "hello"
expected_token_count = 2
assert completer.get_token_count(prompt) == expected_token_count
prompt = "This is a test to see how many tokens we have"
expected_token_count = 14
assert completer.get_token_count(prompt) == expected_token_count