cocoindex 0.1.75__cp313-cp313-manylinux_2_28_x86_64.whl → 0.1.77__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/op.py

@@ -114,8 +114,8 @@ class _FunctionExecutorFactory:
     ) -> tuple[dict[str, Any], Executor]:
         spec = _load_spec_from_engine(self._spec_cls, spec)
         executor = self._executor_cls(spec)
-        result_type = executor.analyze(*args, **kwargs)
-        return (encode_enriched_type(result_type), executor)
+        result_type = executor.analyze_schema(*args, **kwargs)
+        return (result_type, executor)
 
 
 _gpu_dispatch_lock = asyncio.Lock()
@@ -156,6 +156,12 @@ def _to_async_call(call: Callable[..., Any]) -> Callable[..., Awaitable[Any]]:
     return lambda *args, **kwargs: asyncio.to_thread(lambda: call(*args, **kwargs))
 
 
+@dataclasses.dataclass
+class _ArgInfo:
+    decoder: Callable[[Any], Any]
+    is_required: bool
+
+
 def _register_op_factory(
     category: OpCategory,
     expected_args: list[tuple[str, inspect.Parameter]],
@@ -176,8 +182,8 @@ def _register_op_factory(
             return op_args.behavior_version
 
     class _WrappedClass(executor_cls, _Fallback):  # type: ignore[misc]
-        _args_decoders: list[Callable[[Any], Any]]
-        _kwargs_decoders: dict[str, Callable[[Any], Any]]
+        _args_info: list[_ArgInfo]
+        _kwargs_info: dict[str, _ArgInfo]
         _acall: Callable[..., Awaitable[Any]]
 
         def __init__(self, spec: Any) -> None:
@@ -185,28 +191,45 @@ def _register_op_factory(
             self.spec = spec
             self._acall = _to_async_call(super().__call__)
 
-        def analyze(
+        def analyze_schema(
             self, *args: _engine.OpArgSchema, **kwargs: _engine.OpArgSchema
         ) -> Any:
             """
             Analyze the spec and arguments. In this phase, argument types should be validated.
             It should return the expected result type for the current op.
             """
-            self._args_decoders = []
-            self._kwargs_decoders = {}
+            self._args_info = []
+            self._kwargs_info = {}
             attributes = []
-
-            def process_attribute(arg_name: str, arg: _engine.OpArgSchema) -> None:
+            potentially_missing_required_arg = False
+
+            def process_arg(
+                arg_name: str,
+                arg_param: inspect.Parameter,
+                actual_arg: _engine.OpArgSchema,
+            ) -> _ArgInfo:
+                nonlocal potentially_missing_required_arg
                 if op_args.arg_relationship is not None:
                     related_attr, related_arg_name = op_args.arg_relationship
                     if related_arg_name == arg_name:
                         attributes.append(
-                            TypeAttr(related_attr.value, arg.analyzed_value)
+                            TypeAttr(related_attr.value, actual_arg.analyzed_value)
                         )
+                type_info = analyze_type_info(arg_param.annotation)
+                decoder = make_engine_value_decoder(
+                    [arg_name], actual_arg.value_type["type"], type_info
+                )
+                is_required = not type_info.nullable
+                if is_required and actual_arg.value_type.get("nullable", False):
+                    potentially_missing_required_arg = True
+                return _ArgInfo(
+                    decoder=decoder,
+                    is_required=is_required,
+                )
 
             # Match arguments with parameters.
             next_param_idx = 0
-            for arg in args:
+            for actual_arg in args:
                 if next_param_idx >= len(expected_args):
                     raise ValueError(
                         f"Too many arguments passed in: {len(args)} > {len(expected_args)}"
@@ -219,20 +242,13 @@ def _register_op_factory(
                     raise ValueError(
                         f"Too many positional arguments passed in: {len(args)} > {next_param_idx}"
                     )
-                self._args_decoders.append(
-                    make_engine_value_decoder(
-                        [arg_name],
-                        arg.value_type["type"],
-                        analyze_type_info(arg_param.annotation),
-                    )
-                )
-                process_attribute(arg_name, arg)
+                self._args_info.append(process_arg(arg_name, arg_param, actual_arg))
                 if arg_param.kind != inspect.Parameter.VAR_POSITIONAL:
                     next_param_idx += 1
 
             expected_kwargs = expected_args[next_param_idx:]
 
-            for kwarg_name, kwarg in kwargs.items():
+            for kwarg_name, actual_arg in kwargs.items():
                 expected_arg = next(
                     (
                         arg
@@ -254,12 +270,9 @@ def _register_op_factory(
                         f"Unexpected keyword argument passed in: {kwarg_name}"
                     )
                 arg_param = expected_arg[1]
-                self._kwargs_decoders[kwarg_name] = make_engine_value_decoder(
-                    [kwarg_name],
-                    kwarg.value_type["type"],
-                    analyze_type_info(arg_param.annotation),
+                self._kwargs_info[kwarg_name] = process_arg(
+                    kwarg_name, arg_param, actual_arg
                 )
-                process_attribute(kwarg_name, kwarg)
 
             missing_args = [
                 name
@@ -280,32 +293,45 @@ def _register_op_factory(
             if len(missing_args) > 0:
                 raise ValueError(f"Missing arguments: {', '.join(missing_args)}")
 
-            prepare_method = getattr(executor_cls, "analyze", None)
-            if prepare_method is not None:
-                result = prepare_method(self, *args, **kwargs)
+            base_analyze_method = getattr(self, "analyze", None)
+            if base_analyze_method is not None:
+                result = base_analyze_method(*args, **kwargs)
             else:
                 result = expected_return
             if len(attributes) > 0:
                 result = Annotated[result, *attributes]
-            return result
+
+            encoded_type = encode_enriched_type(result)
+            if potentially_missing_required_arg:
+                encoded_type["nullable"] = True
+            return encoded_type
 
         async def prepare(self) -> None:
             """
             Prepare for execution.
             It's executed after `analyze` and before any `__call__` execution.
             """
-            setup_method = getattr(super(), "prepare", None)
-            if setup_method is not None:
-                await _to_async_call(setup_method)()
+            prepare_method = getattr(super(), "prepare", None)
+            if prepare_method is not None:
+                await _to_async_call(prepare_method)()
 
         async def __call__(self, *args: Any, **kwargs: Any) -> Any:
-            decoded_args = (
-                decoder(arg) for decoder, arg in zip(self._args_decoders, args)
-            )
-            decoded_kwargs = {
-                arg_name: self._kwargs_decoders[arg_name](arg)
-                for arg_name, arg in kwargs.items()
-            }
+            decoded_args = []
+            for arg_info, arg in zip(self._args_info, args):
+                if arg_info.is_required and arg is None:
+                    return None
+                decoded_args.append(arg_info.decoder(arg))
+
+            decoded_kwargs = {}
+            for kwarg_name, arg in kwargs.items():
+                kwarg_info = self._kwargs_info.get(kwarg_name)
+                if kwarg_info is None:
+                    raise ValueError(
+                        f"Unexpected keyword argument passed in: {kwarg_name}"
+                    )
+                if kwarg_info.is_required and arg is None:
+                    return None
+                decoded_kwargs[kwarg_name] = kwarg_info.decoder(arg)
 
             if op_args.gpu:
                 # For GPU executions, data-level parallelism is applied, so we don't want to

cocoindex/tests/test_transform_flow.py

@@ -101,3 +101,44 @@ async def test_for_each_transform_flow_async() -> None:
     }
 
     assert result == expected, f"Expected {expected}, got {result}"
+
+
+def test_none_arg_yield_none_result() -> None:
+    """Test that None arguments yield None results."""
+
+    @cocoindex.op.function()
+    def custom_fn(
+        required_arg: int,
+        optional_arg: int | None,
+        required_kwarg: int,
+        optional_kwarg: int | None,
+    ) -> int:
+        return (
+            required_arg + (optional_arg or 0) + required_kwarg + (optional_kwarg or 0)
+        )
+
+    @cocoindex.transform_flow()
+    def transform_flow(
+        required_arg: cocoindex.DataSlice[int | None],
+        optional_arg: cocoindex.DataSlice[int | None],
+        required_kwarg: cocoindex.DataSlice[int | None],
+        optional_kwarg: cocoindex.DataSlice[int | None],
+    ) -> cocoindex.DataSlice[int | None]:
+        return required_arg.transform(
+            custom_fn,
+            optional_arg,
+            required_kwarg=required_kwarg,
+            optional_kwarg=optional_kwarg,
+        )
+
+    result = transform_flow.eval(1, 2, 4, 8)
+    assert result == 15, f"Expected 15, got {result}"
+
+    result = transform_flow.eval(1, None, 4, None)
+    assert result == 5, f"Expected 5, got {result}"
+
+    result = transform_flow.eval(None, 2, 4, 8)
+    assert result is None, f"Expected None, got {result}"
+
+    result = transform_flow.eval(1, 2, None, None)
+    assert result is None, f"Expected None, got {result}"

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cocoindex
-Version: 0.1.75
+Version: 0.1.77
 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.
@@ -211,6 +214,7 @@ It defines an index flow like this:
 | [Image Search with Vision API](examples/image_search) | Generates detailed captions for images using a vision model, embeds them, enables live-updating semantic search via FastAPI and served on a React frontend|
 | [Face Recognition](examples/face_recognition) | Recognize faces in images and build embedding index |
 | [Paper Metadata](examples/paper_metadata) | Index papers in PDF files, and build metadata tables for each paper |
+| [Multi Format Indexing](examples/multi_format_indexing) | Build visual document index from PDFs and images with ColPali for semantic search |
 | [Custom Output Files](examples/custom_output_files) | Convert markdown files to HTML files and save them to a local directory, using *CocoIndex Custom Targets* |
 
 More coming and stay tuned 👀!

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

@@ -1,18 +1,18 @@
-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.77.dist-info/METADATA,sha256=009CbNRu3D4ifri3AGpphTt29ENF0F1HNeu9O1HkVXI,11799
+cocoindex-0.1.77.dist-info/WHEEL,sha256=9Ee4MwqZpMDLH1_kZE8rvruLKRVRs9cmbXRSBB0h-_M,108
+cocoindex-0.1.77.dist-info/entry_points.txt,sha256=_NretjYVzBdNTn7dK-zgwr7YfG2afz1u1uSE-5bZXF8,46
+cocoindex-0.1.77.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=E6gZzpSdni8nWN4QiOKbccBqohglMKSjc8kmUz5pCng,71118256
 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
-cocoindex/op.py,sha256=oiG1rjxz6ad1jGS7DMya4NStrA_6LV3RbcVSR75XUl0,21516
+cocoindex/op.py,sha256=EZ3QrDGJciYAxsqBeZOc3SstOL_34a6cke3WXWwz4-M,22538
 cocoindex/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 cocoindex/runtime.py,sha256=povilB3HH3y1JF-yxKwU-pD8n2WnAqyQxIgvXXHNc60,1080
 cocoindex/setting.py,sha256=TwhQ6pEeZmvc8ZXlnT9d8Wn8Vz_u7Z5LJUkGsKmKSno,4859
@@ -22,10 +22,10 @@ cocoindex/targets.py,sha256=Nfh_tpFd1goTnS_cxBjIs4j9zl3Z4Z1JomAQ1dl3Sic,2796
 cocoindex/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 cocoindex/tests/test_convert.py,sha256=l7LqD7duV9-xkYTaKOsEPdqw7v14dUzE40f4VVLlBCQ,49423
 cocoindex/tests/test_optional_database.py,sha256=snAmkNa6wtOSaxoZE1HgjvL5v_ylitt3Jt_9df4Cgdc,8506
-cocoindex/tests/test_transform_flow.py,sha256=VvT5b895MH5kwT-h4OpdDTl545SU4nxeIm7E_QANmAk,2894
+cocoindex/tests/test_transform_flow.py,sha256=M63Noo29Kl5CawTosCeZHjkn4ObjTP2ycHCGO5p8WHM,4230
 cocoindex/tests/test_typing.py,sha256=9OF3lO2uSpZBefkEJx7WRbnkXjwQtvlQIeeARYQID68,12391
 cocoindex/tests/test_validation.py,sha256=X6AQzVs-hVKIXcrHMEMQnhfUE8at7iXQnPq8nHNhZ2Q,4543
 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.77.dist-info/RECORD,,