cocoindex 0.2.16__cp311-abi3-manylinux_2_28_aarch64.whl → 0.2.17__cp311-abi3-manylinux_2_28_aarch64.whl

cocoindex/functions/colpali.py

@@ -0,0 +1,250 @@
+"""ColPali image and query embedding functions for multimodal document retrieval."""
+
+import functools
+from dataclasses import dataclass
+from typing import Any, Optional, TYPE_CHECKING, Literal
+import numpy as np
+
+from .. import op
+from ..typing import Vector
+
+if TYPE_CHECKING:
+    import torch
+
+
+@dataclass
+class ColPaliModelInfo:
+    """Shared model information for ColPali embedding functions."""
+
+    model: Any
+    processor: Any
+    device: Any
+    dimension: int
+
+
+@functools.lru_cache(maxsize=None)
+def _get_colpali_model_and_processor(model_name: str) -> ColPaliModelInfo:
+    """Load and cache ColPali model and processor with shared device setup."""
+    try:
+        from colpali_engine import (  # type: ignore[import-untyped]
+            ColPali,
+            ColPaliProcessor,
+            ColQwen2,
+            ColQwen2Processor,
+            ColSmol,
+            ColSmolProcessor,
+        )
+        import torch
+    except ImportError as e:
+        raise ImportError(
+            "ColPali support requires the optional 'colpali' dependency. "
+            "Install it with: pip install 'cocoindex[colpali]'"
+        ) from e
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    # Determine model type from name
+    if "colpali" in model_name.lower():
+        model = ColPali.from_pretrained(
+            model_name, torch_dtype=torch.bfloat16, device_map=device
+        )
+        processor = ColPaliProcessor.from_pretrained(model_name)
+    elif "colqwen" in model_name.lower():
+        model = ColQwen2.from_pretrained(
+            model_name, torch_dtype=torch.bfloat16, device_map=device
+        )
+        processor = ColQwen2Processor.from_pretrained(model_name)
+    elif "colsmol" in model_name.lower():
+        model = ColSmol.from_pretrained(
+            model_name, torch_dtype=torch.bfloat16, device_map=device
+        )
+        processor = ColSmolProcessor.from_pretrained(model_name)
+    else:
+        # Fallback to ColPali for backwards compatibility
+        model = ColPali.from_pretrained(
+            model_name, torch_dtype=torch.bfloat16, device_map=device
+        )
+        processor = ColPaliProcessor.from_pretrained(model_name)
+
+    # Detect dimension
+    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 ColVision multimodal models.
+
+    Supports ALL models available in the colpali-engine library, including:
+    - ColPali models (colpali-*): PaliGemma-based, best for general document retrieval
+    - ColQwen2 models (colqwen-*): Qwen2-VL-based, excellent for multilingual text (29+ languages) and general vision
+    - ColSmol models (colsmol-*): Lightweight, good for resource-constrained environments
+    - Any future ColVision models supported by colpali-engine
+
+    These models use late interaction between image patch embeddings and text token
+    embeddings for retrieval.
+
+    Args:
+        model: Any ColVision model name supported by colpali-engine
+               (e.g., "vidore/colpali-v1.2", "vidore/colqwen2.5-v0.2", "vidore/colsmol-v1.0")
+               See https://github.com/illuin-tech/colpali for the complete list of supported models.
+
+    Note:
+        This function requires the optional colpali-engine dependency.
+        Install it with: pip install 'cocoindex[colpali]'
+    """
+
+    model: str
+
+
+@op.executor_class(
+    gpu=True,
+    cache=True,
+    behavior_version=1,
+)
+class ColPaliEmbedImageExecutor:
+    """Executor for ColVision image embedding (ColPali, ColQwen2, ColSmol, etc.)."""
+
+    spec: ColPaliEmbedImage
+    _model_info: ColPaliModelInfo
+
+    def analyze(self) -> 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 ColVision 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 ColVision multimodal models.
+
+    Supports ALL models available in the colpali-engine library, including:
+    - ColPali models (colpali-*): PaliGemma-based, best for general document retrieval
+    - ColQwen2 models (colqwen-*): Qwen2-VL-based, excellent for multilingual text (29+ languages) and general vision
+    - ColSmol models (colsmol-*): Lightweight, good for resource-constrained environments
+    - Any future ColVision models supported by colpali-engine
+
+    This produces query embeddings compatible with ColVision image embeddings
+    for late interaction scoring (MaxSim).
+
+    Args:
+        model: Any ColVision model name supported by colpali-engine
+               (e.g., "vidore/colpali-v1.2", "vidore/colqwen2.5-v0.2", "vidore/colsmol-v1.0")
+               See https://github.com/illuin-tech/colpali for the complete list of supported models.
+
+    Note:
+        This function requires the optional colpali-engine dependency.
+        Install it with: pip install 'cocoindex[colpali]'
+    """
+
+    model: str
+
+
+@op.executor_class(
+    gpu=True,
+    cache=True,
+    behavior_version=1,
+)
+class ColPaliEmbedQueryExecutor:
+    """Executor for ColVision query embedding (ColPali, ColQwen2, ColSmol, etc.)."""
+
+    spec: ColPaliEmbedQuery
+    _model_info: ColPaliModelInfo
+
+    def analyze(self) -> 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 ColVision 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/functions/sbert.py

@@ -0,0 +1,63 @@
+"""SentenceTransformer embedding functionality."""
+
+import dataclasses
+from typing import Any, Literal
+
+import numpy as np
+from numpy.typing import NDArray
+
+from .. import op
+from ..typing import Vector
+
+
+class SentenceTransformerEmbed(op.FunctionSpec):
+    """
+    `SentenceTransformerEmbed` embeds a text into a vector space using the [SentenceTransformer](https://huggingface.co/sentence-transformers) library.
+
+    Args:
+
+        model: The name of the SentenceTransformer model to use.
+        args: Additional arguments to pass to the SentenceTransformer constructor. e.g. {"trust_remote_code": True}
+
+    Note:
+        This function requires the optional sentence-transformers dependency.
+        Install it with: pip install 'cocoindex[embeddings]'
+    """
+
+    model: str
+    args: dict[str, Any] | None = None
+
+
+@op.executor_class(
+    gpu=True,
+    cache=True,
+    behavior_version=1,
+    arg_relationship=(op.ArgRelationship.EMBEDDING_ORIGIN_TEXT, "text"),
+)
+class SentenceTransformerEmbedExecutor:
+    """Executor for SentenceTransformerEmbed."""
+
+    spec: SentenceTransformerEmbed
+    _model: Any | None = None
+
+    def analyze(self) -> type:
+        try:
+            # Only import sentence_transformers locally when it's needed, as its import is very slow.
+            import sentence_transformers  # pylint: disable=import-outside-toplevel
+        except ImportError as e:
+            raise ImportError(
+                "sentence_transformers is required for SentenceTransformerEmbed function. "
+                "Install it with one of these commands:\n"
+                "  pip install 'cocoindex[embeddings]'\n"
+                "  pip install sentence-transformers"
+            ) from e
+
+        args = self.spec.args or {}
+        self._model = sentence_transformers.SentenceTransformer(self.spec.model, **args)
+        dim = self._model.get_sentence_embedding_dimension()
+        return Vector[np.float32, Literal[dim]]  # type: ignore
+
+    def __call__(self, text: str) -> NDArray[np.float32]:
+        assert self._model is not None
+        result: NDArray[np.float32] = self._model.encode(text, convert_to_numpy=True)
+        return result

cocoindex/lib.py

@@ -7,7 +7,7 @@ import warnings
 
 from . import _engine  # type: ignore
 from . import flow, setting
-from .convert import dump_engine_object
+from .engine_object import dump_engine_object
 from .validation import validate_app_namespace_name
 from typing import Any, Callable, overload
 

cocoindex/op.py

@@ -17,9 +17,8 @@ from typing import (
 
 from . import _engine  # type: ignore
 from .subprocess_exec import executor_stub
-from .convert import (
-    dump_engine_object,
-    load_engine_object,
+from .engine_object import dump_engine_object, load_engine_object
+from .engine_value import (
     make_engine_value_encoder,
     make_engine_value_decoder,
     make_engine_key_decoder,
@@ -47,6 +46,7 @@ class OpCategory(Enum):
     SOURCE = "source"
     TARGET = "target"
     DECLARATION = "declaration"
+    TARGET_ATTACHMENT = "target_attachment"
 
 
 @dataclass_transform()
@@ -82,6 +82,10 @@ class TargetSpec(metaclass=SpecMeta, category=OpCategory.TARGET):  # pylint: dis
     """A target spec. All its subclass can be instantiated similar to a dataclass, i.e. ClassName(field1=value1, field2=value2, ...)"""
 
 
+class TargetAttachmentSpec(metaclass=SpecMeta, category=OpCategory.TARGET_ATTACHMENT):  # pylint: disable=too-few-public-methods
+    """A target attachment spec. All its subclass can be instantiated similar to a dataclass, i.e. ClassName(field1=value1, field2=value2, ...)"""
+
+
 class DeclarationSpec(metaclass=SpecMeta, category=OpCategory.DECLARATION):  # pylint: disable=too-few-public-methods
     """A declaration spec. All its subclass can be instantiated similar to a dataclass, i.e. ClassName(field1=value1, field2=value2, ...)"""
 

cocoindex/sources/__init__.py

@@ -0,0 +1,5 @@
+"""
+Sources supported by CocoIndex.
+"""
+
+from ._engine_builtin_specs import *

cocoindex/{sources.py → sources/_engine_builtin_specs.py}

@@ -1,8 +1,8 @@
 """All builtin sources."""
 
-from . import op
-from .auth_registry import TransientAuthEntryReference
-from .setting import DatabaseConnectionSpec
+from .. import op
+from ..auth_registry import TransientAuthEntryReference
+from ..setting import DatabaseConnectionSpec
 from dataclasses import dataclass
 import datetime
 

cocoindex/targets/_engine_builtin_specs.py

@@ -14,6 +14,15 @@ class Postgres(op.TargetSpec):
 
     database: AuthEntryReference[DatabaseConnectionSpec] | None = None
     table_name: str | None = None
+    schema: str | None = None
+
+
+class PostgresSqlCommand(op.TargetAttachmentSpec):
+    """Attachment to execute specified SQL statements for Postgres targets."""
+
+    name: str
+    setup_sql: str
+    teardown_sql: str | None = None
 
 
 @dataclass