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

embedkit/__init__.py

@@ -5,7 +5,6 @@ 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, EmbeddingResponse
@@ -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,
+            api_key=api_key,
+            model=model,
             text_batch_size=text_batch_size,
             image_batch_size=image_batch_size,
-            text_input_type=text_input_type
+            text_input_type=text_input_type,
         )
         return cls(provider)
 

embedkit/base.py

@@ -7,11 +7,15 @@ 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
@@ -29,6 +33,67 @@ class EmbeddingResponse:
 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) -> EmbeddingResponse:
         """Generate document text embeddings using the configured provider."""
@@ -41,10 +106,14 @@ class EmbeddingProvider(ABC):
         """Generate image embeddings using the configured provider."""
         pass
 
-    @abstractmethod
-    def embed_pdf(self, pdf: Union[Path, str]) -> EmbeddingResponse:
-        """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

@@ -13,9 +13,4 @@ from . import EmbeddingResponse, EmbeddingError
 from .models import Model
 from .providers.cohere import CohereInputType
 
-__all__ = [
-    "EmbeddingResponse",
-    "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, EmbeddingResponse, EmbeddingObject
+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."""
@@ -54,9 +60,7 @@ class CohereProvider(EmbeddingProvider):
     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 = []
@@ -72,16 +76,7 @@ class CohereProvider(EmbeddingProvider):
                 )
                 all_embeddings.extend(np.array(response.embeddings.float_))
 
-            return EmbeddingResponse(
-                model_name=self.model_name,
-                model_provider=self.provider_name,
-                input_type=self.input_type.value,
-                objects=[
-                    EmbeddingObject(
-                        embedding=e,
-                    ) for e in all_embeddings
-                ]
-            )
+            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
@@ -92,12 +87,9 @@ class CohereProvider(EmbeddingProvider):
     ) -> 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 = []
@@ -106,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,
@@ -121,24 +118,14 @@ class CohereProvider(EmbeddingProvider):
                 )
 
                 all_embeddings.extend(np.array(response.embeddings.float_))
-                all_b64_images.extend(b64_images)
-
-            return EmbeddingResponse(
-                model_name=self.model_name,
-                model_provider=self.provider_name,
-                input_type=input_type,
-                objects=[
-                    EmbeddingObject(
-                        embedding=all_embeddings[i],
-                        source_b64=all_b64_images[i]
-                    ) for i in range(len(all_embeddings))
-                ]
+
+            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) -> EmbeddingResponse:
-        """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, EmbeddingResponse, EmbeddingObject
+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,38 +86,26 @@ 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]]) -> EmbeddingResponse:
+    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 EmbeddingResponse(
-                model_name=self.model_name,
-                model_provider=self.provider_name,
-                input_type="text",
-                objects=[
-                    EmbeddingObject(
-                        embedding=e,
-                    ) for e in final_embeddings
-                ]
-            )
+            return self._create_text_response(final_embeddings)
 
         except Exception as e:
             raise EmbeddingError(f"Failed to embed text: {e}") from e
@@ -105,21 +115,22 @@ class ColPaliProvider(EmbeddingProvider):
     ) -> 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():
@@ -127,34 +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 EmbeddingResponse(
-                model_name=self.model_name,
-                model_provider=self.provider_name,
-                input_type="image",
-                objects=[
-                    EmbeddingObject(
-                        embedding=final_embeddings[i],
-                        source_b64=all_b64_images[i]
-                    ) for i in range(len(final_embeddings))
-                ]
+            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) -> EmbeddingResponse:
-        """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.4.dist-info → embedkit-0.1.6.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: embedkit
-Version: 0.1.4
+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
@@ -22,7 +22,7 @@ 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
+Requires-Dist: transformers>=4.46.2
 Description-Content-Type: text/markdown
 
 # EmbedKit
@@ -45,7 +45,7 @@ from embedkit.classes import Model, CohereInputType
 
 # Initialize with ColPali
 kit = EmbedKit.colpali(
-    model=Model.ColPali.V1_3,
+    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
 )
@@ -54,7 +54,7 @@ kit = EmbedKit.colpali(
 result = kit.embed_text("Hello world")
 print(result.model_provider)
 print(result.input_type)
-print(result.objects[0].embedding.shape)
+print(result.objects[0].embedding.shape)  # Returns 2D array for ColPali
 print(result.objects[0].source_b64)
 
 # Initialize with Cohere
@@ -70,7 +70,7 @@ kit = EmbedKit.cohere(
 result = kit.embed_text("Hello world")
 print(result.model_provider)
 print(result.input_type)
-print(result.objects[0].embedding.shape)
+print(result.objects[0].embedding.shape)  # Returns 1D array for Cohere
 print(result.objects[0].source_b64)
 ```
 
@@ -85,8 +85,8 @@ result = kit.embed_image(image_path)
 
 print(result.model_provider)
 print(result.input_type)
-print(result.objects[0].embedding.shape)
-print(result.objects[0].source_b64)
+print(result.objects[0].embedding.shape)  # 2D for ColPali, 1D for Cohere
+print(result.objects[0].source_b64)  # Base64 encoded image
 ```
 
 ### PDF Embeddings
@@ -100,8 +100,8 @@ result = kit.embed_pdf(pdf_path)
 
 print(result.model_provider)
 print(result.input_type)
-print(result.objects[0].embedding.shape)
-print(result.objects[0].source_b64)
+print(result.objects[0].embedding.shape)  # 2D for ColPali, 1D for Cohere
+print(result.objects[0].source_b64)  # Base64 encoded PDF page
 ```
 
 ## Response Format
@@ -116,17 +116,23 @@ class EmbeddingResponse:
     objects: List[EmbeddingObject]
 
 class EmbeddingObject:
-    embedding: np.ndarray
-    source_b64: Optional[str]
+    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.V1_3`
+- `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
 

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.4.dist-info/RECORD

@@ -1,13 +0,0 @@
-embedkit/__init__.py,sha256=O7C_e30VMrLliDamyNM_srdyuGz50PvGzBJo9NgePoU,4555
-embedkit/base.py,sha256=QkP-Ih7aYjIaxP5Hgp7MVypQlI5mhyOqLWKe0lOpIrQ,1344
-embedkit/classes.py,sha256=TlCl58Vo8p0q8SfJGINtBmCXY2yVFLtdlfQqBKrxivw,600
-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=5-ux5UzpRqlRPm2PjgmRfcxH4NhRRjXa3cOZGWp0jDc,4880
-embedkit/providers/colpali.py,sha256=PUjTESyF0r__JdgfT22kK1DGrf7XwYC2EuyBelspTsc,5500
-embedkit-0.1.4.dist-info/METADATA,sha256=KPCFTWyeh_8unP1zBmVCUmA0c-yfXwxeMdWf9zO4CjA,3316
-embedkit-0.1.4.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-embedkit-0.1.4.dist-info/licenses/LICENSE,sha256=-g2Rad7b3rb2oVwOTwfMOIpscHT1zuaJoguamLRCBJs,1072
-embedkit-0.1.4.dist-info/RECORD,,