Infinity improvements (#26)

- Uses v2 command to launch Infinity
- Added embedding model "BAAI/bge-small-en-v1.5"
- Added reranking model "mixedbread-ai/mxbai-rerank-xsmall-v1"

Author: michaelfeil

.env.example

@@ -27,4 +27,4 @@ EMBEDDING_MODEL=sentence-transformers/all-MiniLM-L6-v2
 RERANKER_MODEL=cross-encoder/ms-marco-TinyBERT-L-2
 OWL_CONCURRENT_ROWS_BATCH_SIZE=5
 OWL_CONCURRENT_COLS_BATCH_SIZE=5
-OWL_MAX_WRITE_BATCH_SIZE=1000
+OWL_MAX_WRITE_BATCH_SIZE=1000

docker/compose.cpu.yml

@@ -1,14 +1,17 @@
 services:
   infinity:
-    image: michaelf34/infinity:0.0.55
-    entrypoint:
-      [
-        "/bin/sh",
-        "-c",
-        "(. /app/.venv/bin/activate && infinity_emb v2 --port 6909 --model-id $${EMBEDDING_MODEL} --model-warmup --device cpu &);(. /app/.venv/bin/activate && infinity_emb v2 --port 6919 --model-id $${RERANKER_MODEL} --model-warmup --device cpu )",
-      ]
+    image: michaelf34/infinity:0.0.70-cpu
+    container_name: jamai_infinity
+    command: [
+        "v2",
+        "--engine", "torch",
+        "--port", "6909",
+        "--model-warmup",
+        "--model-id", "${EMBEDDING_MODEL}", 
+        "--model-id", "${RERANKER_MODEL}",
+    ]
     healthcheck:
-      test: ["CMD-SHELL", "curl --fail http://localhost:6909/health && curl --fail http://localhost:6919/health || exit 1"]
+      test: ["CMD-SHELL", "curl --fail http://localhost:6909/health"]
       interval: 10s
       timeout: 2s
       retries: 20

docker/compose.nvidia.yml

@@ -1,4 +1,4 @@
 include:
   - path:
       - compose.cpu.yml
-      - nvidia.yml
+      - nvidia.yml

docker/nvidia.yml

@@ -1,11 +1,6 @@
 services:
   infinity:
-    entrypoint:
-      [
-        "/bin/sh",
-        "-c",
-        "(. /app/.venv/bin/activate && infinity_emb v2 --port 6909 --model-id $${EMBEDDING_MODEL} --model-warmup --device cuda &);(. /app/.venv/bin/activate && infinity_emb v2 --port 6919 --model-id $${RERANKER_MODEL} --model-warmup --device cuda )",
-      ]
+    image: michaelf34/infinity:0.0.70  # Use GPU-compatible image, instead of torch cpu-only.
     deploy:
       resources:
         reservations:

services/api/src/owl/configs/models.json

@@ -59,6 +59,21 @@
         }
       ]
     },
+    {
+      "id": "ellm/BAAI/bge-small-en-v1.5",
+      "litellm_id": "openai/BAAI/bge-small-en-v1.5",
+      "context_length": 512,
+      "embedding_size": 1024,
+      "languages": ["mul"],
+      "capabilities": ["embed"],
+      "deployments": [
+        {
+          "litellm_id": "openai/BAAI/bge-small-en-v1.5",
+          "api_base": "http://infinity:6909",
+          "provider": "ellm"
+        }
+      ]
+    },
     {
       "id": "openai/text-embedding-3-large-3072",
       "name": "OpenAI Text Embedding 3 Large (3072-dim)",
@@ -123,6 +138,20 @@
     }
   ],
   "rerank_models": [
+    {
+      "id": "ellm/mixedbread-ai/mxbai-rerank-xsmall-v1",
+      "name": "ELLM mxbai-rerank-xsmall v1",
+      "context_length": 512,
+      "languages": ["en"],
+      "capabilities": ["rerank"],
+      "deployments": [
+        {
+          "litellm_id": "",
+          "api_base": "http://infinity:6919",
+          "provider": "ellm"
+        }
+      ]
+    },
     {
       "id": "ellm/cross-encoder/ms-marco-TinyBERT-L-2",
       "name": "ELLM TinyBERT L2",
@@ -152,4 +181,4 @@
       ]
     }
   ]
-}
+}