embedkit 0.1.3__py3-none-any.whl → 0.1.6__py3-none-any.whl

embedkit/__init__.py

@@ -5,10 +5,9 @@ EmbedKit: A unified toolkit for generating vector embeddings.
 
 from typing import Union, List, Optional
 from pathlib import Path
-import numpy as np
 
 from .models import Model
-from .base import EmbeddingError, EmbeddingResult
+from .base import EmbeddingError, EmbeddingResponse
 from .providers import ColPaliProvider, CohereProvider
 from .providers.cohere import CohereInputType
 
@@ -28,7 +27,7 @@ class EmbedKit:
     @classmethod
     def colpali(
         cls,
-        model: Model = Model.ColPali.V1_3,
+        model: Model = Model.ColPali.COLPALI_V1_3,
         device: Optional[str] = None,
         text_batch_size: int = 32,
         image_batch_size: int = 8,
@@ -42,13 +41,13 @@ class EmbedKit:
             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}")
+        if not isinstance(model, Model.ColPali):
+            raise ValueError(
+                f"Unsupported model: {model}. Must be a Model.ColPali enum value."
+            )
 
         provider = ColPaliProvider(
-            model_name=model_name,
+            model=model,
             device=device,
             text_batch_size=text_batch_size,
             image_batch_size=image_batch_size,
@@ -77,16 +76,15 @@ class EmbedKit:
         if not api_key:
             raise ValueError("API key is required")
 
-        if model == Model.Cohere.EMBED_V4_0:
-            model_name = "embed-v4.0"
-        else:
+        if not isinstance(model, Model.Cohere):
             raise ValueError(f"Unsupported model: {model}")
 
         provider = CohereProvider(
-            api_key=api_key, model_name=model_name,
-            text_batch_size=48,
-            image_batch_size=8,
-            text_input_type=text_input_type
+            api_key=api_key,
+            model=model,
+            text_batch_size=text_batch_size,
+            image_batch_size=image_batch_size,
+            text_input_type=text_input_type,
         )
         return cls(provider)
 
@@ -103,7 +101,7 @@ class EmbedKit:
     #     provider = HuggingFaceProvider(model_name=model_name, device=device)
     #     return cls(provider)
 
-    def embed_text(self, texts: Union[str, List[str]], **kwargs) -> EmbeddingResult:
+    def embed_text(self, texts: Union[str, List[str]], **kwargs) -> EmbeddingResponse:
         """Generate document text embeddings using the configured provider.
 
         Args:
@@ -117,11 +115,11 @@ class EmbedKit:
 
     def embed_image(
         self, images: Union[Path, str, List[Union[Path, str]]]
-    ) -> EmbeddingResult:
+    ) -> EmbeddingResponse:
         """Generate image embeddings using the configured provider."""
         return self._provider.embed_image(images)
 
-    def embed_pdf(self, pdf: Union[Path, str]) -> EmbeddingResult:
+    def embed_pdf(self, pdf: Union[Path, str]) -> EmbeddingResponse:
         """Generate image embeddings from PDFsusing the configured provider. Takes a single PDF file."""
         return self._provider.embed_pdf(pdf)
 

embedkit/base.py

@@ -7,39 +7,113 @@ from pathlib import Path
 import numpy as np
 from dataclasses import dataclass
 
+from .models import Model
+from .utils import with_pdf_cleanup
+
+
+@dataclass
+class EmbeddingObject:
+    embedding: np.ndarray
+    source_b64: str = None
+    source_content_type: str = None  # e.g., "image/png", "image/jpeg"
+
 
 @dataclass
-class EmbeddingResult:
-    embeddings: np.ndarray
+class EmbeddingResponse:
     model_name: str
     model_provider: str
     input_type: str
-    source_images_b64: Optional[List[str]] = None
+    objects: List[EmbeddingObject]
 
     @property
     def shape(self) -> tuple:
-        return self.embeddings.shape
+        return self.objects[0].embedding.shape
 
 
 class EmbeddingProvider(ABC):
     """Abstract base class for embedding providers."""
 
+    def __init__(
+        self,
+        model_name: str,
+        text_batch_size: int,
+        image_batch_size: int,
+        provider_name: str,
+    ):
+        self.model_name = model_name
+        self.provider_name = provider_name
+        self.text_batch_size = text_batch_size
+        self.image_batch_size = image_batch_size
+
+    def _normalize_text_input(self, texts: Union[str, List[str]]) -> List[str]:
+        """Normalize text input to a list of strings."""
+        if isinstance(texts, str):
+            return [texts]
+        return texts
+
+    def _normalize_image_input(
+        self, images: Union[Path, str, List[Union[Path, str]]]
+    ) -> List[Path]:
+        """Normalize image input to a list of Path objects."""
+        if isinstance(images, (str, Path)):
+            return [Path(images)]
+        return [Path(img) for img in images]
+
+    def _create_text_response(
+        self, embeddings: List[np.ndarray], input_type: str = "text"
+    ) -> EmbeddingResponse:
+        """Create a standardized text embedding response."""
+        return EmbeddingResponse(
+            model_name=self.model_name,
+            model_provider=self.provider_name,
+            input_type=input_type,
+            objects=[EmbeddingObject(embedding=e) for e in embeddings],
+        )
+
+    def _create_image_response(
+        self,
+        embeddings: List[np.ndarray],
+        b64_data: List[str],
+        content_types: List[str],
+        input_type: str = "image",
+    ) -> EmbeddingResponse:
+        """Create a standardized image embedding response."""
+        return EmbeddingResponse(
+            model_name=self.model_name,
+            model_provider=self.provider_name,
+            input_type=input_type,
+            objects=[
+                EmbeddingObject(
+                    embedding=embedding,
+                    source_b64=b64_data,
+                    source_content_type=content_type,
+                )
+                for embedding, b64_data, content_type in zip(
+                    embeddings, b64_data, content_types
+                )
+            ],
+        )
+
     @abstractmethod
-    def embed_text(self, texts: Union[str, List[str]], **kwargs) -> EmbeddingResult:
+    def embed_text(self, texts: Union[str, List[str]], **kwargs) -> EmbeddingResponse:
         """Generate document text embeddings using the configured provider."""
         pass
 
     @abstractmethod
     def embed_image(
         self, images: Union[Path, str, List[Union[Path, str]]]
-    ) -> EmbeddingResult:
+    ) -> EmbeddingResponse:
         """Generate image embeddings using the configured provider."""
         pass
 
-    @abstractmethod
-    def embed_pdf(self, pdf: Union[Path, str]) -> EmbeddingResult:
-        """Generate image embeddings from PDFsusing the configured provider. Takes a single PDF file."""
-        pass
+    def embed_pdf(self, pdf_path: Path) -> EmbeddingResponse:
+        """Generate embeddings for a PDF file."""
+        return self._embed_pdf_impl(pdf_path)
+
+    @with_pdf_cleanup
+    def _embed_pdf_impl(self, pdf_path: List[Path]) -> EmbeddingResponse:
+        """Internal implementation of PDF embedding with cleanup handled by decorator."""
+        return self.embed_image(pdf_path)
 
 
 class EmbeddingError(Exception):

embedkit/classes.py

@@ -9,13 +9,8 @@ This module provides the main types and enums that users should interact with:
 - CohereInputType: Enum for Cohere's input types
 """
 
-from . import EmbeddingResult, EmbeddingError
+from . import EmbeddingResponse, EmbeddingError
 from .models import Model
 from .providers.cohere import CohereInputType
 
-__all__ = [
-    "EmbeddingResult",
-    "EmbeddingError",
-    "Model",
-    "CohereInputType"
-]
+__all__ = ["EmbeddingResponse", "EmbeddingError", "Model", "CohereInputType"]

embedkit/models.py

@@ -6,7 +6,13 @@ from enum import Enum
 
 class Model:
     class ColPali(Enum):
-        V1_3 = "colpali-v1.3"
+        COLPALI_V1_3 = "vidore/colpali-v1.3"
+        COLSMOL_500M = "vidore/colSmol-500M"
+        COLSMOL_256M = "vidore/colSmol-256M"
 
     class Cohere(Enum):
         EMBED_V4_0 = "embed-v4.0"
+        EMBED_ENGLISH_V3_0 = "embed-english-v3.0"
+        EMBED_ENGLISH_LIGHT_V3_0 = "embed-english-light-v3.0"
+        EMBED_MULTILINGUAL_V3_0 = "embed-multilingual-v3.0"
+        EMBED_MULTILINGUAL_LIGHT_V3_0 = "embed-multilingual-light-v3.0"

embedkit/providers/cohere.py

@@ -5,9 +5,13 @@ from typing import Union, List
 from pathlib import Path
 import numpy as np
 from enum import Enum
+import logging
 
-from ..utils import pdf_to_images, image_to_base64
-from ..base import EmbeddingProvider, EmbeddingError, EmbeddingResult
+from ..models import Model
+from ..utils import image_to_base64
+from ..base import EmbeddingProvider, EmbeddingError, EmbeddingResponse
+
+logger = logging.getLogger(__name__)
 
 
 class CohereInputType(Enum):
@@ -23,18 +27,20 @@ class CohereProvider(EmbeddingProvider):
     def __init__(
         self,
         api_key: str,
-        model_name: str,
+        model: Model.Cohere,
         text_batch_size: int,
         image_batch_size: int,
         text_input_type: CohereInputType = CohereInputType.SEARCH_DOCUMENT,
     ):
+        super().__init__(
+            model_name=model.value,
+            text_batch_size=text_batch_size,
+            image_batch_size=image_batch_size,
+            provider_name="Cohere",
+        )
         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"
 
     def _get_client(self):
         """Lazy load the Cohere client."""
@@ -51,12 +57,10 @@ class CohereProvider(EmbeddingProvider):
                 raise EmbeddingError(f"Failed to initialize Cohere client: {e}") from e
         return self._client
 
-    def embed_text(self, texts: Union[str, List[str]], **kwargs) -> EmbeddingResult:
+    def embed_text(self, texts: Union[str, List[str]], **kwargs) -> EmbeddingResponse:
         """Generate text embeddings using the Cohere API."""
         client = self._get_client()
-
-        if isinstance(texts, str):
-            texts = [texts]
+        texts = self._normalize_text_input(texts)
 
         try:
             all_embeddings = []
@@ -70,14 +74,9 @@ class CohereProvider(EmbeddingProvider):
                     input_type=self.input_type.value,
                     embedding_types=["float"],
                 )
-                all_embeddings.extend(response.embeddings.float_)
+                all_embeddings.extend(np.array(response.embeddings.float_))
 
-            return EmbeddingResult(
-                embeddings=np.array(all_embeddings),
-                model_name=self.model_name,
-                model_provider=self.provider_name,
-                input_type=self.input_type.value,
-            )
+            return self._create_text_response(all_embeddings, self.input_type.value)
 
         except Exception as e:
             raise EmbeddingError(f"Failed to embed text with Cohere: {e}") from e
@@ -85,15 +84,12 @@ class CohereProvider(EmbeddingProvider):
     def embed_image(
         self,
         images: Union[Path, str, List[Union[Path, str]]],
-    ) -> EmbeddingResult:
+    ) -> EmbeddingResponse:
         """Generate embeddings for images using Cohere API."""
         client = self._get_client()
-        input_type = "image"
-
-        if isinstance(images, (str, Path)):
-            images = [Path(images)]
-        else:
-            images = [Path(img) for img in images]
+        images = self._normalize_image_input(images)
+        total_images = len(images)
+        logger.info(f"Starting to process {total_images} images")
 
         try:
             all_embeddings = []
@@ -102,12 +98,17 @@ class CohereProvider(EmbeddingProvider):
             # Process images in batches
             for i in range(0, len(images), self.image_batch_size):
                 batch_images = images[i : i + self.image_batch_size]
+                logger.info(f"Processing batch {i//self.image_batch_size + 1} of {(total_images + self.image_batch_size - 1)//self.image_batch_size} ({len(batch_images)} images)")
                 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))
+                    b64_data, content_type = image_to_base64(image)
+                    # Construct full data URI for API
+                    data_uri = f"data:{content_type};base64,{b64_data}"
+                    b64_images.append(data_uri)
+                    all_b64_images.append((b64_data, content_type))
 
                 response = client.embed(
                     model=self.model_name,
@@ -116,21 +117,15 @@ class CohereProvider(EmbeddingProvider):
                     embedding_types=["float"],
                 )
 
-                all_embeddings.extend(response.embeddings.float_)
-                all_b64_images.extend(b64_images)
+                all_embeddings.extend(np.array(response.embeddings.float_))
 
-            return EmbeddingResult(
-                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,
+            logger.info(f"Successfully processed all {total_images} images")
+            return self._create_image_response(
+                all_embeddings,
+                [b64 for b64, _ in all_b64_images],
+                [content_type for _, content_type in all_b64_images],
             )
 
         except Exception as e:
+            logger.error(f"Failed to embed images: {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)
-        return self.embed_image(image_paths)

embedkit/providers/colpali.py

@@ -8,8 +8,9 @@ import numpy as np
 import torch
 from PIL import Image
 
-from ..utils import pdf_to_images, image_to_base64
-from ..base import EmbeddingProvider, EmbeddingError, EmbeddingResult
+from ..models import Model
+from ..utils import image_to_base64
+from ..base import EmbeddingProvider, EmbeddingError, EmbeddingResponse
 
 logger = logging.getLogger(__name__)
 
@@ -19,15 +20,17 @@ class ColPaliProvider(EmbeddingProvider):
 
     def __init__(
         self,
-        model_name: str,
+        model: Model.ColPali,
         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
+        super().__init__(
+            model_name=model.value,
+            text_batch_size=text_batch_size,
+            image_batch_size=image_batch_size,
+            provider_name="ColPali",
+        )
 
         # Auto-detect device
         if device is None:
@@ -38,24 +41,43 @@ class ColPaliProvider(EmbeddingProvider):
             else:
                 device = "cpu"
 
-        self.device = device
-        self._model = None
-        self._processor = None
+        self._hf_device = device
+        self._hf_model = None
+        self._hf_processor = None
 
     def _load_model(self):
         """Lazy load the model."""
-        if self._model is None:
+        if self._hf_model is None:
             try:
-                from colpali_engine.models import ColPali, ColPaliProcessor
-
-                self._model = ColPali.from_pretrained(
-                    self.model_name,
-                    torch_dtype=torch.bfloat16,
-                    device_map=self.device,
-                ).eval()
-
-                self._processor = ColPaliProcessor.from_pretrained(self.model_name)
-                logger.info(f"Loaded ColPali model on {self.device}")
+                if self.model_name in [Model.ColPali.COLPALI_V1_3.value]:
+                    from colpali_engine.models import ColPali, ColPaliProcessor
+
+                    self._hf_model = ColPali.from_pretrained(
+                        self.model_name,
+                        torch_dtype=torch.bfloat16,
+                        device_map=self._hf_device,
+                    ).eval()
+
+                    self._hf_processor = ColPaliProcessor.from_pretrained(self.model_name)
+
+                elif self.model_name in [
+                    Model.ColPali.COLSMOL_500M.value,
+                    Model.ColPali.COLSMOL_256M.value,
+                ]:
+                    from colpali_engine.models import ColIdefics3, ColIdefics3Processor
+
+                    self._hf_model = ColIdefics3.from_pretrained(
+                        self.model_name,
+                        torch_dtype=torch.bfloat16,
+                        device_map=self._hf_device,
+                    ).eval()
+                    self._hf_processor = ColIdefics3Processor.from_pretrained(
+                        self.model_name
+                    )
+                else:
+                    raise ValueError(f"Unable to load model for: {self.model_name}.")
+
+                logger.info(f"Loaded {self.model_name} on {self._hf_device}")
 
             except ImportError as e:
                 raise EmbeddingError(
@@ -64,58 +86,51 @@ class ColPaliProvider(EmbeddingProvider):
             except Exception as e:
                 raise EmbeddingError(f"Failed to load model: {e}") from e
 
-    def embed_text(self, texts: Union[str, List[str]]) -> np.ndarray:
+    def embed_text(self, texts: Union[str, List[str]], **kwargs) -> EmbeddingResponse:
         """Generate embeddings for text inputs."""
         self._load_model()
-
-        if isinstance(texts, str):
-            texts = [texts]
+        texts = self._normalize_text_input(texts)
 
         try:
             # Process texts in batches
-            all_embeddings = []
+            all_embeddings: List[np.ndarray] = []
 
             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)
+                processed = self._hf_processor.process_queries(batch_texts).to(self._hf_device)
 
                 with torch.no_grad():
-                    batch_embeddings = self._model(**processed)
+                    batch_embeddings = self._hf_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=final_embeddings,
-                model_name=self.model_name,
-                model_provider=self.provider_name,
-                input_type="text",
-            )
+            return self._create_text_response(final_embeddings)
 
         except Exception as e:
             raise EmbeddingError(f"Failed to embed text: {e}") from e
 
     def embed_image(
         self, images: Union[Path, str, List[Union[Path, str]]]
-    ) -> np.ndarray:
+    ) -> EmbeddingResponse:
         """Generate embeddings for images."""
         self._load_model()
-
-        if isinstance(images, (str, Path)):
-            images = [Path(images)]
-        else:
-            images = [Path(img) for img in images]
+        images = self._normalize_image_input(images)
+        total_images = len(images)
+        logger.info(f"Starting to process {total_images} images")
 
         try:
             # Process images in batches
-            all_embeddings = []
-            all_b64_images = []
+            all_embeddings: List[np.ndarray] = []
+            all_b64_data: List[str] = []
+            all_content_types: List[str] = []
 
             for i in range(0, len(images), self.image_batch_size):
                 batch_images = images[i : i + self.image_batch_size]
+                logger.info(f"Processing batch {i//self.image_batch_size + 1} of {(total_images + self.image_batch_size - 1)//self.image_batch_size} ({len(batch_images)} images)")
                 pil_images = []
-                b64_images = []
+                batch_b64_data = []
+                batch_content_types = []
 
                 for img_path in batch_images:
                     if not img_path.exists():
@@ -123,30 +138,25 @@ class ColPaliProvider(EmbeddingProvider):
 
                     with Image.open(img_path) as img:
                         pil_images.append(img.convert("RGB"))
-                    b64_images.append(image_to_base64(img_path))
+                    b64, content_type = image_to_base64(img_path)
+                    batch_b64_data.append(b64)
+                    batch_content_types.append(content_type)
 
-                processed = self._processor.process_images(pil_images).to(self.device)
+                processed = self._hf_processor.process_images(pil_images).to(self._hf_device)
 
                 with torch.no_grad():
-                    batch_embeddings = self._model(**processed)
+                    batch_embeddings = self._hf_model(**processed)
                     all_embeddings.append(batch_embeddings.cpu().float().numpy())
-                    all_b64_images.extend(b64_images)
+                    all_b64_data.extend(batch_b64_data)
+                    all_content_types.extend(batch_content_types)
 
             # Concatenate all batch embeddings
             final_embeddings = np.concatenate(all_embeddings, axis=0)
-
-            return EmbeddingResult(
-                embeddings=final_embeddings,
-                model_name=self.model_name,
-                model_provider=self.provider_name,
-                input_type="image",
-                source_images_b64=all_b64_images,
+            logger.info(f"Successfully processed all {total_images} images")
+            return self._create_image_response(
+                final_embeddings, all_b64_data, all_content_types
             )
 
         except Exception as e:
+            logger.error(f"Failed to embed images: {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)
-        return self.embed_image(images)

embedkit/utils.py

@@ -1,37 +1,84 @@
+import tempfile
+import shutil
+import logging
+from contextlib import contextmanager
 from pdf2image import convert_from_path
 from pathlib import Path
 from .config import get_temp_dir
-from typing import Union
+from typing import Union, List, Iterator, Callable, TypeVar, Any
 
+logger = logging.getLogger(__name__)
 
-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 = []
+T = TypeVar("T")
 
-    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
+@contextmanager
+def temporary_directory() -> Iterator[Path]:
+    """Create a temporary directory that is automatically cleaned up when done.
 
+    Yields:
+        Path: Path to the temporary directory
+    """
+    temp_dir = Path(tempfile.mkdtemp())
+    try:
+        yield temp_dir
+    finally:
+        shutil.rmtree(temp_dir)
+
+
+def pdf_to_images(pdf_path: Path) -> List[Path]:
+    """Convert a PDF file to a list of images.
+
+    The images are stored in a temporary directory that will be automatically
+    cleaned up when the process exits.
+
+    Args:
+        pdf_path: Path to the PDF file
+
+    Returns:
+        List[Path]: List of paths to the generated images
+
+    Note:
+        The temporary files will be automatically cleaned up when the process exits.
+        Do not rely on these files persisting after the function returns.
+    """
+    with temporary_directory() as temp_dir:
+        images = convert_from_path(pdf_path=str(pdf_path), output_folder=str(temp_dir))
+        image_paths = []
+
+        for i, image in enumerate(images):
+            output_path = temp_dir / f"{pdf_path.stem}_{i}.png"
+            image.save(output_path)
+            final_path = Path(tempfile.mktemp(suffix=".png"))
+            shutil.move(output_path, final_path)
+            image_paths.append(final_path)
+
+        return image_paths
+
+
+def image_to_base64(image_path: Union[str, Path]) -> tuple[str, str]:
+    """Convert an image to base64 and return the base64 data and content type.
 
-def image_to_base64(image_path: Union[str, Path]):
+    Args:
+        image_path: Path to the image file
+
+    Returns:
+        tuple[str, str]: (base64_data, content_type)
+
+    Raises:
+        ValueError: If the image cannot be read or has an unsupported format
+    """
     import base64
 
     try:
-        base64_only = base64.b64encode(Path(image_path).read_bytes()).decode("utf-8")
+        base64_data = 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)
+    else:
+        image_path_str = image_path
 
     if image_path_str.lower().endswith(".png"):
         content_type = "image/png"
@@ -43,6 +90,53 @@ def image_to_base64(image_path: Union[str, Path]):
         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
+    return base64_data, content_type
+
+
+def with_pdf_cleanup(embed_func: Callable[..., T]) -> Callable[..., T]:
+    """Decorator to handle PDF to image conversion with automatic cleanup.
+
+    This decorator handles the common pattern of:
+    1. Converting PDF to images
+    2. Passing images to an embedding function
+    3. Cleaning up temporary files
+
+    Args:
+        embed_func: Function that takes a list of image paths and returns embeddings
+
+    Returns:
+        Callable that takes a PDF path and returns embeddings
+    """
+
+    def wrapper(*args, **kwargs) -> T:
+        # First argument is self for instance methods
+        pdf_path = args[-1] if args else kwargs.get("pdf_path")
+        if not pdf_path:
+            raise ValueError(
+                "PDF path must be provided as the last positional argument or as 'pdf_path' keyword argument"
+            )
+
+        images = []  # Initialize images as empty list
+        try:
+            images = pdf_to_images(pdf_path)
+            # Call the original function with the images instead of pdf_path
+            if args:
+                # For instance methods, replace the last argument (pdf_path) with images
+                args = list(args)
+                args[-1] = images
+            else:
+                kwargs["pdf_path"] = images
+            return embed_func(*args, **kwargs)
+        finally:
+            # Clean up temporary files created by pdf_to_images
+            for img_path in images:
+                try:
+                    if img_path.exists() and str(img_path).startswith(
+                        tempfile.gettempdir()
+                    ):
+                        img_path.unlink()
+                except Exception as e:
+                    logger.warning(f"Failed to clean up temporary file {img_path}: {e}")
+
+    return wrapper

embedkit-0.1.6.dist-info/METADATA

@@ -0,0 +1,143 @@
+Metadata-Version: 2.4
+Name: embedkit
+Version: 0.1.6
+Summary: A simple toolkit for generating vector embeddings across multiple providers and models
+Author-email: JP Hwang <me@jphwang.com>
+License: MIT
+License-File: LICENSE
+Keywords: ai,cohere,colpali,embeddings,machine-learning,vector
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Requires-Python: >=3.10
+Requires-Dist: accelerate>=1.7.0
+Requires-Dist: cohere>=5.15.0
+Requires-Dist: colpali-engine<0.4.0,>=0.3.0
+Requires-Dist: pdf2image>=1.17.0
+Requires-Dist: pillow>=11.2.1
+Requires-Dist: torch<=2.5
+Requires-Dist: transformers>=4.46.2
+Description-Content-Type: text/markdown
+
+# EmbedKit
+
+A unified interface for text and image embeddings, supporting multiple providers.
+
+## Installation
+
+```bash
+pip install embedkit
+```
+
+## Usage
+
+### Text Embeddings
+
+```python
+from embedkit import EmbedKit
+from embedkit.classes import Model, CohereInputType
+
+# Initialize with ColPali
+kit = EmbedKit.colpali(
+    model=Model.ColPali.COLPALI_V1_3,  # or COLSMOL_256M, COLSMOL_500M
+    text_batch_size=16,  # Optional: process text in batches of 16
+    image_batch_size=8,  # Optional: process images in batches of 8
+)
+
+# Get embeddings
+result = kit.embed_text("Hello world")
+print(result.model_provider)
+print(result.input_type)
+print(result.objects[0].embedding.shape)  # Returns 2D array for ColPali
+print(result.objects[0].source_b64)
+
+# Initialize with Cohere
+kit = EmbedKit.cohere(
+    model=Model.Cohere.EMBED_V4_0,
+    api_key="your-api-key",
+    text_input_type=CohereInputType.SEARCH_QUERY,  # or SEARCH_DOCUMENT
+    text_batch_size=64,  # Optional: process text in batches of 64
+    image_batch_size=8,  # Optional: process images in batches of 8
+)
+
+# Get embeddings
+result = kit.embed_text("Hello world")
+print(result.model_provider)
+print(result.input_type)
+print(result.objects[0].embedding.shape)  # Returns 1D array for Cohere
+print(result.objects[0].source_b64)
+```
+
+### Image Embeddings
+
+```python
+from pathlib import Path
+
+# Get embeddings for an image
+image_path = Path("path/to/image.png")
+result = kit.embed_image(image_path)
+
+print(result.model_provider)
+print(result.input_type)
+print(result.objects[0].embedding.shape)  # 2D for ColPali, 1D for Cohere
+print(result.objects[0].source_b64)  # Base64 encoded image
+```
+
+### PDF Embeddings
+
+```python
+from pathlib import Path
+
+# Get embeddings for a PDF
+pdf_path = Path("path/to/document.pdf")
+result = kit.embed_pdf(pdf_path)
+
+print(result.model_provider)
+print(result.input_type)
+print(result.objects[0].embedding.shape)  # 2D for ColPali, 1D for Cohere
+print(result.objects[0].source_b64)  # Base64 encoded PDF page
+```
+
+## Response Format
+
+The embedding methods return an `EmbeddingResponse` object with the following structure:
+
+```python
+class EmbeddingResponse:
+    model_name: str
+    model_provider: str
+    input_type: str
+    objects: List[EmbeddingObject]
+
+class EmbeddingObject:
+    embedding: np.ndarray  # 1D array for Cohere, 2D array for ColPali
+    source_b64: Optional[str]  # Base64 encoded source for images and PDFs
+```
+
+## Supported Models
+
+### ColPali
+- `Model.ColPali.COLPALI_V1_3`
+- `Model.ColPali.COLSMOL_256M`
+- `Model.ColPali.COLSMOL_500M`
+
+### Cohere
+- `Model.Cohere.EMBED_V4_0`
+- `Model.Cohere.EMBED_ENGLISH_V3_0`
+- `Model.Cohere.EMBED_ENGLISH_LIGHT_V3_0`
+- `Model.Cohere.EMBED_MULTILINGUAL_V3_0`
+- `Model.Cohere.EMBED_MULTILINGUAL_LIGHT_V3_0`
+
+## Requirements
+
+- Python 3.10+
+
+## License
+
+MIT

embedkit-0.1.6.dist-info/RECORD

@@ -0,0 +1,13 @@
+embedkit/__init__.py,sha256=ahMyC4SjYLr3QpM39obwDKpJO_HQ4_kqvG4d7jD5XtA,4503
+embedkit/base.py,sha256=FfyzCqG7azs4yJj4RU6QKQRy9_wfd3KUFidThplmEo0,3739
+embedkit/classes.py,sha256=LI5egTsfcBglUqIxBrDV_ymA1eKcxOUHJeMAhKVtU_Y,582
+embedkit/config.py,sha256=EVGODSKxQAr46bU8dyORFunsfRuj6dnvtSqa4MxUZCo,138
+embedkit/models.py,sha256=fndxt6QoV7nuyeXijs39K3QLs1l8ATw2PX2X-64oKt8,575
+embedkit/utils.py,sha256=xqIjwkWDgDr1mCtbb5UbZo_Ci50griubxYJq3pSBv34,4560
+embedkit/providers/__init__.py,sha256=HaS-HNQabvhn9xLNZCq3VUqPCb7rGG4pvgvpKP4AXcw,201
+embedkit/providers/cohere.py,sha256=B5hgOIY_rexgc0s_DWjZ-yuuOXaYLQV-jmrY-gdBhz8,4726
+embedkit/providers/colpali.py,sha256=8LsV-kkBmNYzCGXWeZy-KLEyMOxTvzOAS4gnzN5q_f4,6260
+embedkit-0.1.6.dist-info/METADATA,sha256=tf0SLNvH5OOYFudiMmgEALYE6fnV716WJO0pPZIQJPw,3871
+embedkit-0.1.6.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+embedkit-0.1.6.dist-info/licenses/LICENSE,sha256=-g2Rad7b3rb2oVwOTwfMOIpscHT1zuaJoguamLRCBJs,1072
+embedkit-0.1.6.dist-info/RECORD,,

embedkit-0.1.3.dist-info/METADATA

@@ -1,59 +0,0 @@
-Metadata-Version: 2.4
-Name: embedkit
-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
-License-File: LICENSE
-Keywords: ai,cohere,colpali,embeddings,machine-learning,vector
-Classifier: Development Status :: 4 - Beta
-Classifier: Intended Audience :: Developers
-Classifier: License :: OSI Approved :: MIT License
-Classifier: Programming Language :: Python :: 3
-Classifier: Programming Language :: Python :: 3.10
-Classifier: Programming Language :: Python :: 3.11
-Classifier: Programming Language :: Python :: 3.12
-Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
-Classifier: Topic :: Software Development :: Libraries :: Python Modules
-Requires-Python: >=3.10
-Requires-Dist: accelerate>=1.7.0
-Requires-Dist: cohere>=5.15.0
-Requires-Dist: colpali-engine<0.4.0,>=0.3.0
-Requires-Dist: pdf2image>=1.17.0
-Requires-Dist: pillow>=11.2.1
-Requires-Dist: torch<=2.5
-Requires-Dist: transformers
-Description-Content-Type: text/markdown
-
-# EmbedKit
-
-A Python library for generating embeddings from text, images, and PDFs using various models (e.g. from Cohere, ColPali).
-
-## Usage
-
-See [main.py](main.py) for examples.
-
-```python
-from embedkit import EmbedKit
-from embedkit.models import Model
-
-# Instantiate a kit
-# Using ColPali
-kit = EmbedKit.colpali(model=Model.ColPali.V1_3)
-
-# Using Cohere
-kit = EmbedKit.cohere(
-    model=Model.Cohere.EMBED_V4_0,
-    api_key="your_api_key",
-    text_input_type=CohereInputType.SEARCH_DOCUMENT,
-)
-
-# Then - the embedding API is consistent
-embeddings = kit.embed_text("Hello world") or kit.embed_text(["Hello world", "Hello world"])
-embeddings = kit.embed_image("path/to/image.png") or kit.embed_image(["path/to/image1.png", "path/to/image2.png"])
-embeddings = kit.embed_pdf("path/to/pdf.pdf")  # Single PDF only
-```
-
-## License
-
-MIT

embedkit-0.1.3.dist-info/RECORD

@@ -1,13 +0,0 @@
-embedkit/__init__.py,sha256=E9BiBw4ikP0FOnbj0Y7qeF71XijghYwTUnjh-_Ogvfg,4519
-embedkit/base.py,sha256=d6-899ffow4fGsbIJGBO6RdFVDuo-4hKH_MHR8uTNe8,1279
-embedkit/classes.py,sha256=iaMS0fZqDm4iL39qhtivyjMY8r7rGANnW8ZLfRjQQzA,596
-embedkit/config.py,sha256=EVGODSKxQAr46bU8dyORFunsfRuj6dnvtSqa4MxUZCo,138
-embedkit/models.py,sha256=EBIYkyZeIhGaOPL-9bslHHdLaZ7qzOYLd0qxVZ7VX7w,226
-embedkit/utils.py,sha256=91BzzvbYSrUsWeW3CTAw3yK-M3S5FgQXov16gxffkUo,1572
-embedkit/providers/__init__.py,sha256=HaS-HNQabvhn9xLNZCq3VUqPCb7rGG4pvgvpKP4AXcw,201
-embedkit/providers/cohere.py,sha256=e4MLe_SaeoCuS9JNyzPIhMxMGxBq4521TiMIWI6Mhes,4578
-embedkit/providers/colpali.py,sha256=vyCs_Q5DfvsPDGOGYsOA_oGRi5Z-Lvb07GR9pNT_OH4,5186
-embedkit-0.1.3.dist-info/METADATA,sha256=2Taju_C4Cv3Vugh3RVVBLRSAhVlMU-i9DirGvQT1Gjw,1893
-embedkit-0.1.3.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-embedkit-0.1.3.dist-info/licenses/LICENSE,sha256=-g2Rad7b3rb2oVwOTwfMOIpscHT1zuaJoguamLRCBJs,1072
-embedkit-0.1.3.dist-info/RECORD,,