cocoindex 0.1.75__cp313-cp313-win_amd64.whl → 0.1.76__cp313-cp313-win_amd64.whl

cocoindex/functions.py

@@ -1,6 +1,7 @@
 """All builtin functions."""
 
 import dataclasses
+import functools
 from typing import Annotated, Any, Literal
 
 import numpy as np
@@ -23,6 +24,16 @@ class CustomLanguageSpec:
     aliases: list[str] = dataclasses.field(default_factory=list)
 
 
+@dataclasses.dataclass
+class ColPaliModelInfo:
+    """Data structure for ColPali model and processor."""
+
+    model: Any
+    processor: Any
+    dimension: int
+    device: Any
+
+
 class SplitRecursively(op.FunctionSpec):
     """Split a document (in string) recursively."""
 
@@ -99,3 +110,189 @@ class SentenceTransformerEmbedExecutor:
         assert self._model is not None
         result: NDArray[np.float32] = self._model.encode(text, convert_to_numpy=True)
         return result
+
+
+@functools.cache
+def _get_colpali_model_and_processor(model_name: str) -> ColPaliModelInfo:
+    """Get or load ColPali model and processor, with caching."""
+    try:
+        from colpali_engine.models import ColPali, ColPaliProcessor  # type: ignore[import-untyped]
+        from colpali_engine.utils.torch_utils import get_torch_device  # type: ignore[import-untyped]
+        import torch
+    except ImportError as e:
+        raise ImportError(
+            "ColPali is not available. Make sure cocoindex is installed with ColPali support."
+        ) from e
+
+    device = get_torch_device("auto")
+    model = ColPali.from_pretrained(
+        model_name, device_map=device, torch_dtype=torch.bfloat16
+    ).eval()
+    processor = ColPaliProcessor.from_pretrained(model_name)
+
+    # Get dimension from the actual model
+    dimension = _detect_colpali_dimension(model, processor, device)
+
+    return ColPaliModelInfo(
+        model=model,
+        processor=processor,
+        dimension=dimension,
+        device=device,
+    )
+
+
+def _detect_colpali_dimension(model: Any, processor: Any, device: Any) -> int:
+    """Detect ColPali embedding dimension from the actual model config."""
+    # Try to access embedding dimension
+    if hasattr(model.config, "embedding_dim"):
+        dim = model.config.embedding_dim
+    else:
+        # Fallback: infer from output shape with dummy data
+        from PIL import Image
+        import numpy as np
+        import torch
+
+        dummy_img = Image.fromarray(np.zeros((224, 224, 3), np.uint8))
+        # Use the processor to process the dummy image
+        processed = processor.process_images([dummy_img]).to(device)
+        with torch.no_grad():
+            output = model(**processed)
+        dim = int(output.shape[-1])
+    if isinstance(dim, int):
+        return dim
+    else:
+        raise ValueError(f"Expected integer dimension, got {type(dim)}: {dim}")
+    return dim
+
+
+class ColPaliEmbedImage(op.FunctionSpec):
+    """
+    `ColPaliEmbedImage` embeds images using the ColPali multimodal model.
+
+    ColPali (Contextual Late-interaction over Patches) uses late interaction
+    between image patch embeddings and text token embeddings for retrieval.
+
+    Args:
+        model: The ColPali model name to use (e.g., "vidore/colpali-v1.2")
+
+    Note:
+        This function requires the optional colpali-engine dependency.
+        Install it with: pip install 'cocoindex[embeddings]'
+    """
+
+    model: str
+
+
+@op.executor_class(
+    gpu=True,
+    cache=True,
+    behavior_version=1,
+)
+class ColPaliEmbedImageExecutor:
+    """Executor for ColPaliEmbedImage."""
+
+    spec: ColPaliEmbedImage
+    _model_info: ColPaliModelInfo
+
+    def analyze(self, _img_bytes: Any) -> type:
+        # Get shared model and dimension
+        self._model_info = _get_colpali_model_and_processor(self.spec.model)
+
+        # Return multi-vector type: Variable patches x Fixed hidden dimension
+        dimension = self._model_info.dimension
+        return Vector[Vector[np.float32, Literal[dimension]]]  # type: ignore
+
+    def __call__(self, img_bytes: bytes) -> Any:
+        try:
+            from PIL import Image
+            import torch
+            import io
+        except ImportError as e:
+            raise ImportError(
+                "Required dependencies (PIL, torch) are missing for ColPali image embedding."
+            ) from e
+
+        model = self._model_info.model
+        processor = self._model_info.processor
+        device = self._model_info.device
+
+        pil_image = Image.open(io.BytesIO(img_bytes)).convert("RGB")
+        inputs = processor.process_images([pil_image]).to(device)
+        with torch.no_grad():
+            embeddings = model(**inputs)
+
+        # Return multi-vector format: [patches, hidden_dim]
+        if len(embeddings.shape) != 3:
+            raise ValueError(
+                f"Expected 3D tensor [batch, patches, hidden_dim], got shape {embeddings.shape}"
+            )
+
+        # Keep patch-level embeddings: [batch, patches, hidden_dim] -> [patches, hidden_dim]
+        patch_embeddings = embeddings[0]  # Remove batch dimension
+
+        return patch_embeddings.cpu().to(torch.float32).numpy()
+
+
+class ColPaliEmbedQuery(op.FunctionSpec):
+    """
+    `ColPaliEmbedQuery` embeds text queries using the ColPali multimodal model.
+
+    This produces query embeddings compatible with ColPali image embeddings
+    for late interaction scoring (MaxSim).
+
+    Args:
+        model: The ColPali model name to use (e.g., "vidore/colpali-v1.2")
+
+    Note:
+        This function requires the optional colpali-engine dependency.
+        Install it with: pip install 'cocoindex[embeddings]'
+    """
+
+    model: str
+
+
+@op.executor_class(
+    gpu=True,
+    cache=True,
+    behavior_version=1,
+)
+class ColPaliEmbedQueryExecutor:
+    """Executor for ColPaliEmbedQuery."""
+
+    spec: ColPaliEmbedQuery
+    _model_info: ColPaliModelInfo
+
+    def analyze(self, _query: Any) -> type:
+        # Get shared model and dimension
+        self._model_info = _get_colpali_model_and_processor(self.spec.model)
+
+        # Return multi-vector type: Variable tokens x Fixed hidden dimension
+        dimension = self._model_info.dimension
+        return Vector[Vector[np.float32, Literal[dimension]]]  # type: ignore
+
+    def __call__(self, query: str) -> Any:
+        try:
+            import torch
+        except ImportError as e:
+            raise ImportError(
+                "Required dependencies (torch) are missing for ColPali query embedding."
+            ) from e
+
+        model = self._model_info.model
+        processor = self._model_info.processor
+        device = self._model_info.device
+
+        inputs = processor.process_queries([query]).to(device)
+        with torch.no_grad():
+            embeddings = model(**inputs)
+
+        # Return multi-vector format: [tokens, hidden_dim]
+        if len(embeddings.shape) != 3:
+            raise ValueError(
+                f"Expected 3D tensor [batch, tokens, hidden_dim], got shape {embeddings.shape}"
+            )
+
+        # Keep token-level embeddings: [batch, tokens, hidden_dim] -> [tokens, hidden_dim]
+        token_embeddings = embeddings[0]  # Remove batch dimension
+
+        return token_embeddings.cpu().to(torch.float32).numpy()

{cocoindex-0.1.75.dist-info → cocoindex-0.1.76.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cocoindex
-Version: 0.1.75
+Version: 0.1.76
 Requires-Dist: click>=8.1.8
 Requires-Dist: rich>=14.0.0
 Requires-Dist: python-dotenv>=1.1.0
@@ -12,9 +12,12 @@ Requires-Dist: ruff ; extra == 'dev'
 Requires-Dist: mypy ; extra == 'dev'
 Requires-Dist: pre-commit ; extra == 'dev'
 Requires-Dist: sentence-transformers>=3.3.1 ; extra == 'embeddings'
+Requires-Dist: colpali-engine ; extra == 'colpali'
 Requires-Dist: sentence-transformers>=3.3.1 ; extra == 'all'
+Requires-Dist: colpali-engine ; extra == 'all'
 Provides-Extra: dev
 Provides-Extra: embeddings
+Provides-Extra: colpali
 Provides-Extra: all
 License-File: LICENSE
 Summary: With CocoIndex, users declare the transformation, CocoIndex creates & maintains an index, and keeps the derived index up to date based on source update, with minimal computation and changes.

{cocoindex-0.1.75.dist-info → cocoindex-0.1.76.dist-info}/RECORD

@@ -1,14 +1,14 @@
-cocoindex-0.1.75.dist-info/METADATA,sha256=mc2W-NDPW0z5rlfOkSmYNH3duucqIzUtRD767vGQxgo,11745
-cocoindex-0.1.75.dist-info/WHEEL,sha256=OASh95bEME3UNqnnxl-DTJ7jIUFeKxREESTsmRdKzNQ,96
-cocoindex-0.1.75.dist-info/entry_points.txt,sha256=_NretjYVzBdNTn7dK-zgwr7YfG2afz1u1uSE-5bZXF8,46
-cocoindex-0.1.75.dist-info/licenses/LICENSE,sha256=HrhfyXIkWY2tGFK11kg7vPCqhgh5DcxleloqdhrpyMY,11558
+cocoindex-0.1.76.dist-info/METADATA,sha256=1S5SYoO7DCKsw7aMEIXrJki83RclbmSdxzxf4p84gQY,11867
+cocoindex-0.1.76.dist-info/WHEEL,sha256=OASh95bEME3UNqnnxl-DTJ7jIUFeKxREESTsmRdKzNQ,96
+cocoindex-0.1.76.dist-info/entry_points.txt,sha256=_NretjYVzBdNTn7dK-zgwr7YfG2afz1u1uSE-5bZXF8,46
+cocoindex-0.1.76.dist-info/licenses/LICENSE,sha256=HrhfyXIkWY2tGFK11kg7vPCqhgh5DcxleloqdhrpyMY,11558
 cocoindex/__init__.py,sha256=5zwuS_X-n7QAE5uBSufqXp77OW8KVVD8E5t_6koDLRc,2293
-cocoindex/_engine.cp313-win_amd64.pyd,sha256=hzhbXMK2kw45-4fOzhoYpQ5Evr71CRsIQfkoeYpIy8M,70278144
+cocoindex/_engine.cp313-win_amd64.pyd,sha256=40SXaZVZd2a6GuVXu-wlb_1etHOxmi_sQmttHCmDPL0,70297088
 cocoindex/auth_registry.py,sha256=Qq1IVZb-7K4luRrQSDlOPbISnGEZ4kIDsrCU8H2ARw0,1529
 cocoindex/cli.py,sha256=zjZv7EH-ZLoO1-3Ua2U7Yt3n2A_T7pN2NCuv7FB8WA0,23097
 cocoindex/convert.py,sha256=HVnnfanWxEeSiPq8SBrk9Tf6vXFHgD8LtBE-Umi5LcU,18793
 cocoindex/flow.py,sha256=sZRe3wqwvZWPCdqtfNURCkPGCh3OMMeX6H1DX8NbhjM,37734
-cocoindex/functions.py,sha256=ERtlJJ_D5xNSL4nE6FiqMseX4aAuqsTG6RvJVBC_LPE,3210
+cocoindex/functions.py,sha256=mKOYe-XozuOykrsmFE1cNLZjO2PlXiD6UfFEEHmuGT8,9826
 cocoindex/index.py,sha256=GrqTm1rLwICQ8hadtNvJAxVg7GWMvtMmFcbiNtNzmP0,569
 cocoindex/lib.py,sha256=cZosix4nwROvod4QJOwCzrm6U1CVy_wKMMk7sDDG_Z0,849
 cocoindex/llm.py,sha256=TDUyTxW7ooFmpA_jDNezABhReXksOXN8F8zY9aiuNTw,709
@@ -28,4 +28,4 @@ cocoindex/tests/test_validation.py,sha256=I4wr8lAMAjmy5xgG5N_OJKveXt8XIa96MsQTXh
 cocoindex/typing.py,sha256=gMNJIpGGe-SiXlihDQ-Dw2YdebQvOyG-bWovR-veO6g,13817
 cocoindex/utils.py,sha256=U3W39zD2uZpXX8v84tJD7sRmbC5ar3z_ljAP1cJrYXI,618
 cocoindex/validation.py,sha256=4ZjsW-SZT8X_TEEhEE6QG6D-8Oq_TkPAhTqP0mdFYSE,3194
-cocoindex-0.1.75.dist-info/RECORD,,
+cocoindex-0.1.76.dist-info/RECORD,,