embedkit 0.1.3__tar.gz → 0.1.5__tar.gz

embedkit-0.1.5/PKG-INFO

@@ -0,0 +1,143 @@
+Metadata-Version: 2.4
+Name: embedkit
+Version: 0.1.5
+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.5/README.md

@@ -0,0 +1,116 @@
+# 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.5/main.py

@@ -0,0 +1,107 @@
+# ./main.py
+from embedkit import EmbedKit
+from embedkit.classes import Model, 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")
+longer_pdf = Path("tests/fixtures/2407.01449v6_p1_p5.pdf")
+
+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,
+)
+
+print(f"Trying out Cohere")
+results = kit.embed_text("Hello world")
+assert len(results.objects) == 1
+assert len(results.objects[0].embedding.shape) == 1
+assert results.objects[0].source_b64 == None
+
+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 len(results.objects) == 1
+assert len(results.objects[0].embedding.shape) == 1
+assert results.objects[0].source_b64 == None
+
+results = kit.embed_image(sample_image)
+assert len(results.objects) == 1
+assert len(results.objects[0].embedding.shape) == 1
+assert type(results.objects[0].source_b64) == str
+
+results = kit.embed_pdf(sample_pdf)
+assert len(results.objects) == 1
+assert len(results.objects[0].embedding.shape) == 1
+assert type(results.objects[0].source_b64) == str
+
+results = kit.embed_pdf(longer_pdf)
+assert len(results.objects) == 5
+assert len(results.objects[0].embedding.shape) == 1
+assert type(results.objects[0].source_b64) == str
+
+for colpali_model in [
+    Model.ColPali.COLSMOL_256M,
+    Model.ColPali.COLSMOL_500M,
+    Model.ColPali.COLPALI_V1_3,
+]:
+    print(f"Trying out {colpali_model}")
+    kit = EmbedKit.colpali(
+        model=colpali_model, text_batch_size=16, image_batch_size=8
+    )
+
+    results = kit.embed_text("Hello world")
+    assert len(results.objects) == 1
+    assert len(results.objects[0].embedding.shape) == 2
+    assert results.objects[0].source_b64 == None
+
+    results = kit.embed_image(sample_image)
+    assert len(results.objects) == 1
+    assert len(results.objects[0].embedding.shape) == 2
+    assert type(results.objects[0].source_b64) == str
+
+    results = kit.embed_pdf(sample_pdf)
+    assert len(results.objects) == 1
+    assert len(results.objects[0].embedding.shape) == 2
+    assert type(results.objects[0].source_b64) == str
+
+    results = kit.embed_pdf(longer_pdf)
+    assert len(results.objects) == 5
+    assert len(results.objects[0].embedding.shape) == 2
+    assert type(results.objects[0].source_b64) == str

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

@@ -1,6 +1,6 @@
 [project]
 name = "embedkit"
-version = "0.1.3"
+version = "0.1.5"
 description = "A simple toolkit for generating vector embeddings across multiple providers and models"
 readme = "README.md"
 requires-python = ">=3.10"
@@ -11,7 +11,7 @@ dependencies = [
     "pdf2image>=1.17.0",
     "pillow>=11.2.1",
     "torch<=2.5",
-    "transformers",
+    "transformers>=4.46.2",
 ]
 authors = [
     {name = "JP Hwang", email = "me@jphwang.com"},

{embedkit-0.1.3 → embedkit-0.1.5}/src/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-0.1.5/src/embedkit/base.py

@@ -0,0 +1,122 @@
+# ./src/embedkit/base.py
+"""Base classes for EmbedKit."""
+
+from abc import ABC, abstractmethod
+from typing import Union, List, Optional
+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 EmbeddingResponse:
+    model_name: str
+    model_provider: str
+    input_type: str
+    objects: List[EmbeddingObject]
+
+    @property
+    def shape(self) -> tuple:
+        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) -> EmbeddingResponse:
+        """Generate document text embeddings using the configured provider."""
+        pass
+
+    @abstractmethod
+    def embed_image(
+        self, images: Union[Path, str, List[Union[Path, str]]]
+    ) -> EmbeddingResponse:
+        """Generate image embeddings using the configured provider."""
+        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):
+    """Base exception for embedding-related errors."""
+
+    pass

{embedkit-0.1.3 → embedkit-0.1.5}/src/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-0.1.5/src/embedkit/models.py

@@ -0,0 +1,18 @@
+# ./src/embedkit/models.py
+"""Model definitions and enum for EmbedKit."""
+
+from enum import Enum
+
+
+class Model:
+    class ColPali(Enum):
+        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"