tuning code

This commit is contained in:
geekan 2024-05-18 15:23:21 +08:00
parent 6d906ce1ed
commit e03db313a2
9 changed files with 40 additions and 39 deletions

View file

@ -66,7 +66,7 @@ class TestRAGEmbeddingFactory:
@pytest.mark.parametrize(
"model, embed_batch_size, expected_params",
[("test_model", 100, {"model_name": "test_model", "embed_batch_size": 100}), (None, None, {})],
[("test_model", 100, {"model": "test_model", "embed_batch_size": 100}), (None, None, {})],
)
def test_try_set_model_and_batch_size(self, mock_config, model, embed_batch_size, expected_params):
# Mock

View file

@ -20,7 +20,7 @@ def _paragraphs(n):
@pytest.mark.parametrize(
"msgs, model_name, system_text, reserved, expected",
"msgs, model, system_text, reserved, expected",
[
(_msgs(), "gpt-3.5-turbo-0613", "System", 1500, 1),
(_msgs(), "gpt-3.5-turbo-16k", "System", 3000, 6),
@ -37,7 +37,7 @@ def test_reduce_message_length(msgs, model_name, system_text, reserved, expected
@pytest.mark.parametrize(
"text, prompt_template, model_name, system_text, reserved, expected",
"text, prompt_template, model, system_text, reserved, expected",
[
(" ".join("Hello World." for _ in range(1000)), "Prompt: {}", "gpt-3.5-turbo-0613", "System", 1500, 2),
(" ".join("Hello World." for _ in range(1000)), "Prompt: {}", "gpt-3.5-turbo-16k", "System", 3000, 1),

View file

@ -7,7 +7,7 @@
"""
import pytest
from metagpt.utils.token_counter import count_message_tokens, count_string_tokens
from metagpt.utils.token_counter import count_input_tokens, count_output_tokens
def test_count_message_tokens():
@ -15,7 +15,7 @@ def test_count_message_tokens():
{"role": "user", "content": "Hello"},
{"role": "assistant", "content": "Hi there!"},
]
assert count_message_tokens(messages) == 15
assert count_input_tokens(messages) == 15
def test_count_message_tokens_with_name():
@ -23,12 +23,12 @@ def test_count_message_tokens_with_name():
{"role": "user", "content": "Hello", "name": "John"},
{"role": "assistant", "content": "Hi there!"},
]
assert count_message_tokens(messages) == 17
assert count_input_tokens(messages) == 17
def test_count_message_tokens_empty_input():
"""Empty input should return 3 tokens"""
assert count_message_tokens([]) == 3
assert count_input_tokens([]) == 3
def test_count_message_tokens_invalid_model():
@ -38,7 +38,7 @@ def test_count_message_tokens_invalid_model():
{"role": "assistant", "content": "Hi there!"},
]
with pytest.raises(NotImplementedError):
count_message_tokens(messages, model="invalid_model")
count_input_tokens(messages, model="invalid_model")
def test_count_message_tokens_gpt_4():
@ -46,27 +46,27 @@ def test_count_message_tokens_gpt_4():
{"role": "user", "content": "Hello"},
{"role": "assistant", "content": "Hi there!"},
]
assert count_message_tokens(messages, model="gpt-4-0314") == 15
assert count_input_tokens(messages, model="gpt-4-0314") == 15
def test_count_string_tokens():
"""Test that the string tokens are counted correctly."""
string = "Hello, world!"
assert count_string_tokens(string, model_name="gpt-3.5-turbo-0301") == 4
assert count_output_tokens(string, model="gpt-3.5-turbo-0301") == 4
def test_count_string_tokens_empty_input():
"""Test that the string tokens are counted correctly."""
assert count_string_tokens("", model_name="gpt-3.5-turbo-0301") == 0
assert count_output_tokens("", model="gpt-3.5-turbo-0301") == 0
def test_count_string_tokens_gpt_4():
"""Test that the string tokens are counted correctly."""
string = "Hello, world!"
assert count_string_tokens(string, model_name="gpt-4-0314") == 4
assert count_output_tokens(string, model="gpt-4-0314") == 4
if __name__ == "__main__":