Update README.md

README.md

@@ -19,34 +19,41 @@ Benchmarking on an Nvidia-L4 instance. Note: CPU uses bert-small, CUDA uses Bert
 
 ```python
 from embed import BatchedInference
+from concurrent.futures import Future
 
 # Run any model
-register = BatchedInference(model_id=[
-  # sentence-embeddings
-  "michaelfeil/bge-small-en-v1.5",
-  # sentence-embeddings and image-embeddings
-  "jinaai/jina-clip-v1",
-  # classification models
-  "philschmid/tiny-bert-sst2-distilled",
-  # rerankers
-  "mixedbread-ai/mxbai-rerank-xsmall-v1"
-],
-# engine to `torch` or `optimum`
-engine="torch",
-# device `cuda` (Nvidia/AMD) or `cpu`
-device="cpu"
+register = BatchedInference(
+    model_id=[
+        # sentence-embeddings
+        "BAAI/bge-small-en-v1.5",
+        # sentence-embeddings and image-embeddings
+        "jinaai/jina-clip-v1",
+        # classification models
+        "philschmid/tiny-bert-sst2-distilled",
+        # rerankers
+        "mixedbread-ai/mxbai-rerank-xsmall-v1",
+    ],
+    # engine to `torch` or `optimum`
+    engine="torch",
+    # device `cuda` (Nvidia/AMD) or `cpu`
+    device="cpu",
 )
 
-sentences = ["Paris is in France.", "Berlin is in Germany.", "I love SF"]
+sentences = ["Paris is in France.", "Berlin is in Germany.", "A image of two cats."]
 images = ["http://images.cocodataset.org/val2017/000000039769.jpg"]
 question = "Where is Paris?"
 
-register.embed(sentences=sentences, model_id="michaelfeil/bge-small-en-v1.5")
-register.rerank(query=question, docs=sentences, model_id="mixedbread-ai/mxbai-rerank-xsmall-v1")
+future: "Future" = register.embed(
+    sentences=sentences, model_id="michaelfeil/bge-small-en-v1.5"
+)
+future.result()
+register.rerank(
+    query=question, docs=sentences, model_id="mixedbread-ai/mxbai-rerank-xsmall-v1"
+)
 register.classify(model_id="philschmid/tiny-bert-sst2-distilled", sentences=sentences)
 register.image_embed(model_id="jinaai/jina-clip-v1", images=images)
 
-# Always manually stop the register upon termination to free model memory.
+# manually stop the register upon termination to free model memory.
 register.stop()
 ```