Skip to content

Commit 7093a37

Browse files
committed
test(neuron): use greedy for stop sequences
1 parent 4870824 commit 7093a37

File tree

1 file changed

+6
-10
lines changed

1 file changed

+6
-10
lines changed

integration-tests/neuron/test_generate.py

Lines changed: 6 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -29,15 +29,15 @@ async def test_model_single_request(tgi_service):
2929
assert response.generated_text == greedy_expectations[service_name]
3030

3131
# Greedy bounded with input
32-
response = await tgi_service.client.text_generation(
32+
greedy_response = await tgi_service.client.text_generation(
3333
"What is Deep Learning?",
3434
max_new_tokens=17,
3535
return_full_text=True,
3636
details=True,
3737
decoder_input_details=True,
3838
)
39-
assert response.details.generated_tokens == 17
40-
assert response.generated_text == prompt + greedy_expectations[service_name]
39+
assert greedy_response.details.generated_tokens == 17
40+
assert greedy_response.generated_text == prompt + greedy_expectations[service_name]
4141

4242
# Sampling
4343
response = await tgi_service.client.text_generation(
@@ -52,16 +52,12 @@ async def test_model_single_request(tgi_service):
5252
# The response must be different
5353
assert not response.startswith(greedy_expectations[service_name])
5454

55-
# Sampling with stop sequence (using one of the words returned from the previous test)
56-
stop_sequence = response.split(" ")[-5]
55+
# Greedy with stop sequence (using one of the words returned from the previous test)
56+
stop_sequence = greedy_response.generated_text.split(" ")[-5]
5757
response = await tgi_service.client.text_generation(
5858
"What is Deep Learning?",
59-
do_sample=True,
60-
top_k=50,
61-
top_p=0.9,
62-
repetition_penalty=1.2,
59+
do_sample=False,
6360
max_new_tokens=128,
64-
seed=42,
6561
stop_sequences=[stop_sequence],
6662
)
6763
assert response.endswith(stop_sequence)

0 commit comments

Comments
 (0)