From 2a94c330284b6bf2c4629199ea9d26be6c3c95fc Mon Sep 17 00:00:00 2001
From: Georgi Gerganov <ggerganov@gmail.com>
Date: Tue, 17 Dec 2024 11:45:18 +0200
Subject: [PATCH] server : be explicit about the pooling type in the tests

ggml-ci
---
 examples/server/tests/unit/test_embedding.py | 9 ++++++++-
 examples/server/tests/utils.py               | 3 +--
 2 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/examples/server/tests/unit/test_embedding.py b/examples/server/tests/unit/test_embedding.py
index d6a3b6125..f60034b4a 100644
--- a/examples/server/tests/unit/test_embedding.py
+++ b/examples/server/tests/unit/test_embedding.py
@@ -14,6 +14,7 @@ def create_server():
 
 def test_embedding_single():
     global server
+    server.pooling = 'last'
     server.start()
     res = server.make_request("POST", "/embeddings", data={
         "input": "I believe the meaning of life is",
@@ -29,6 +30,7 @@ def test_embedding_single():
 
 def test_embedding_multiple():
     global server
+    server.pooling = 'last'
     server.start()
     res = server.make_request("POST", "/embeddings", data={
         "input": [
@@ -75,7 +77,8 @@ def test_embedding_mixed_input(content, is_multi_prompt: bool):
 
 
 def test_embedding_pooling_none():
-    server = ServerPreset.bert_bge_small(pooling = 'none')
+    global server
+    server.pooling = 'none'
     server.start()
     res = server.make_request("POST", "/embeddings", data={
         "input": "hello hello hello",
@@ -88,6 +91,7 @@ def test_embedding_pooling_none():
 
 def test_embedding_openai_library_single():
     global server
+    server.pooling = 'last'
     server.start()
     client = OpenAI(api_key="dummy", base_url=f"http://{server.server_host}:{server.server_port}")
     res = client.embeddings.create(model="text-embedding-3-small", input="I believe the meaning of life is")
@@ -97,6 +101,7 @@ def test_embedding_openai_library_single():
 
 def test_embedding_openai_library_multiple():
     global server
+    server.pooling = 'last'
     server.start()
     client = OpenAI(api_key="dummy", base_url=f"http://{server.server_host}:{server.server_port}")
     res = client.embeddings.create(model="text-embedding-3-small", input=[
@@ -112,6 +117,7 @@ def test_embedding_openai_library_multiple():
 
 def test_embedding_error_prompt_too_long():
     global server
+    server.pooling = 'last'
     server.start()
     res = server.make_request("POST", "/embeddings", data={
         "input": "This is a test " * 512,
@@ -121,6 +127,7 @@ def test_embedding_error_prompt_too_long():
 
 
 def test_same_prompt_give_same_result():
+    server.pooling = 'last'
     server.start()
     res = server.make_request("POST", "/embeddings", data={
         "input": [
diff --git a/examples/server/tests/utils.py b/examples/server/tests/utils.py
index da95c830b..277125e88 100644
--- a/examples/server/tests/utils.py
+++ b/examples/server/tests/utils.py
@@ -275,7 +275,7 @@ class ServerPreset:
         return server
 
     @staticmethod
-    def bert_bge_small(pooling = 'last') -> ServerProcess:
+    def bert_bge_small() -> ServerProcess:
         server = ServerProcess()
         server.model_hf_repo = "ggml-org/models"
         server.model_hf_file = "bert-bge-small/ggml-model-f16.gguf"
@@ -286,7 +286,6 @@ class ServerPreset:
         server.n_slots = 2
         server.seed = 42
         server.server_embeddings = True
-        server.pooling = pooling
         return server
 
     @staticmethod