More fixes

tools/server/tests/unit/test_speculative.py

@@ -25,34 +25,37 @@ def fixture_create_server():
 
 def test_with_and_without_draft():
     global server
-    server.model_draft = None  # disable draft model
-    server.spec_type = None
-    server.start()
-    res = server.make_request("POST", "/completion", data={
+    request = {
         "prompt": "I believe the meaning of life is",
         "temperature": 0.0,
         "top_k": 1,
+        "seed": 4242,
         "n_predict": 16,
-    })
+        "return_tokens": True,
+    }
+
+    server.model_draft = None  # disable draft model
+    server.spec_type = None
+    server.backend_sampling = True
+    server.start()
+    res = server.make_request("POST", "/completion", data=request)
     assert res.status_code == 200
     content_no_draft = res.body["content"]
+    tokens_no_draft = res.body["tokens"]
     server.stop()
 
     # create new server with draft model
     create_server()
     server.backend_sampling = True
     server.start()
-    res = server.make_request("POST", "/completion", data={
-        "prompt": "I believe the meaning of life is",
-        "temperature": 0.0,
-        "top_k": 1,
-        "n_predict": 16,
-    })
+    res = server.make_request("POST", "/completion", data=request)
     assert res.status_code == 200
     assert res.body["timings"]["draft_n"] > 0
     content_draft = res.body["content"]
+    tokens_draft = res.body["tokens"]
 
     assert content_no_draft == content_draft
+    assert tokens_no_draft == tokens_draft
 
 
 def test_different_draft_min_draft_max():