Config batch size for embedder

embeder.py

@@ -6,6 +6,7 @@ import torch.nn.functional as F
 from transformers import AutoImageProcessor, AutoModel, AutoTokenizer
 from PIL import Image
 from transformers.utils import ModelOutput
+from tqdm import tqdm
 
 
 class MultimodalEmbedder:
@@ -13,7 +14,8 @@ class MultimodalEmbedder:
     def __init__(
             self,
             text_model: str = 'nomic-ai/nomic-embed-text-v1.5',
-            image_model: str = 'nomic-ai/nomic-embed-vision-v1.5'
+            image_model: str = 'nomic-ai/nomic-embed-vision-v1.5',
+            batch_size: int = 64
     ):
         self.tokenizer = AutoTokenizer.from_pretrained(text_model)
         self.text_model = AutoModel.from_pretrained(text_model, trust_remote_code=True)
@@ -24,6 +26,8 @@ class MultimodalEmbedder:
         self.image_model = AutoModel.from_pretrained(image_model, trust_remote_code=True)
         self.image_embedding_size = self.image_model.config.hidden_size
 
+        self.batch_size = batch_size
+
     def embed_texts(
             self,
             texts: list[str],
@@ -32,29 +36,47 @@ class MultimodalEmbedder:
     ) -> list[list[float]]:
         """Embed a list of texts"""
         texts = [f'search_query: {text}' if kind == 'query' else f'search_document: {text}' for text in texts]
-        inputs = self.tokenizer(texts, padding=True, truncation=True, return_tensors='pt')
+        
+        all_embeddings = []
+        for start in tqdm(range(0, len(texts), self.batch_size), desc='Embed texts'):
+            batch_texts = texts[start:start + self.batch_size]
+
+            inputs = self.tokenizer(
+                batch_texts,
+                padding=True,
+                truncation=True,
+                return_tensors='pt'
+            ).to(device)
 
-        with torch.no_grad():
-            outputs = self.text_model.to(device)(**inputs.to(device), batch_size=64)
+            with torch.no_grad():
+                outputs = self.text_model(**inputs)
 
-        embeddings = mean_pooling(outputs, inputs['attention_mask'])
-        embeddings = F.layer_norm(embeddings, normalized_shape=(embeddings.shape[1],))
-        embeddings = F.normalize(embeddings, p=2, dim=1)
+            embeddings = mean_pooling(outputs, inputs['attention_mask'])
+            embeddings = F.layer_norm(embeddings, normalized_shape=(embeddings.shape[1],))
+            embeddings = F.normalize(embeddings, p=2, dim=1)
+            all_embeddings.append(embeddings.cpu())
 
-        return embeddings.cpu().tolist()
+        return torch.cat(all_embeddings, dim=0).tolist()
 
     def embed_images(self, images: list[str | Image.Image], device: str = 'cpu') -> list[list[float]]:
         """Embed a list of images, which can be file paths or PIL Image objects."""
         images = [Image.open(img) if isinstance(img, str) else img for img in images]
         images = [img.convert('RGB') for img in images]
 
-        inputs = self.processor(images, return_tensors='pt')
+        all_embeddings = []
+        for start in tqdm(range(0, len(images), self.batch_size), desc='Embed images'):
+            batch_images = images[start:start + self.batch_size]
+
+            inputs = self.processor(batch_images, return_tensors='pt').to(device)
 
-        embeddings = self.image_model.to(device)(**inputs.to(device), batch_size=64).last_hidden_state
+            with torch.no_grad():
+                outputs = self.image_model(**inputs)
 
-        embeddings = F.normalize(embeddings[:, 0], p=2, dim=1)
+            embeddings = outputs.last_hidden_state[:, 0]  # CLS token
+            embeddings = F.normalize(embeddings, p=2, dim=1)
+            all_embeddings.append(embeddings.cpu())
 
-        return embeddings.cpu().tolist()
+        return torch.cat(all_embeddings, dim=0).tolist()
 
     def similarity(
             self,