From ad9e3e5c0f23c6f72dd77f0dfe37d61c65e84788 Mon Sep 17 00:00:00 2001 From: Alvaro Moran Date: Wed, 20 Nov 2024 09:56:50 +0000 Subject: [PATCH] review: add docstring to explain tests goals --- text-generation-inference/tests/test_tinyllama.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/text-generation-inference/tests/test_tinyllama.py b/text-generation-inference/tests/test_tinyllama.py index 4ab71d06..5566bccd 100644 --- a/text-generation-inference/tests/test_tinyllama.py +++ b/text-generation-inference/tests/test_tinyllama.py @@ -18,6 +18,7 @@ def model_path(): def test_jetstream_info(model_path): + """Verify the model info is correctly loaded and check expected results.""" if not jetstream_pt_available(): pytest.skip("Jetstream PyTorch is not available") generator = AutoGenerator.from_pretrained(model_path, revision="", max_batch_size=1, max_sequence_length=1) @@ -48,6 +49,8 @@ def test_jetstream_info(model_path): ) @pytest.mark.parametrize("batch_size", [1, 4], ids=["single", "multiple"]) def test_jetstream_prefill(input_text, token_id, token_text, do_sample, batch_size, model_path): + """Verify that prefilling a batch with a single request with different sampling techniques. + """ if not jetstream_pt_available(): pytest.skip("Jetstream PyTorch is not available") generator = AutoGenerator.from_pretrained(model_path, revision="", max_batch_size=batch_size, max_sequence_length=SEQUENCE_LENGTH) @@ -71,6 +74,8 @@ def test_jetstream_prefill(input_text, token_id, token_text, do_sample, batch_si def test_jetstream_prefill_change_sampling(model_path): + """Verify changing the sampling strategy between requests in the same batch works as expected. + """ if not jetstream_pt_available(): pytest.skip("Jetstream PyTorch is not available") input_text = "It was a bright cold day in April, and the clocks were striking thirteen." @@ -102,6 +107,9 @@ def check_request(do_sample, expected_token_id, expected_text): def test_jetstream_decode_multiple(model_path): + """Verify that two requests added to the batch at different generation steps + generate the same outputs (continuous batching). + """ if not jetstream_pt_available(): pytest.skip("Jetstream PyTorch is not available") generator = AutoGenerator.from_pretrained(model_path,