test(tgi): fix sampling expectations

huggingface · Jan 2, 2025 · 00c9655 · 00c9655
1 parent 2066ac8
commit 00c9655
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
@@ -47,7 +47,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
@@ -40,7 +40,7 @@ def _test_decode(config_name, generator, do_sample):
             "llama": "George Orwell",
             "mistral": "The sky is black",
             "qwen2": " I stood in the back yard",
-            "granite": "Aldous Huxley, Brave New World",
+            "granite": "Nineteen eighty-four",
         }[config_name]
         assert expected_text in output.text
     else:

diff --git a/text-generation-inference/tests/server/test_prefill.py b/text-generation-inference/tests/server/test_prefill.py
@@ -39,7 +39,7 @@ def _test_prefill(config_name, generator, batch_size, do_sample):
             "llama": [10058, " George"],
             "mistral": [450, " The"],
             "qwen2": [358, " I"],
-            "granite": [429, " -"],
+            "granite": [489, " N"],
         }[config_name]
     else:
         expectations = {