update embd batch size

Value4AI · Sep 9, 2024 · 887107a · 887107a
1 parent a4a6fe4
commit 887107a
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/gpv/embd.py b/gpv/embd.py
@@ -13,7 +13,7 @@ def __init__(self, model_name_or_path: str='Alibaba-NLP/gte-multilingual-base',
         self.tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
 
 
-    def get_embedding(self, input_texts: list[str], dimension: int = 768, batch_size: int = 16) -> torch.Tensor:
+    def get_embedding(self, input_texts: list[str], dimension: int = 768, batch_size: int = 8) -> torch.Tensor:
         """
         Get the sentence embeddings of the input texts.