huggingface/text-generation-inference
5
0
Fork 0
mirror of https://github.com/huggingface/text-generation-inference.git synced 2025-04-27 13:02:12 +00:00
Code Issues Packages Projects Releases Wiki Activity

test(neuron): simplify sampling test

Browse Source
This commit is contained in:
David Corvoysier 2025-03-12 08:37:28 +00:00
parent 2bb2bdb6a8
commit 23b449d238
1 changed files with 4 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
integration-tests/neuron/test_generate.py
View File

@ -49,17 +49,11 @@ async def test_model_single_request(tgi_service):
max_new_tokens=128,
seed=42,
)
sample_expectations = {
"gpt2": "Deep Learning",
"llama": "Deep Learning",
"mistral": "Deep learning",
"qwen2": "Deep Learning",
"granite": "Deep learning",
}
assert sample_expectations[service_name] in response
# The response must be different
assert not response.startswith(greedy_expectations[service_name])
# Sampling with stop sequence
stop_sequence = sample_expectations[service_name][-5:]
# Sampling with stop sequence (using one of the words returned from the previous test)
stop_sequence = response.split(" ")[-5]
response = await tgi_service.client.text_generation(
"What is Deep Learning?",
do_sample=True,