feat(server): Add model tests (IBM#6)

Author: OlivierDehaene

README.md

@@ -87,9 +87,4 @@ curl 127.0.0.1:3000/generate \
 ```shell
 make server-dev
 make router-dev
-```
-
-## TODO:
-
-- [ ] Add tests for the `server/model` logic
-- [ ] Backport custom CUDA kernels to Transformers
+```

router/src/batcher.rs

@@ -70,7 +70,7 @@ impl Batcher {
 
         // Notify the background task that we have a new entry in the database that needs
         // to be batched
-        self.shared.batching_task.notify_waiters();
+        self.shared.batching_task.notify_one();
 
         // Await on the response from the background task
         // We can safely unwrap as the background task will never drop the sender

server/Makefile

@@ -8,8 +8,9 @@ gen-server:
 
 install-transformers:
 	# Install specific version of transformers with custom cuda kernels
-	rm transformers || true
-	rm transformers-text_generation_inference || true
+	pip uninstall transformers -y || true
+	rm -rf transformers || true
+	rm -rf transformers-text_generation_inference || true
 	curl -L -O https://github.com/OlivierDehaene/transformers/archive/refs/heads/text_generation_inference.zip
 	unzip text_generation_inference.zip
 	rm text_generation_inference.zip

server/poetry.lock

@@ -22,6 +22,7 @@ bnb = ["bitsandbytes"]
 
 [tool.poetry.group.dev.dependencies]
 grpcio-tools = "^1.49.1"
+pytest = "^7.2.0"
 
 [build-system]
 requires = ["poetry-core>=1.0.0"]

server/pyproject.toml

@@ -0,0 +1,36 @@
+import pytest
+
+from transformers import AutoTokenizer
+
+from text_generation.pb import generate_pb2
+
+
+@pytest.fixture
+def default_pb_parameters():
+    return generate_pb2.LogitsWarperParameters(
+        temperature=1.0,
+        top_k=0,
+        top_p=1.0,
+        do_sample=False,
+    )
+
+
+@pytest.fixture(scope="session")
+def bloom_560m_tokenizer():
+    return AutoTokenizer.from_pretrained("bigscience/bloom-560m", padding_side="left")
+
+
+@pytest.fixture(scope="session")
+def gpt2_tokenizer():
+    tokenizer = AutoTokenizer.from_pretrained("gpt2", padding_side="left")
+    tokenizer.pad_token_id = 50256
+    return tokenizer
+
+
+@pytest.fixture(scope="session")
+def mt0_small_tokenizer():
+    tokenizer = AutoTokenizer.from_pretrained(
+        "bigscience/mt0-small", padding_side="left"
+    )
+    tokenizer.bos_token_id = 0
+    return tokenizer