53 lines
1.6 KiB
Python
53 lines
1.6 KiB
Python
# -*- coding: utf-8 -*-
|
|
""" Unit test for token_utils."""
|
|
import unittest
|
|
from agentscope.utils.token_utils import get_openai_max_length
|
|
from agentscope.utils.token_utils import count_openai_token
|
|
|
|
|
|
class TokenUtilsTest(unittest.TestCase):
|
|
"""Token utils test."""
|
|
|
|
def test_get_openai_max_length(self) -> None:
|
|
"""Test the function get_openai_max_length."""
|
|
self.assertEqual(get_openai_max_length("gpt-4"), 8192)
|
|
self.assertEqual(get_openai_max_length("gpt-3.5-turbo"), 4096)
|
|
with self.assertRaises(KeyError):
|
|
get_openai_max_length("non-existing-model")
|
|
|
|
def test_count_openai_token(self) -> None:
|
|
"""Test the function count_openai_token."""
|
|
test_model_1 = "text-davinci-003"
|
|
test_content_str = "This is a test string."
|
|
self.assertEqual(
|
|
count_openai_token(test_content_str, test_model_1),
|
|
6,
|
|
)
|
|
|
|
test_content_list = [
|
|
{
|
|
"role": "system",
|
|
"content": "This is a system content.",
|
|
},
|
|
{
|
|
"role": "user",
|
|
"name": "example_user",
|
|
"content": "This is a user content.",
|
|
},
|
|
]
|
|
|
|
test_model_2 = "gpt-4"
|
|
self.assertEqual(
|
|
count_openai_token(test_content_list, test_model_2),
|
|
26,
|
|
)
|
|
|
|
# Test unsupported model
|
|
unsupported_model = "unsupported-model"
|
|
with self.assertRaises(NotImplementedError):
|
|
count_openai_token(test_content_str, unsupported_model)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|