test(tgi): fix sampling expectations

huggingface · Dec 24, 2024 · 7714081 · 7714081
1 parent fef3fa8
commit 7714081
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 3 deletions.
diff --git a/text-generation-inference/tests/integration/test_generate.py b/text-generation-inference/tests/integration/test_generate.py
@@ -46,7 +46,7 @@ async def test_model_single_request(tgi_service):
         seed=42,
     )
     sample_expectations = {
-        "gpt2": "Deep Learning",
+        "gpt2": "Machine Learning",
         "llama": "Deep Learning",
         "mistral": "Deep learning",
         "qwen2": "Deep Learning",

diff --git a/text-generation-inference/tests/server/test_decode.py b/text-generation-inference/tests/server/test_decode.py
@@ -36,7 +36,7 @@ def _test_decode(config_name, generator, do_sample):
     assert output.finish_reason == 0
     if do_sample:
         expected_text = {
-            "gpt2": " The sun was set",
+            "gpt2": "The only things",
             "llama": "George Orwell, 1984",
             "mistral": "The sky was",
             "qwen2": " A young woman with",

diff --git a/text-generation-inference/tests/server/test_prefill.py b/text-generation-inference/tests/server/test_prefill.py
@@ -35,7 +35,7 @@ def _test_prefill(config_name, generator, batch_size, do_sample):
     assert len(generations) == batch_size
     if do_sample:
         expectations = {
-            "gpt2": [383, " The"],
+            "gpt2": [198, "\n"],
             "llama": [10058, " George"],
             "mistral": [450, " The"],
             "qwen2": [362, " A"],