cocoindex 0.3.0__cp311-abi3-win_amd64.whl → 0.3.1__cp311-abi3-win_amd64.whl

cocoindex/functions/colpali.py

@@ -2,7 +2,7 @@
 
 import functools
 from dataclasses import dataclass
-from typing import Any, Optional, TYPE_CHECKING, Literal
+from typing import Any, TYPE_CHECKING, Literal
 import numpy as np
 
 from .. import op
@@ -22,18 +22,11 @@ class ColPaliModelInfo:
     dimension: int
 
 
-@functools.lru_cache(maxsize=None)
+@functools.cache
 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 colpali_engine as ce  # type: ignore[import-untyped]
         import torch
     except ImportError as e:
         raise ImportError(
@@ -42,29 +35,30 @@ def _get_colpali_model_and_processor(model_name: str) -> ColPaliModelInfo:
         ) from e
 
     device = "cuda" if torch.cuda.is_available() else "cpu"
+    lower_model_name = model_name.lower()
 
     # Determine model type from name
-    if "colpali" in model_name.lower():
-        model = ColPali.from_pretrained(
+    if lower_model_name.startswith("colpali"):
+        model = ce.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(
+        processor = ce.ColPaliProcessor.from_pretrained(model_name)
+    elif lower_model_name.startswith("colqwen2.5"):
+        model = ce.ColQwen2_5.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(
+        processor = ce.ColQwen2_5_Processor.from_pretrained(model_name)
+    elif lower_model_name.startswith("colqwen"):
+        model = ce.ColQwen2.from_pretrained(
             model_name, torch_dtype=torch.bfloat16, device_map=device
         )
-        processor = ColSmolProcessor.from_pretrained(model_name)
+        processor = ce.ColQwen2Processor.from_pretrained(model_name)
     else:
         # Fallback to ColPali for backwards compatibility
-        model = ColPali.from_pretrained(
+        model = ce.ColPali.from_pretrained(
             model_name, torch_dtype=torch.bfloat16, device_map=device
         )
-        processor = ColPaliProcessor.from_pretrained(model_name)
+        processor = ce.ColPaliProcessor.from_pretrained(model_name)
 
     # Detect dimension
     dimension = _detect_colpali_dimension(model, processor, device)
@@ -130,6 +124,7 @@ class ColPaliEmbedImage(op.FunctionSpec):
 @op.executor_class(
     gpu=True,
     cache=True,
+    batching=True,
     behavior_version=1,
 )
 class ColPaliEmbedImageExecutor:
@@ -146,7 +141,7 @@ class ColPaliEmbedImageExecutor:
         dimension = self._model_info.dimension
         return Vector[Vector[np.float32, Literal[dimension]]]  # type: ignore
 
-    def __call__(self, img_bytes: bytes) -> Any:
+    def __call__(self, img_bytes_list: list[bytes]) -> Any:
         try:
             from PIL import Image
             import torch
@@ -160,8 +155,11 @@ class ColPaliEmbedImageExecutor:
         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)
+        pil_images = [
+            Image.open(io.BytesIO(img_bytes)).convert("RGB")
+            for img_bytes in img_bytes_list
+        ]
+        inputs = processor.process_images(pil_images).to(device)
         with torch.no_grad():
             embeddings = model(**inputs)
 
@@ -171,10 +169,8 @@ class ColPaliEmbedImageExecutor:
                 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()
+        # [patches, hidden_dim]
+        return embeddings.cpu().to(torch.float32).numpy()
 
 
 class ColPaliEmbedQuery(op.FunctionSpec):
@@ -207,6 +203,7 @@ class ColPaliEmbedQuery(op.FunctionSpec):
     gpu=True,
     cache=True,
     behavior_version=1,
+    batching=True,
 )
 class ColPaliEmbedQueryExecutor:
     """Executor for ColVision query embedding (ColPali, ColQwen2, ColSmol, etc.)."""
@@ -222,7 +219,7 @@ class ColPaliEmbedQueryExecutor:
         dimension = self._model_info.dimension
         return Vector[Vector[np.float32, Literal[dimension]]]  # type: ignore
 
-    def __call__(self, query: str) -> Any:
+    def __call__(self, queries: list[str]) -> Any:
         try:
             import torch
         except ImportError as e:
@@ -234,7 +231,7 @@ class ColPaliEmbedQueryExecutor:
         processor = self._model_info.processor
         device = self._model_info.device
 
-        inputs = processor.process_queries([query]).to(device)
+        inputs = processor.process_queries(queries).to(device)
         with torch.no_grad():
             embeddings = model(**inputs)
 
@@ -244,7 +241,5 @@ class ColPaliEmbedQueryExecutor:
                 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()
+        # [tokens, hidden_dim]
+        return embeddings.cpu().to(torch.float32).numpy()

{cocoindex-0.3.0.dist-info → cocoindex-0.3.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cocoindex
-Version: 0.3.0
+Version: 0.3.1
 Classifier: Development Status :: 3 - Alpha
 Classifier: License :: OSI Approved :: Apache Software License
 Classifier: Operating System :: OS Independent

{cocoindex-0.3.0.dist-info → cocoindex-0.3.1.dist-info}/RECORD

@@ -1,9 +1,9 @@
-cocoindex-0.3.0.dist-info/METADATA,sha256=lqHIyy_Usb-cEtNh0cgLLusV2ZNL7DYzvtz1qX0KEQU,14426
-cocoindex-0.3.0.dist-info/WHEEL,sha256=QC3zdlgimDC1GtRrc0qfjqbzuc7G6nDvPrjaINbNOTw,95
-cocoindex-0.3.0.dist-info/entry_points.txt,sha256=_NretjYVzBdNTn7dK-zgwr7YfG2afz1u1uSE-5bZXF8,46
-cocoindex-0.3.0.dist-info/licenses/THIRD_PARTY_NOTICES.html,sha256=VRDb6qOsN808v5fkXRUNfGUqSvD_OJWjO6hD3uWfKFg,750831
+cocoindex-0.3.1.dist-info/METADATA,sha256=XVa358pgFtptG2qP5yJTbxgzqFe_1XjVm0_69HXchTg,14426
+cocoindex-0.3.1.dist-info/WHEEL,sha256=QC3zdlgimDC1GtRrc0qfjqbzuc7G6nDvPrjaINbNOTw,95
+cocoindex-0.3.1.dist-info/entry_points.txt,sha256=_NretjYVzBdNTn7dK-zgwr7YfG2afz1u1uSE-5bZXF8,46
+cocoindex-0.3.1.dist-info/licenses/THIRD_PARTY_NOTICES.html,sha256=_9FT9tv3L_NmA0GOtdkqRZGf_4o5A16ui5YvPPi5-WM,750831
 cocoindex/__init__.py,sha256=MsjYflfJHL_sKL4OxcExfRwl57JFKwqvt9AWbiHsZ3Q,2744
-cocoindex/_engine.pyd,sha256=iHuBUOawJoNhA3zjKNG9LYU2J_5rFpr6VfsZzUknXVM,77715456
+cocoindex/_engine.pyd,sha256=5SgV6FV9mGZGq6hw0cUg5I8f9lf9lWuGxf2CMwGAkq0,77837312
 cocoindex/auth_registry.py,sha256=HK1vfKQh_6z310c8kgFDIQf9RdoiA3vWUwvFYbgybr0,1384
 cocoindex/cli.py,sha256=1Q_c4FzW41FhlCrdJNNkD93Z4zP48hmFzsaU2FJ3wzc,25622
 cocoindex/engine_object.py,sha256=JH27f2MLNUw3HQy6JuCpvRhMxuKZBkxplLIo0m7PU04,10322
@@ -11,7 +11,7 @@ cocoindex/engine_value.py,sha256=CYS_rdJQZoLEcKVcUXupcH0u8k-6ce80Xrzd3rgUBJE,239
 cocoindex/flow.py,sha256=MznA23dsIOlu-1uEBe6XVTfYlDjvH3bvky1BCUEnEng,41357
 cocoindex/functions/__init__.py,sha256=AZ4f7dBVZMkWyR85z9Gy1AgOAB0f1tex62CMP8MDLX4,1061
 cocoindex/functions/_engine_builtin_specs.py,sha256=6ZV91MUosqGSolnKNUjyRHZ_oTOOsQ_jMMEBGZW1EYo,1852
-cocoindex/functions/colpali.py,sha256=kinPow46CSW0PW_-a5PT9JfsCxpvxZG4ZFzQ_OYt9GU,9115
+cocoindex/functions/colpali.py,sha256=Got7es-Em9KwWGwJAAXMJNFp1wFcbtEOj6vkV0cqSAA,8869
 cocoindex/functions/sbert.py,sha256=HPieCBNuKRaWRIBIH8hW4ACFtGN2rHpg6N7mfURpUcQ,2289
 cocoindex/index.py,sha256=C__LzwIC918VIDGsBsyLwvNBO-4BiC5Coq01Fp1zXkI,1032
 cocoindex/lib.py,sha256=cyKGdn8cfH9bkYfrnJ7dlUBO8OVZkKyrkYhWHsMFW_g,2365
@@ -39,4 +39,4 @@ cocoindex/typing.py,sha256=dC8CmltIbT85TAgiFXwmHT_lePdiHRQLrURjgZWF6oo,24716
 cocoindex/user_app_loader.py,sha256=ZkvUG9aJNNECAjwTY0ZYtNpFd9dNBPVoPKGTtB7dSZg,1926
 cocoindex/utils.py,sha256=U3W39zD2uZpXX8v84tJD7sRmbC5ar3z_ljAP1cJrYXI,618
 cocoindex/validation.py,sha256=4ZjsW-SZT8X_TEEhEE6QG6D-8Oq_TkPAhTqP0mdFYSE,3194
-cocoindex-0.3.0.dist-info/RECORD,,
+cocoindex-0.3.1.dist-info/RECORD,,

{cocoindex-0.3.0.dist-info → cocoindex-0.3.1.dist-info}/licenses/THIRD_PARTY_NOTICES.html

@@ -2846,7 +2846,7 @@ Software.
                 <h3 id="Apache-2.0">Apache License 2.0</h3>
                 <h4>Used by:</h4>
                 <ul class="license-used-by">
-                    <li><a href=" https://crates.io/crates/cocoindex ">cocoindex 0.3.0</a></li>
+                    <li><a href=" https://crates.io/crates/cocoindex ">cocoindex 0.3.1</a></li>
                     <li><a href=" https://github.com/awesomized/crc-fast-rust ">crc-fast 1.3.0</a></li>
                     <li><a href=" https://github.com/qdrant/rust-client ">qdrant-client 1.15.0</a></li>
                 </ul>