cocoindex 0.1.75__cp313-cp313-manylinux_2_28_x86_64.whl → 0.1.76__cp313-cp313-manylinux_2_28_x86_64.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=kdYHdHSv-kunSaI3MoIUCoZ88Ne3pgeV5jS5WliHGrc,11533
-cocoindex-0.1.75.dist-info/WHEEL,sha256=9Ee4MwqZpMDLH1_kZE8rvruLKRVRs9cmbXRSBB0h-_M,108
-cocoindex-0.1.75.dist-info/entry_points.txt,sha256=_NretjYVzBdNTn7dK-zgwr7YfG2afz1u1uSE-5bZXF8,46
-cocoindex-0.1.75.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+cocoindex-0.1.76.dist-info/METADATA,sha256=w_VMnPWkx5iMgpwgAWTB3KRxSfcGaU1sVES6-jXiAjQ,11655
+cocoindex-0.1.76.dist-info/WHEEL,sha256=9Ee4MwqZpMDLH1_kZE8rvruLKRVRs9cmbXRSBB0h-_M,108
+cocoindex-0.1.76.dist-info/entry_points.txt,sha256=_NretjYVzBdNTn7dK-zgwr7YfG2afz1u1uSE-5bZXF8,46
+cocoindex-0.1.76.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
 cocoindex/__init__.py,sha256=sLpSVO5Cotgn_82lawxvXnaqfa-qj33rytWBAe2MTtU,2201
-cocoindex/_engine.cpython-313-x86_64-linux-gnu.so,sha256=ay8KAvvTUo9j0webtFd6Hwr2mqWNg1QiMbnweVeDg0Q,71097144
+cocoindex/_engine.cpython-313-x86_64-linux-gnu.so,sha256=sIEYde0TfFKjUS1vPsht_7DPZjbvlApMxb_U2CbiIJw,71114160
 cocoindex/auth_registry.py,sha256=PE1-kVkcyC1G2C_V7b1kvYzeq73OFQehWKQP7ln7fJ8,1478
 cocoindex/cli.py,sha256=-gp639JSyQN6YjnhGqCakIzYoSSqXxQMbxbkcYGP0QY,22359
 cocoindex/convert.py,sha256=HodeDl1HVX8nnBH02lQKarw5i3xmkjB0nGj-DXt7Ifc,18284
 cocoindex/flow.py,sha256=egKbBG2X9DjAqmcATcndyRhe9zMZHRd-YxKCpt9BsUg,36551
-cocoindex/functions.py,sha256=LLu_ausirvqnsx_k3euZpv8sLCpBZ4DF77h2HOzbinE,3109
+cocoindex/functions.py,sha256=34sZWoS0zGnaKyooIODQgc6QEPZKiJoWhfb8jKIWwps,9528
 cocoindex/index.py,sha256=j93B9jEvvLXHtpzKWL88SY6wCGEoPgpsQhEGHlyYGFg,540
 cocoindex/lib.py,sha256=f--9dAYd84CZosbDZqNW0oGbBLsY3dXiUTR1VrfQ_QY,817
 cocoindex/llm.py,sha256=WxmWUbNcf9HOCM5xkbDeFs9lF67M3mr810B7deDDc-8,673
@@ -28,4 +28,4 @@ cocoindex/tests/test_validation.py,sha256=X6AQzVs-hVKIXcrHMEMQnhfUE8at7iXQnPq8nH
 cocoindex/typing.py,sha256=qQ0ANF3iuQDeSqipHgL2SDiiXL2reTMUN0aj4ve_T0w,13359
 cocoindex/utils.py,sha256=hUhX-XV6XGCtJSEIpBOuDv6VvqImwPlgBxztBTw7u0U,598
 cocoindex/validation.py,sha256=PZnJoby4sLbsmPv9fOjOQXuefjfZ7gmtsiTGU8SH-tc,3090
-cocoindex-0.1.75.dist-info/RECORD,,
+cocoindex-0.1.76.dist-info/RECORD,,