add reference tensors

Author: fxmarty

tests/README.md

@@ -8,4 +8,27 @@ pip install -r requirements.txt
 and mounting a volume for the tests, they can be run from within the container with
 ```
 pytest tests/ -s -vvvvv
+```
+
+## Reference outputs
+
+For example, collecting the reference on an RTX 4090 on Candle backend:
+```
+docker run --rm -it --gpus all --net host --entrypoint "/bin/bash" -v $(pwd):/tei ghcr.io/huggingface/text-embeddings-inference:89-1.2.3
+```
+and
+```
+text-embeddings-router --model-id sentence-transformers/all-MiniLM-L6-v2
+```
+
+and then
+```
+python collect.py --model-id sentence-transformers/all-MiniLM-L6-v2 --n_inp 1 --flash
+python collect.py --model-id sentence-transformers/all-MiniLM-L6-v2 --n_inp 3 --flash
+```
+
+Restart server with `USE_FLASH_ATTENTION=0`, and
+```
+python collect.py --model-id sentence-transformers/all-MiniLM-L6-v2 --n_inp 1
+python collect.py --model-id sentence-transformers/all-MiniLM-L6-v2 --n_inp 3
 ```

tests/assets/default_bert.pt

@@ -0,0 +1,37 @@
+
+import requests
+import torch
+import argparse
+import json
+import os
+
+parser = argparse.ArgumentParser(description='Assets collection')
+parser.add_argument('--model-id', help='Model id', required=True)
+parser.add_argument('--n_inp', help='Number of inputs', required=True, type=int)
+parser.add_argument('--flash', action='store_true')
+
+args = parser.parse_args()
+
+url = f"http://0.0.0.0:80/embed"
+
+INPUTS = [
+    "What is Deep Learning?",
+    "Today I am in Paris and I would like to",
+    "Paris weather is",
+    "Great job"
+]
+
+data = {"inputs": INPUTS[:args.n_inp]}
+headers = {"Content-Type": "application/json"}
+
+response = requests.post(url, json=data, headers=headers)
+
+embedding = torch.Tensor(json.loads(response.text))
+
+postfix = ""
+if not args.flash:
+    postfix = "_no_flash"
+
+save_path = f"./assets/{args.model_id.replace('/', '-')}_inp{args.n_inp}{postfix}.pt"
+print(f"Saving embedding of shape {embedding.shape} to {save_path}")
+torch.save(embedding, save_path)

tests/assets/flash_bert.pt

@@ -23,6 +23,6 @@ async def test_single_query(default_model):
     response = requests.post(url, json=data, headers=headers)
 
     embedding = torch.Tensor(json.loads(response.text))
-    # reference_embedding = torch.load("assets/default_model.pt")
+    reference_embedding = torch.load("./tests/assets/sentence-transformers-all-MiniLM-L6-v2_inp1_no_flash.pt")
 
-    # assert torch.allclose(embedding, reference_embedding)
+    assert torch.allclose(embedding, reference_embedding, atol=1e-3, rtol=1e-3)

tests/assets/sentence-transformers-all-MiniLM-L6-v2_inp1.pt

@@ -23,6 +23,6 @@ async def test_single_query(default_model):
     response = requests.post(url, json=data, headers=headers)
 
     embedding = torch.Tensor(json.loads(response.text))
-    # reference_embedding = torch.load("assets/default_model.pt")
+    reference_embedding = torch.load("./tests/assets/sentence-transformers-all-MiniLM-L6-v2_inp1.pt")
 
-    # assert torch.allclose(embedding, reference_embedding)
+    assert torch.allclose(embedding, reference_embedding, atol=1e-3, rtol=1e-3)