| from unittest.mock import MagicMock, patch |
|
|
| import pytest |
|
|
| from lm_eval.models.openai_completions import LocalCompletionsAPI |
|
|
|
|
| @pytest.fixture |
| def api(): |
| return LocalCompletionsAPI( |
| base_url="http://test-url.com", tokenizer_backend=None, model="gpt-3.5-turbo" |
| ) |
|
|
|
|
| @pytest.fixture |
| def api_tokenized(): |
| return LocalCompletionsAPI( |
| base_url="http://test-url.com", |
| model="EleutherAI/pythia-1b", |
| tokenizer_backend="huggingface", |
| ) |
|
|
|
|
| def test_create_payload_generate(api): |
| messages = ["Generate a story"] |
| gen_kwargs = { |
| "max_tokens": 100, |
| "temperature": 0.7, |
| "until": ["The End"], |
| "do_sample": True, |
| "seed": 1234, |
| } |
| payload = api._create_payload(messages, generate=True, gen_kwargs=gen_kwargs) |
|
|
| assert payload == { |
| "prompt": ["Generate a story"], |
| "model": "gpt-3.5-turbo", |
| "max_tokens": 100, |
| "temperature": 0.7, |
| "stop": ["The End"], |
| "seed": 1234, |
| } |
|
|
|
|
| def test_create_payload_loglikelihood(api): |
| messages = ["The capital of France is"] |
| payload = api._create_payload(messages, generate=False, gen_kwargs=None) |
|
|
| assert payload == { |
| "model": "gpt-3.5-turbo", |
| "prompt": ["The capital of France is"], |
| "max_tokens": 1, |
| "logprobs": 1, |
| "echo": True, |
| "temperature": 0, |
| "seed": 1234, |
| } |
|
|
|
|
| @pytest.mark.parametrize( |
| "input_messages, generate, gen_kwargs, expected_payload", |
| [ |
| ( |
| ["Hello, how are"], |
| True, |
| {"max_gen_toks": 100, "temperature": 0.7, "until": ["hi"]}, |
| { |
| "prompt": "Hello, how are", |
| "model": "gpt-3.5-turbo", |
| "max_tokens": 100, |
| "temperature": 0.7, |
| "stop": ["hi"], |
| "seed": 1234, |
| }, |
| ), |
| ( |
| ["Hello, how are", "you"], |
| True, |
| {}, |
| { |
| "prompt": "Hello, how are", |
| "model": "gpt-3.5-turbo", |
| "max_tokens": 256, |
| "temperature": 0, |
| "stop": [], |
| "seed": 1234, |
| }, |
| ), |
| ], |
| ) |
| def test_model_generate_call_usage( |
| api, input_messages, generate, gen_kwargs, expected_payload |
| ): |
| with patch("requests.post") as mock_post: |
| mock_response = MagicMock() |
| mock_response.json.return_value = {"result": "success"} |
| mock_post.return_value = mock_response |
|
|
| |
| result = api.model_call( |
| input_messages, generate=generate, gen_kwargs=gen_kwargs |
| ) |
|
|
| |
| mock_post.assert_called_once() |
| _, kwargs = mock_post.call_args |
| assert "json" in kwargs |
| assert kwargs["json"] == expected_payload |
| assert result == {"result": "success"} |
|
|
|
|
| @pytest.mark.parametrize( |
| "input_messages, generate, gen_kwargs, expected_payload", |
| [ |
| ( |
| [[1, 2, 3, 4, 5]], |
| False, |
| None, |
| { |
| "model": "EleutherAI/pythia-1b", |
| "prompt": [[1, 2, 3, 4, 5]], |
| "max_tokens": 1, |
| "logprobs": 1, |
| "echo": True, |
| "seed": 1234, |
| "temperature": 0, |
| }, |
| ), |
| ], |
| ) |
| def test_model_tokenized_call_usage( |
| api_tokenized, input_messages, generate, gen_kwargs, expected_payload |
| ): |
| with patch("requests.post") as mock_post: |
| mock_response = MagicMock() |
| mock_response.json.return_value = {"result": "success"} |
| mock_post.return_value = mock_response |
|
|
| |
| result = api_tokenized.model_call( |
| input_messages, generate=generate, gen_kwargs=gen_kwargs |
| ) |
|
|
| |
| mock_post.assert_called_once() |
| _, kwargs = mock_post.call_args |
| assert "json" in kwargs |
| assert kwargs["json"] == expected_payload |
| assert result == {"result": "success"} |
|
|
