Skip to content

Commit

Permalink
Update tests for the new tokenizer
Browse files Browse the repository at this point in the history
  • Loading branch information
AutoPR committed Apr 11, 2023
1 parent 913fcd8 commit 2995f85
Showing 1 changed file with 23 additions and 0 deletions.
23 changes: 23 additions & 0 deletions tests/test_tokenizer.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,23 @@
import unittest
from autopr.utils.tokenizer import num_tokens_from_messages

class TestTokenizer(unittest.TestCase):

def test_num_tokens_from_messages_chat_completions(self):
messages = [
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "What is the capital of France?"},
{"role": "assistant", "content": "The capital of France is Paris."}
]

token_count = num_tokens_from_messages(messages)
self.assertTrue(token_count > 0, "Token count should be greater than 0")

def test_num_tokens_from_messages_ordinary_completions(self):
completion_text = "Lorem ipsum dolor sit amet, consectetur adipiscing elit."

token_count = num_tokens_from_messages(completion_text)
self.assertTrue(token_count > 0, "Token count should be greater than 0")

if __name__ == '__main__':
unittest.main()

0 comments on commit 2995f85

Please sign in to comment.