Optimize batch size for models (#46)

lanterndata · Dec 5, 2023 · 084b28e · 084b28e
1 parent 8b0858e
commit 084b28e
Show file tree

Hide file tree

Showing 4 changed files with 13 additions and 13 deletions.
diff --git a/.github/workflows/publish-cli-docker.yaml b/.github/workflows/publish-cli-docker.yaml
@@ -6,7 +6,7 @@ on:
         type: string
         description: "CLI version"
         required: true
-        default: "0.0.3"
+        default: "0.0.15"
       IMAGE_NAME:
         type: string
         description: "Container image name to tag"

diff --git a/lantern_cli/Cargo.toml b/lantern_cli/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "lantern_cli"
-version = "0.0.14"
+version = "0.0.15"
 edition = "2021"
 
 [[bin]]

diff --git a/lantern_embeddings/Cargo.toml b/lantern_embeddings/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "lantern_embeddings"
-version = "0.0.3"
+version = "0.0.4"
 edition = "2021"
 
 [lib]

diff --git a/lantern_embeddings/src/lib.rs b/lantern_embeddings/src/lib.rs
@@ -395,19 +395,19 @@ fn csv_exporter_worker(
 
 fn get_default_batch_size(model: &str) -> usize {
     match model {
-        "clip/ViT-B-32-textual" => 500,
+        "clip/ViT-B-32-textual" => 2000,
         "clip/ViT-B-32-visual" => 100,
-        "BAAI/bge-small-en" => 500,
+        "BAAI/bge-small-en" => 300,
         "BAAI/bge-base-en" => 100,
-        "BAAI/bge-large-en" => 40,
-        "infloat/e5-base-v2" => 100,
-        "infloat/e5-large-v2" => 40,
+        "BAAI/bge-large-en" => 60,
+        "intfloat/e5-base-v2" => 300,
+        "intfloat/e5-large-v2" => 100,
         "llmrails/ember-v1" => 100,
-        "thenlper/gte-base" => 100,
-        "thenlper/gte-large" => 40,
-        "microsoft/all-MiniLM-L12-v2" => 500,
-        "microsoft/all-mpnet-base-v2" => 100,
-        "transformers/multi-qa-mpnet-base-dot-v1" => 50,
+        "thenlper/gte-base" => 1000,
+        "thenlper/gte-large" => 800,
+        "microsoft/all-MiniLM-L12-v2" => 1000,
+        "microsoft/all-mpnet-base-v2" => 400,
+        "transformers/multi-qa-mpnet-base-dot-v1" => 300,
         _ => 100,
     }
 }