embedkit 0.1.1__tar.gz → 0.1.3__tar.gz

{embedkit-0.1.1 → embedkit-0.1.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: embedkit
-Version: 0.1.1
+Version: 0.1.3
 Summary: A simple toolkit for generating vector embeddings across multiple providers and models
 Author-email: JP Hwang <me@jphwang.com>
 License: MIT

{embedkit-0.1.1 → embedkit-0.1.3}/main.py

@@ -1,7 +1,6 @@
 # ./main.py
 from embedkit import EmbedKit
-from embedkit.models import Model
-from embedkit.providers.cohere import CohereInputType
+from embedkit.classes import Model, CohereInputType
 from pathlib import Path
 import os
 
@@ -33,8 +32,9 @@ def get_sample_image() -> Path:
 sample_image = get_sample_image()
 
 sample_pdf = Path("tests/fixtures/2407.01449v6_p1.pdf")
+long_pdf = Path("tmp/2407.01449v6.pdf")
 
-kit = EmbedKit.colpali(model=Model.ColPali.V1_3)
+kit = EmbedKit.colpali(model=Model.ColPali.V1_3, text_batch_size=16, image_batch_size=8)
 
 results = kit.embed_text("Hello world")
 assert results.shape[0] == 1
@@ -50,10 +50,17 @@ assert results.shape[0] == 1
 assert len(results.shape) == 3
 assert len(results.source_images_b64) > 0
 
+results = kit.embed_pdf(long_pdf)
+assert results.shape[0] == 26
+assert len(results.shape) == 3
+assert len(results.source_images_b64) > 0
+
 
 kit = EmbedKit.cohere(
     model=Model.Cohere.EMBED_V4_0,
     api_key=os.getenv("COHERE_API_KEY"),
+    text_batch_size=64,
+    image_batch_size=8,
     text_input_type=CohereInputType.SEARCH_QUERY,
 )
 
@@ -64,6 +71,8 @@ assert len(results.shape) == 2
 kit = EmbedKit.cohere(
     model=Model.Cohere.EMBED_V4_0,
     api_key=os.getenv("COHERE_API_KEY"),
+    text_batch_size=64,
+    image_batch_size=8,
     text_input_type=CohereInputType.SEARCH_DOCUMENT,
 )
 
@@ -80,3 +89,8 @@ results = kit.embed_pdf(sample_pdf)
 assert results.shape[0] == 1
 assert len(results.shape) == 2
 assert len(results.source_images_b64) > 0
+
+results = kit.embed_pdf(long_pdf)
+assert results.shape[0] == 26
+assert len(results.shape) == 2
+assert len(results.source_images_b64) > 0

{embedkit-0.1.1 → embedkit-0.1.3}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "embedkit"
-version = "0.1.1"
+version = "0.1.3"
 description = "A simple toolkit for generating vector embeddings across multiple providers and models"
 readme = "README.md"
 requires-python = ">=3.10"

{embedkit-0.1.1 → embedkit-0.1.3}/src/embedkit/__init__.py

@@ -26,21 +26,33 @@ class EmbedKit:
         self._provider = provider_instance
 
     @classmethod
-    def colpali(cls, model: Model = Model.ColPali.V1_3, device: Optional[str] = None):
+    def colpali(
+        cls,
+        model: Model = Model.ColPali.V1_3,
+        device: Optional[str] = None,
+        text_batch_size: int = 32,
+        image_batch_size: int = 8,
+    ):
         """
         Create EmbedKit instance with ColPali provider.
 
         Args:
             model: ColPali model enum
             device: Device to run on ('cuda', 'mps', 'cpu', or None for auto-detect)
+            text_batch_size: Batch size for text embedding generation
+            image_batch_size: Batch size for image embedding generation
         """
         if model == Model.ColPali.V1_3:
             model_name = "vidore/colpali-v1.3"
         else:
             raise ValueError(f"Unsupported model: {model}")
 
-
-        provider = ColPaliProvider(model_name=model_name, device=device)
+        provider = ColPaliProvider(
+            model_name=model_name,
+            device=device,
+            text_batch_size=text_batch_size,
+            image_batch_size=image_batch_size,
+        )
         return cls(provider)
 
     @classmethod
@@ -48,6 +60,8 @@ class EmbedKit:
         cls,
         api_key: str,
         model: Model = Model.Cohere.EMBED_V4_0,
+        text_batch_size: int = 32,
+        image_batch_size: int = 8,
         text_input_type: CohereInputType = CohereInputType.SEARCH_DOCUMENT,
     ):
         """
@@ -56,6 +70,8 @@ class EmbedKit:
         Args:
             api_key: Cohere API key
             model: Cohere model enum
+            text_batch_size: Batch size for text embedding generation
+            image_batch_size: Batch size for image embedding generation
             input_type: Type of input for embedding (search_document or search_query)
         """
         if not api_key:
@@ -67,7 +83,10 @@ class EmbedKit:
             raise ValueError(f"Unsupported model: {model}")
 
         provider = CohereProvider(
-            api_key=api_key, model_name=model_name, text_input_type=text_input_type
+            api_key=api_key, model_name=model_name,
+            text_batch_size=48,
+            image_batch_size=8,
+            text_input_type=text_input_type
         )
         return cls(provider)
 

{embedkit-0.1.1 → embedkit-0.1.3}/src/embedkit/base.py

@@ -37,9 +37,7 @@ class EmbeddingProvider(ABC):
         pass
 
     @abstractmethod
-    def embed_pdf(
-        self, pdf: Union[Path, str]
-    ) -> EmbeddingResult:
+    def embed_pdf(self, pdf: Union[Path, str]) -> EmbeddingResult:
         """Generate image embeddings from PDFsusing the configured provider. Takes a single PDF file."""
         pass
 

embedkit-0.1.3/src/embedkit/classes.py

@@ -0,0 +1,21 @@
+# ./src/embedkit/classes.py
+
+"""Core types and enums for the EmbedKit library.
+
+This module provides the main types and enums that users should interact with:
+- EmbeddingResult: The result type returned by embedding operations
+- EmbeddingError: Exception type for embedding operations
+- Model: Enum of supported embedding models
+- CohereInputType: Enum for Cohere's input types
+"""
+
+from . import EmbeddingResult, EmbeddingError
+from .models import Model
+from .providers.cohere import CohereInputType
+
+__all__ = [
+    "EmbeddingResult",
+    "EmbeddingError",
+    "Model",
+    "CohereInputType"
+]

{embedkit-0.1.1 → embedkit-0.1.3}/src/embedkit/providers/cohere.py

@@ -24,10 +24,14 @@ class CohereProvider(EmbeddingProvider):
         self,
         api_key: str,
         model_name: str,
+        text_batch_size: int,
+        image_batch_size: int,
         text_input_type: CohereInputType = CohereInputType.SEARCH_DOCUMENT,
     ):
         self.api_key = api_key
         self.model_name = model_name
+        self.text_batch_size = text_batch_size
+        self.image_batch_size = image_batch_size
         self.input_type = text_input_type
         self._client = None
         self.provider_name = "Cohere"
@@ -55,15 +59,21 @@ class CohereProvider(EmbeddingProvider):
             texts = [texts]
 
         try:
-            response = client.embed(
-                texts=texts,
-                model=self.model_name,
-                input_type=self.input_type.value,
-                embedding_types=["float"],
-            )
+            all_embeddings = []
+
+            # Process texts in batches
+            for i in range(0, len(texts), self.text_batch_size):
+                batch_texts = texts[i : i + self.text_batch_size]
+                response = client.embed(
+                    texts=batch_texts,
+                    model=self.model_name,
+                    input_type=self.input_type.value,
+                    embedding_types=["float"],
+                )
+                all_embeddings.extend(response.embeddings.float_)
 
             return EmbeddingResult(
-                embeddings=np.array(response.embeddings.float_),
+                embeddings=np.array(all_embeddings),
                 model_name=self.model_name,
                 model_provider=self.provider_name,
                 input_type=self.input_type.value,
@@ -81,34 +91,45 @@ class CohereProvider(EmbeddingProvider):
         input_type = "image"
 
         if isinstance(images, (str, Path)):
-            images = [images]
+            images = [Path(images)]
+        else:
+            images = [Path(img) for img in images]
 
         try:
-            b64_images = []
-            for image in images:
-                b64_image = image_to_base64(image)
+            all_embeddings = []
+            all_b64_images = []
 
-            b64_images.append(b64_image)
+            # Process images in batches
+            for i in range(0, len(images), self.image_batch_size):
+                batch_images = images[i : i + self.image_batch_size]
+                b64_images = []
 
-            response = client.embed(
-                model=self.model_name,
-                input_type="image",
-                images=b64_images,
-                embedding_types=["float"],
-            )
+                for image in batch_images:
+                    if not image.exists():
+                        raise EmbeddingError(f"Image not found: {image}")
+                    b64_images.append(image_to_base64(image))
+
+                response = client.embed(
+                    model=self.model_name,
+                    input_type="image",
+                    images=b64_images,
+                    embedding_types=["float"],
+                )
+
+                all_embeddings.extend(response.embeddings.float_)
+                all_b64_images.extend(b64_images)
 
             return EmbeddingResult(
-                embeddings=np.array(response.embeddings.float_),
+                embeddings=np.array(all_embeddings),
                 model_name=self.model_name,
                 model_provider=self.provider_name,
                 input_type=input_type,
-                source_images_b64=b64_images,
+                source_images_b64=all_b64_images,
             )
 
         except Exception as e:
             raise EmbeddingError(f"Failed to embed image with Cohere: {e}") from e
 
-
     def embed_pdf(self, pdf_path: Path) -> EmbeddingResult:
         """Generate embeddings for a PDF file using Cohere API."""
         image_paths = pdf_to_images(pdf_path)

{embedkit-0.1.1 → embedkit-0.1.3}/src/embedkit/providers/colpali.py

@@ -17,9 +17,17 @@ logger = logging.getLogger(__name__)
 class ColPaliProvider(EmbeddingProvider):
     """ColPali embedding provider for document understanding."""
 
-    def __init__(self, model_name: str, device: Optional[str] = None):
+    def __init__(
+        self,
+        model_name: str,
+        text_batch_size: int,
+        image_batch_size: int,
+        device: Optional[str] = None,
+    ):
         self.model_name = model_name
         self.provider_name = "ColPali"
+        self.text_batch_size = text_batch_size
+        self.image_batch_size = image_batch_size
 
         # Auto-detect device
         if device is None:
@@ -64,13 +72,22 @@ class ColPaliProvider(EmbeddingProvider):
             texts = [texts]
 
         try:
-            processed = self._processor.process_queries(texts).to(self.device)
+            # Process texts in batches
+            all_embeddings = []
 
-            with torch.no_grad():
-                embeddings = self._model(**processed)
+            for i in range(0, len(texts), self.text_batch_size):
+                batch_texts = texts[i : i + self.text_batch_size]
+                processed = self._processor.process_queries(batch_texts).to(self.device)
+
+                with torch.no_grad():
+                    batch_embeddings = self._model(**processed)
+                    all_embeddings.append(batch_embeddings.cpu().float().numpy())
+
+            # Concatenate all batch embeddings
+            final_embeddings = np.concatenate(all_embeddings, axis=0)
 
             return EmbeddingResult(
-                embeddings=embeddings.cpu().float().numpy(),
+                embeddings=final_embeddings,
                 model_name=self.model_name,
                 model_provider=self.provider_name,
                 input_type="text",
@@ -91,38 +108,44 @@ class ColPaliProvider(EmbeddingProvider):
             images = [Path(img) for img in images]
 
         try:
-            pil_images = []
-            b64_images = []
-            for img_path in images:
-                if not img_path.exists():
-                    raise EmbeddingError(f"Image not found: {img_path}")
+            # Process images in batches
+            all_embeddings = []
+            all_b64_images = []
 
-                with Image.open(img_path) as img:
-                    pil_images.append(img.convert("RGB"))
+            for i in range(0, len(images), self.image_batch_size):
+                batch_images = images[i : i + self.image_batch_size]
+                pil_images = []
+                b64_images = []
 
-                for image in images:
-                    b64_image = image_to_base64(image)
+                for img_path in batch_images:
+                    if not img_path.exists():
+                        raise EmbeddingError(f"Image not found: {img_path}")
 
-                b64_images.append(b64_image)
+                    with Image.open(img_path) as img:
+                        pil_images.append(img.convert("RGB"))
+                    b64_images.append(image_to_base64(img_path))
 
-            processed = self._processor.process_images(pil_images).to(self.device)
+                processed = self._processor.process_images(pil_images).to(self.device)
 
+                with torch.no_grad():
+                    batch_embeddings = self._model(**processed)
+                    all_embeddings.append(batch_embeddings.cpu().float().numpy())
+                    all_b64_images.extend(b64_images)
 
-            with torch.no_grad():
-                embeddings = self._model(**processed)
+            # Concatenate all batch embeddings
+            final_embeddings = np.concatenate(all_embeddings, axis=0)
 
             return EmbeddingResult(
-                embeddings=embeddings.cpu().float().numpy(),
+                embeddings=final_embeddings,
                 model_name=self.model_name,
                 model_provider=self.provider_name,
                 input_type="image",
-                source_images_b64=b64_images,
+                source_images_b64=all_b64_images,
             )
 
         except Exception as e:
             raise EmbeddingError(f"Failed to embed images: {e}") from e
 
-
     def embed_pdf(self, pdf_path: Path) -> EmbeddingResult:
         """Generate embeddings for a PDF file using ColPali API."""
         images = pdf_to_images(pdf_path)

{embedkit-0.1.1 → embedkit-0.1.3}/src/embedkit/utils.py

@@ -26,13 +26,9 @@ def image_to_base64(image_path: Union[str, Path]):
     import base64
 
     try:
-        base64_only = base64.b64encode(Path(image_path).read_bytes()).decode(
-            "utf-8"
-        )
+        base64_only = base64.b64encode(Path(image_path).read_bytes()).decode("utf-8")
     except Exception as e:
-        raise ValueError(
-            f"Failed to read image {image_path}: {e}"
-        ) from e
+        raise ValueError(f"Failed to read image {image_path}: {e}") from e
 
     if isinstance(image_path, Path):
         image_path_str = str(image_path)