Optimize ColPali with functools.cache and add colpali feature

  - Use @functools.cache for model caching instead of manual dict
  - Add 'colpali' optional dependency separate from 'embeddings'
  - Fix dimension detection and LAN access for frontend

Author: badmonster0

examples/image_search/README.md

@@ -62,7 +62,7 @@ export OLLAMA_MODEL="gemma3"  # Optional, for caption generation
 - Install dependencies:
   ```
   pip install -e .
-  pip install 'cocoindex[embeddings]'  # Adds ColPali and sentence-transformers support
+  pip install 'cocoindex[colpali]'  # Adds ColPali support
   ```
 
 - Configure model (optional):

examples/image_search/frontend/src/App.jsx

@@ -1,6 +1,6 @@
 import React, { useState } from 'react';
 
-const API_URL = 'http://localhost:8000/search'; // Adjust this to your backend search endpoint
+const API_URL = `http://${window.location.hostname}:8000/search`;
 
 export default function App() {
   const [query, setQuery] = useState('');
@@ -42,7 +42,7 @@ export default function App() {
         {results.length === 0 && !loading && <div>No results</div>}
         {results.map((result, idx) => (
           <div key={idx} className="result-card">
-            <img src={`http://localhost:8000/img/${result.filename}`} alt={result.filename} className="result-img" />
+            <img src={`http://${window.location.hostname}:8000/img/${result.filename}`} alt={result.filename} className="result-img" />
             <div className="score">Score: {result.score?.toFixed(3)}</div>
           </div>
         ))}

pyproject.toml

@@ -31,8 +31,12 @@ features = ["pyo3/extension-module"]
 [project.optional-dependencies]
 dev = ["pytest", "pytest-asyncio", "ruff", "mypy", "pre-commit"]
 
-embeddings = ["sentence-transformers>=3.3.1", "colpali-engine", "fastembed"]
-all = ["sentence-transformers>=3.3.1", "colpali-engine", "fastembed"]
+embeddings = ["sentence-transformers>=3.3.1"]
+colpali = ["colpali-engine"]
+
+# We need to repeat the dependency above to make it available for the `all` feature.
+# Indirect dependencies such as "cocoindex[embeddings]" will not work for local development.
+all = ["sentence-transformers>=3.3.1", "colpali-engine"]
 
 [tool.mypy]
 python_version = "3.11"

python/cocoindex/functions.py

@@ -1,6 +1,7 @@
 """All builtin functions."""
 
 import dataclasses
+import functools
 from typing import Annotated, Any, Literal
 
 import numpy as np
@@ -101,69 +102,51 @@ def __call__(self, text: str) -> NDArray[np.float32]:
         return result
 
 
-# Global ColPali model cache to avoid reloading models
-_COLPALI_MODEL_CACHE = {}
-
-
-class _ColPaliModelManager:
-    """Shared model manager for ColPali models to avoid duplicate loading."""
-
-    @staticmethod
-    def get_model_and_processor(model_name: str) -> dict[str, Any]:
-        """Get or load ColPali model and processor, with caching."""
-        if model_name not in _COLPALI_MODEL_CACHE:
-            try:
-                from colpali_engine.models import ColPali, ColPaliProcessor  # type: ignore[import-untyped]
-            except ImportError as e:
-                raise ImportError(
-                    "ColPali is not available. Make sure cocoindex is installed with ColPali support."
-                ) from e
-
-            model = ColPali.from_pretrained(model_name)
-            processor = ColPaliProcessor.from_pretrained(model_name)
-
-            # Get dimension from FastEmbed API
-            dimension = _ColPaliModelManager._detect_dimension()
-
-            _COLPALI_MODEL_CACHE[model_name] = {
-                "model": model,
-                "processor": processor,
-                "dimension": dimension,
-            }
-
-        return _COLPALI_MODEL_CACHE[model_name]
-
-    @staticmethod
-    def _detect_dimension() -> int:
-        """Detect ColPali embedding dimension using FastEmbed API."""
-        try:
-            from fastembed import LateInteractionMultimodalEmbedding
-
-            # Use the standard FastEmbed ColPali model for dimension detection
-            standard_colpali_model = "Qdrant/colpali-v1.3-fp16"
-
-            supported_models = (
-                LateInteractionMultimodalEmbedding.list_supported_models()
-            )
-            for supported_model in supported_models:
-                if supported_model["model"] == standard_colpali_model:
-                    dim = supported_model["dim"]
-                    if isinstance(dim, int):
-                        return dim
-                    else:
-                        raise ValueError(
-                            f"Expected integer dimension, got {type(dim)}: {dim}"
-                        )
-
-            raise ValueError(
-                f"Could not find dimension for ColPali model in FastEmbed supported models"
-            )
-
-        except ImportError:
-            raise ImportError(
-                "FastEmbed is required for ColPali dimension detection. "
-                "Install it with: pip install fastembed"
-            )
+@functools.cache
+def _get_colpali_model_and_processor(model_name: str) -> dict[str, Any]:
+    """Get or load ColPali model and processor, with caching."""
+    try:
+        from colpali_engine.models import ColPali, ColPaliProcessor  # type: ignore[import-untyped]
+    except ImportError as e:
+        raise ImportError(
+            "ColPali is not available. Make sure cocoindex is installed with ColPali support."
+        ) from e
+
+    model = ColPali.from_pretrained(model_name)
+    processor = ColPaliProcessor.from_pretrained(model_name)
+
+    # Get dimension from the actual model
+    dimension = _detect_colpali_dimension(model, processor)
+
+    return {
+        "model": model,
+        "processor": processor,
+        "dimension": dimension,
+    }
+
+
+def _detect_colpali_dimension(model: Any, processor: 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])
+        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):
@@ -198,9 +181,7 @@ class ColPaliEmbedImageExecutor:
     def analyze(self, _img_bytes: Any) -> type:
         # Get shared model and dimension
         if self._cached_model_data is None:
-            self._cached_model_data = _ColPaliModelManager.get_model_and_processor(
-                self.spec.model
-            )
+            self._cached_model_data = _get_colpali_model_and_processor(self.spec.model)
 
         # Return multi-vector type: Variable patches x Fixed hidden dimension
         dimension = self._cached_model_data["dimension"]
@@ -218,9 +199,7 @@ def __call__(self, img_bytes: bytes) -> list[list[float]]:
 
         # Get shared model and processor
         if self._cached_model_data is None:
-            self._cached_model_data = _ColPaliModelManager.get_model_and_processor(
-                self.spec.model
-            )
+            self._cached_model_data = _get_colpali_model_and_processor(self.spec.model)
 
         model = self._cached_model_data["model"]
         processor = self._cached_model_data["processor"]
@@ -279,9 +258,7 @@ class ColPaliEmbedQueryExecutor:
     def analyze(self, _query: Any) -> type:
         # Get shared model and dimension
         if self._cached_model_data is None:
-            self._cached_model_data = _ColPaliModelManager.get_model_and_processor(
-                self.spec.model
-            )
+            self._cached_model_data = _get_colpali_model_and_processor(self.spec.model)
 
         # Return multi-vector type: Variable tokens x Fixed hidden dimension
         dimension = self._cached_model_data["dimension"]
@@ -297,9 +274,7 @@ def __call__(self, query: str) -> list[list[float]]:
 
         # Get shared model and processor
         if self._cached_model_data is None:
-            self._cached_model_data = _ColPaliModelManager.get_model_and_processor(
-                self.spec.model
-            )
+            self._cached_model_data = _get_colpali_model_and_processor(self.spec.model)
 
         model = self._cached_model_data["model"]
         processor = self._cached_model_data["processor"]