Skip to content

Commit

Permalink
fix(tests): correct slow-tests expected texts for do_sample=True
Browse files Browse the repository at this point in the history
Add the top_k filtering to increase variability when doing tests with
do_sample.
  • Loading branch information
tengomucho committed Sep 24, 2024
1 parent 2d62758 commit 221a5da
Showing 1 changed file with 11 additions and 2 deletions.
13 changes: 11 additions & 2 deletions text-generation-inference/tests/test_decode.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@

from dataclasses import dataclass

import pytest
Expand All @@ -17,6 +16,7 @@ class DecodeTestParams:
expected_text: str
do_sample: bool = False
max_new_tokens: int = 20
top_k: int = 50


@pytest.mark.parametrize("params",
Expand Down Expand Up @@ -70,7 +70,13 @@ def _test_decode_single(params):
generator = AutoGenerator.from_pretrained(
model_path, revision="", max_batch_size=1, max_sequence_length=params.sequence_length
)
request = create_request(id=0, inputs=input_text, max_new_tokens=max_new_tokens, do_sample=params.do_sample)
request = create_request(
id=0,
inputs=input_text,
max_new_tokens=max_new_tokens,
do_sample=params.do_sample,
top_k=params.top_k,
)
batch = Batch(id=0, requests=[request], size=1, max_tokens=params.sequence_length)
generations, next_batch = generator.prefill(batch)
# We already generated one token: call decode max_new_tokens - 1 times
Expand All @@ -88,6 +94,7 @@ def _test_decode_single(params):
output = generations[0].generated_text
assert output.generated_tokens == max_new_tokens
assert output.finish_reason == 0
print(f"Generated text: {output.text}")
if params.do_sample:
assert output.text != params.expected_text
else:
Expand All @@ -102,11 +109,13 @@ def _test_decode_single(params):
model_id="meta-llama/Llama-2-7b-hf",
sequence_length=256,
expected_text="\nWinston Smith, his chin nuzzled into his breast in an effort to escape",
top_k=100,
),
DecodeTestParams(
model_id="meta-llama/Meta-Llama-3-8B",
sequence_length=256,
expected_text=" Winston Smith, his chin nuzzled into his breast in an effort to escape the vile wind,",
top_k=100,
),
],
ids=["Llama-2-7b-hf", "Meta-Llama-3-8B"],
Expand Down

0 comments on commit 221a5da

Please sign in to comment.