embedkit 0.1.0__tar.gz → 0.1.2__tar.gz

{embedkit-0.1.0 → embedkit-0.1.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: embedkit
-Version: 0.1.0
+Version: 0.1.2
 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.2/main.py

@@ -0,0 +1,97 @@
+# ./main.py
+from embedkit import EmbedKit
+from embedkit.models import Model
+from embedkit.providers.cohere import CohereInputType
+from pathlib import Path
+import os
+
+
+def get_online_image(url: str) -> Path:
+    """Download an image from a URL and return its local path."""
+    import requests
+    from tempfile import NamedTemporaryFile
+
+    # Add User-Agent header to comply with Wikipedia's policy
+    headers = {"User-Agent": "EmbedKit-Example/1.0"}
+
+    response = requests.get(url, headers=headers)
+    response.raise_for_status()
+
+    temp_file = NamedTemporaryFile(delete=False, suffix=".png")
+    temp_file.write(response.content)
+    temp_file.close()
+
+    return Path(temp_file.name)
+
+
+def get_sample_image() -> Path:
+    """Get a sample image for testing."""
+    url = "https://upload.wikimedia.org/wikipedia/commons/b/b8/English_Wikipedia_HomePage_2001-12-20.png"
+    return get_online_image(url)
+
+
+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, text_batch_size=16, image_batch_size=8)
+
+results = kit.embed_text("Hello world")
+assert results.shape[0] == 1
+assert len(results.shape) == 3
+
+results = kit.embed_image(sample_image)
+assert results.shape[0] == 1
+assert len(results.shape) == 3
+assert len(results.source_images_b64) > 0
+
+results = kit.embed_pdf(sample_pdf)
+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,
+)
+
+results = kit.embed_text("Hello world")
+assert results.shape[0] == 1
+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,
+)
+
+results = kit.embed_text("Hello world")
+assert results.shape[0] == 1
+assert len(results.shape) == 2
+
+results = kit.embed_image(sample_image)
+assert results.shape[0] == 1
+assert len(results.shape) == 2
+assert len(results.source_images_b64) > 0
+
+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.0 → embedkit-0.1.2}/pyproject.toml

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

{embedkit-0.1.0 → embedkit-0.1.2}/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.0 → embedkit-0.1.2}/src/embedkit/base.py

@@ -2,7 +2,7 @@
 """Base classes for EmbedKit."""
 
 from abc import ABC, abstractmethod
-from typing import Union, List
+from typing import Union, List, Optional
 from pathlib import Path
 import numpy as np
 from dataclasses import dataclass
@@ -14,6 +14,7 @@ class EmbeddingResult:
     model_name: str
     model_provider: str
     input_type: str
+    source_images_b64: Optional[List[str]] = None
 
     @property
     def shape(self) -> tuple:
@@ -36,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.0 → embedkit-0.1.2}/src/embedkit/providers/cohere.py

@@ -6,7 +6,7 @@ from pathlib import Path
 import numpy as np
 from enum import Enum
 
-from ..utils import pdf_to_images
+from ..utils import pdf_to_images, image_to_base64
 from ..base import EmbeddingProvider, EmbeddingError, EmbeddingResult
 
 
@@ -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,60 +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:
-            import base64
-
-            b64_images = []
-            for image in images:
-                if isinstance(image, (Path, str)):
-                    try:
-                        base64_only = base64.b64encode(Path(image).read_bytes()).decode(
-                            "utf-8"
-                        )
-                    except Exception as e:
-                        raise EmbeddingError(
-                            f"Failed to read image {image}: {e}"
-                        ) from e
-
-                    if isinstance(image, Path):
-                        image = str(image)
-
-                    if image.lower().endswith(".png"):
-                        content_type = "image/png"
-                    elif image.lower().endswith((".jpg", ".jpeg")):
-                        content_type = "image/jpeg"
-                    elif image.lower().endswith(".gif"):
-                        content_type = "image/gif"
-                    else:
-                        raise EmbeddingError(
-                            f"Unsupported image format for {image}; expected .png, .jpg, .jpeg, or .gif"
-                        )
-                    base64_image = f"data:{content_type};base64,{base64_only}"
-                else:
-                    raise EmbeddingError(f"Unsupported image type: {type(image)}")
-
-                b64_images.append(base64_image)
-
-            response = client.embed(
-                model=self.model_name,
-                input_type="image",
-                images=b64_images,
-                embedding_types=["float"],
-            )
+            all_embeddings = []
+            all_b64_images = []
+
+            # 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 = []
+
+                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=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.0 → embedkit-0.1.2}/src/embedkit/providers/colpali.py

@@ -8,7 +8,7 @@ import numpy as np
 import torch
 from PIL import Image
 
-from ..utils import pdf_to_images
+from ..utils import pdf_to_images, image_to_base64
 from ..base import EmbeddingProvider, EmbeddingError, EmbeddingResult
 
 logger = logging.getLogger(__name__)
@@ -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,30 +108,44 @@ class ColPaliProvider(EmbeddingProvider):
             images = [Path(img) for img in images]
 
         try:
-            pil_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 = []
+
+            for i in range(0, len(images), self.image_batch_size):
+                batch_images = images[i : i + self.image_batch_size]
+                pil_images = []
+                b64_images = []
 
-                with Image.open(img_path) as img:
-                    pil_images.append(img.convert("RGB"))
+                for img_path in batch_images:
+                    if not img_path.exists():
+                        raise EmbeddingError(f"Image not found: {img_path}")
 
-            processed = self._processor.process_images(pil_images).to(self.device)
+                    with Image.open(img_path) as img:
+                        pil_images.append(img.convert("RGB"))
+                    b64_images.append(image_to_base64(img_path))
 
-            with torch.no_grad():
-                embeddings = self._model(**processed)
+                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)
+
+            # 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=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.2/src/embedkit/utils.py

@@ -0,0 +1,48 @@
+from pdf2image import convert_from_path
+from pathlib import Path
+from .config import get_temp_dir
+from typing import Union
+
+
+def pdf_to_images(pdf_path: Path) -> list[Path]:
+    """Convert a PDF file to a list of images."""
+    root_temp_dir = get_temp_dir()
+    img_temp_dir = root_temp_dir / "images"
+    img_temp_dir.mkdir(parents=True, exist_ok=True)
+    images = convert_from_path(pdf_path=str(pdf_path), output_folder=str(img_temp_dir))
+    image_paths = []
+
+    for i, image in enumerate(images):
+        output_path = img_temp_dir / f"{pdf_path.stem}_{i}.png"
+        if output_path.exists():
+            output_path.unlink()
+
+        image.save(output_path)
+        image_paths.append(output_path)
+    return image_paths
+
+
+def image_to_base64(image_path: Union[str, Path]):
+    import base64
+
+    try:
+        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
+
+    if isinstance(image_path, Path):
+        image_path_str = str(image_path)
+
+    if image_path_str.lower().endswith(".png"):
+        content_type = "image/png"
+    elif image_path_str.lower().endswith((".jpg", ".jpeg")):
+        content_type = "image/jpeg"
+    elif image_path_str.lower().endswith(".gif"):
+        content_type = "image/gif"
+    else:
+        raise ValueError(
+            f"Unsupported image format for {image_path}; expected .png, .jpg, .jpeg, or .gif"
+        )
+    base64_image = f"data:{content_type};base64,{base64_only}"
+
+    return base64_image

embedkit-0.1.0/main.py

@@ -1,78 +0,0 @@
-# ./main.py
-from embedkit import EmbedKit
-from embedkit.models import Model
-from embedkit.providers.cohere import CohereInputType
-from pathlib import Path
-import os
-
-
-def get_online_image(url: str) -> Path:
-    """Download an image from a URL and return its local path."""
-    import requests
-    from tempfile import NamedTemporaryFile
-
-    # Add User-Agent header to comply with Wikipedia's policy
-    headers = {"User-Agent": "EmbedKit-Example/1.0"}
-
-    response = requests.get(url, headers=headers)
-    response.raise_for_status()
-
-    temp_file = NamedTemporaryFile(delete=False, suffix=".png")
-    temp_file.write(response.content)
-    temp_file.close()
-
-    return Path(temp_file.name)
-
-
-def get_sample_image() -> Path:
-    """Get a sample image for testing."""
-    url = "https://upload.wikimedia.org/wikipedia/commons/b/b8/English_Wikipedia_HomePage_2001-12-20.png"
-    return get_online_image(url)
-
-
-sample_image = get_sample_image()
-
-sample_pdf = Path("tests/fixtures/2407.01449v6_p1.pdf")
-
-kit = EmbedKit.colpali(model=Model.ColPali.V1_3)
-
-embeddings = kit.embed_text("Hello world")
-assert embeddings.shape[0] == 1
-assert len(embeddings.shape) == 3
-
-embeddings = kit.embed_image(sample_image)
-assert embeddings.shape[0] == 1
-assert len(embeddings.shape) == 3
-
-embeddings = kit.embed_pdf(sample_pdf)
-assert embeddings.shape[0] == 1
-assert len(embeddings.shape) == 3
-
-
-kit = EmbedKit.cohere(
-    model=Model.Cohere.EMBED_V4_0,
-    api_key=os.getenv("COHERE_API_KEY"),
-    text_input_type=CohereInputType.SEARCH_QUERY,
-)
-
-embeddings = kit.embed_text("Hello world")
-assert embeddings.shape[0] == 1
-assert len(embeddings.shape) == 2
-
-kit = EmbedKit.cohere(
-    model=Model.Cohere.EMBED_V4_0,
-    api_key=os.getenv("COHERE_API_KEY"),
-    text_input_type=CohereInputType.SEARCH_DOCUMENT,
-)
-
-embeddings = kit.embed_text("Hello world")
-assert embeddings.shape[0] == 1
-assert len(embeddings.shape) == 2
-
-embeddings = kit.embed_image(sample_image)
-assert embeddings.shape[0] == 1
-assert len(embeddings.shape) == 2
-
-embeddings = kit.embed_pdf(sample_pdf)
-assert embeddings.shape[0] == 1
-assert len(embeddings.shape) == 2

embedkit-0.1.0/src/embedkit/utils.py

@@ -1,21 +0,0 @@
-from pdf2image import convert_from_path
-from pathlib import Path
-from .config import get_temp_dir
-
-
-def pdf_to_images(pdf_path: Path) -> list[Path]:
-    """Convert a PDF file to a list of images."""
-    root_temp_dir = get_temp_dir()
-    img_temp_dir = root_temp_dir / "images"
-    img_temp_dir.mkdir(parents=True, exist_ok=True)
-    images = convert_from_path(pdf_path=str(pdf_path), output_folder=str(img_temp_dir))
-    image_paths = []
-
-    for i, image in enumerate(images):
-        output_path = img_temp_dir / f"{pdf_path.stem}_{i}.png"
-        if output_path.exists():
-            output_path.unlink()
-
-        image.save(output_path)
-        image_paths.append(output_path)
-    return image_paths