fix tests

integration-tests/conftest.py

@@ -332,7 +332,8 @@ def generate_load():
         client: AsyncClient, prompt: str, max_new_tokens: int, n: int
     ) -> List[Response]:
         futures = [
-            client.generate(prompt, max_new_tokens=max_new_tokens) for _ in range(n)
+            client.generate(prompt, max_new_tokens=max_new_tokens, prefill_details=True)
+            for _ in range(n)
         ]
 
         return await asyncio.gather(*futures)

integration-tests/models/test_bloom_560m.py

@@ -19,6 +19,7 @@ async def test_bloom_560m(bloom_560, response_snapshot):
         "Pour déguster un ortolan, il faut tout d'abord",
         max_new_tokens=10,
         top_p=0.9,
+        prefill_details=True,
         seed=0,
     )
 
@@ -40,6 +41,7 @@ async def test_bloom_560m_all_params(bloom_560, response_snapshot):
         truncate=5,
         typical_p=0.9,
         watermark=True,
+        prefill_details=True,
         seed=0,
     )
 

integration-tests/models/test_bloom_560m_sharded.py

@@ -19,6 +19,7 @@ async def test_bloom_560m_sharded(bloom_560m_sharded, response_snapshot):
         "Pour déguster un ortolan, il faut tout d'abord",
         max_new_tokens=10,
         top_p=0.9,
+        prefill_details=True,
         seed=0,
     )
 

integration-tests/models/test_flash_falcon.py

@@ -19,6 +19,7 @@ async def test_flash_falcon(flash_falcon, response_snapshot):
     response = await flash_falcon.generate(
         "Girafatron is obsessed with giraffes, the most glorious animal on the face of this Earth. Giraftron believes all other animals are irrelevant when compared to the glorious majesty of the giraffe.\nDaniel: Hello, Girafatron!\nGirafatron:",
         max_new_tokens=10,
+        prefill_details=True,
     )
 
     assert response.details.generated_tokens == 10
@@ -40,6 +41,7 @@ async def test_flash_falcon_all_params(flash_falcon, response_snapshot):
         truncate=5,
         typical_p=0.9,
         watermark=True,
+        prefill_details=True,
         seed=0,
     )
 

integration-tests/models/test_flash_llama.py

@@ -16,7 +16,9 @@ async def flash_llama(flash_llama_handle):
 @pytest.mark.asyncio
 @pytest.mark.private
 async def test_flash_llama(flash_llama, response_snapshot):
-    response = await flash_llama.generate("Test request", max_new_tokens=10)
+    response = await flash_llama.generate(
+        "Test request", max_new_tokens=10, prefill_details=True
+    )
 
     assert response.details.generated_tokens == 10
     assert response == response_snapshot
@@ -37,6 +39,7 @@ async def test_flash_llama_all_params(flash_llama, response_snapshot):
         truncate=5,
         typical_p=0.9,
         watermark=True,
+        prefill_details=True,
         seed=0,
     )
 

integration-tests/models/test_flash_neox.py

@@ -18,6 +18,7 @@ async def test_flash_neox(flash_neox, response_snapshot):
     response = await flash_neox.generate(
         "<|USER|>What's your mood today?<|ASSISTANT|>",
         max_new_tokens=10,
+        prefill_details=True,
     )
 
     assert response.details.generated_tokens == 10

integration-tests/models/test_flash_neox_sharded.py

@@ -18,6 +18,7 @@ async def test_flash_neox(flash_neox_sharded, response_snapshot):
     response = await flash_neox_sharded.generate(
         "<|prompter|>What is a meme, and what's the history behind this word?<|endoftext|><|assistant|>",
         max_new_tokens=10,
+        prefill_details=True,
     )
 
     assert response.details.generated_tokens == 10

integration-tests/models/test_flash_santacoder.py

@@ -15,7 +15,9 @@ async def flash_santacoder(flash_santacoder_handle):
 
 @pytest.mark.asyncio
 async def test_flash_santacoder(flash_santacoder, response_snapshot):
-    response = await flash_santacoder.generate("def print_hello", max_new_tokens=10)
+    response = await flash_santacoder.generate(
+        "def print_hello", max_new_tokens=10, prefill_details=True
+    )
 
     assert response.details.generated_tokens == 10
     assert response == response_snapshot

integration-tests/models/test_flash_starcoder.py

@@ -16,7 +16,9 @@ async def flash_starcoder(flash_starcoder_handle):
 @pytest.mark.asyncio
 @pytest.mark.private
 async def test_flash_starcoder(flash_starcoder, response_snapshot):
-    response = await flash_starcoder.generate("def print_hello", max_new_tokens=10)
+    response = await flash_starcoder.generate(
+        "def print_hello", max_new_tokens=10, prefill_details=True
+    )
 
     assert response.details.generated_tokens == 10
     assert response == response_snapshot
@@ -26,7 +28,12 @@ async def test_flash_starcoder(flash_starcoder, response_snapshot):
 @pytest.mark.private
 async def test_flash_starcoder_default_params(flash_starcoder, response_snapshot):
     response = await flash_starcoder.generate(
-        "def print_hello", max_new_tokens=60, temperature=0.2, top_p=0.95, seed=0
+        "def print_hello",
+        max_new_tokens=60,
+        temperature=0.2,
+        top_p=0.95,
+        prefill_details=True,
+        seed=0,
     )
 
     assert response.details.generated_tokens == 60

integration-tests/models/test_mt0_base.py

@@ -19,6 +19,7 @@ async def test_mt0_base(mt0_base, response_snapshot):
         "Why is the sky blue?",
         max_new_tokens=10,
         top_p=0.9,
+        prefill_details=True,
         seed=0,
     )
 
@@ -40,6 +41,7 @@ async def test_mt0_base_all_params(mt0_base, response_snapshot):
         truncate=5,
         typical_p=0.9,
         watermark=True,
+        prefill_details=True,
         seed=0,
     )
 

integration-tests/models/test_t5_sharded.py

@@ -18,6 +18,7 @@ async def test_t5_sharded(t5_sharded, response_snapshot):
     response = await t5_sharded.generate(
         "Please answer the following question. What is the boiling point of Nitrogen?",
         max_new_tokens=10,
+        prefill_details=True,
     )
 
     assert response == response_snapshot

server/tests/models/test_bloom.py

@@ -24,6 +24,7 @@ def default_pb_request(default_pb_parameters, default_pb_stop_parameters):
     return generate_pb2.Request(
         id=0,
         inputs="Test",
+        prefill_logprobs=True,
         truncate=100,
         parameters=default_pb_parameters,
         stopping_parameters=default_pb_stop_parameters,

server/tests/models/test_causal_lm.py

@@ -25,6 +25,7 @@ def default_pb_request(default_pb_parameters, default_pb_stop_parameters):
     return generate_pb2.Request(
         id=0,
         inputs="Test",
+        prefill_logprobs=True,
         truncate=100,
         parameters=default_pb_parameters,
         stopping_parameters=default_pb_stop_parameters,

server/tests/models/test_santacoder.py

@@ -15,6 +15,7 @@ def default_pb_request(default_pb_parameters, default_pb_stop_parameters):
     return generate_pb2.Request(
         id=0,
         inputs="def",
+        prefill_logprobs=True,
         truncate=100,
         parameters=default_pb_parameters,
         stopping_parameters=default_pb_stop_parameters,
@@ -31,6 +32,7 @@ def default_fim_pb_request(default_pb_parameters, default_pb_stop_parameters):
     return generate_pb2.Request(
         id=0,
         inputs="<fim-prefix>def<fim-suffix>world<fim-middle>",
+        prefill_logprobs=True,
         truncate=100,
         parameters=default_pb_parameters,
         stopping_parameters=default_pb_stop_parameters,

server/tests/models/test_seq2seq_lm.py

@@ -28,6 +28,7 @@ def default_pb_request(default_pb_parameters, default_pb_stop_parameters):
     return generate_pb2.Request(
         id=0,
         inputs="Test",
+        prefill_logprobs=True,
         truncate=100,
         parameters=default_pb_parameters,
         stopping_parameters=default_pb_stop_parameters,