sie-server 0.3.2__tar.gz → 0.3.3__tar.gz

{sie_server-0.3.2 → sie_server-0.3.3}/Dockerfile.cpu

@@ -18,14 +18,15 @@ ENV DEBIAN_FRONTEND=noninteractive \
     UV_LINK_MODE=copy \
     PIP_DISABLE_PIP_VERSION_CHECK=1
 
-# build-essential + git are builder-only; they do NOT leak into the runtime stage.
+# build-essential + git + zlib1g-dev are builder-only; they do NOT leak into the runtime stage.
 RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
     --mount=type=cache,target=/var/lib/apt/lists,sharing=locked \
     apt-get update && apt-get install -y --no-install-recommends \
     build-essential \
     ca-certificates \
     curl \
-    git
+    git \
+    zlib1g-dev
 
 ARG UV_VERSION=0.9.28
 RUN curl -LsSf https://astral.sh/uv/${UV_VERSION}/install.sh | sh \
@@ -154,13 +155,16 @@ ENV DEBIAN_FRONTEND=noninteractive \
     OMP_NUM_THREADS=4 \
     MKL_NUM_THREADS=4
 
-# Only the shared libs torch + pillow actually dlopen at runtime.
+# Only the shared libs torch + pillow + rtree actually dlopen at runtime.
+# libspatialindex-c6: rtree (docling dep) dlopens libspatialindex_c.so; the
+# rtree==1.4.1 wheel only bundles the C++ core, not the C wrapper.
 RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
     --mount=type=cache,target=/var/lib/apt/lists,sharing=locked \
     apt-get update && apt-get install -y --no-install-recommends \
     libgomp1 \
     libjpeg62-turbo \
-    libpng16-16
+    libpng16-16 \
+    libspatialindex-c6
 
 RUN groupadd -g 1000 sie && useradd -u 1000 -g sie -m sie
 

{sie_server-0.3.2 → sie_server-0.3.3}/Dockerfile.cuda11

@@ -152,6 +152,8 @@ FROM nvidia/cuda:11.8.0-base-ubuntu22.04 AS runtime
 
 ENV DEBIAN_FRONTEND=noninteractive
 
+# libspatialindex-c6: rtree (docling dep) dlopens libspatialindex_c.so; the
+# rtree==1.4.1 wheel only bundles the C++ core, not the C wrapper.
 RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
     --mount=type=cache,target=/var/lib/apt/lists,sharing=locked \
     apt-get update && apt-get install -y --no-install-recommends \
@@ -159,7 +161,8 @@ RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
     gcc \
     libc6-dev \
     libgomp1 \
-    libnuma1
+    libnuma1 \
+    libspatialindex-c6
 
 RUN groupadd -g 1000 sie && useradd -u 1000 -g sie -m sie
 

{sie_server-0.3.2 → sie_server-0.3.3}/Dockerfile.cuda12

@@ -173,6 +173,8 @@ ENV DEBIAN_FRONTEND=noninteractive
 # libnuma1: required by sgl_kernel (SGLang bundle); import fails with a
 #   misleading SM-arch error without it.
 # libgomp1: torch OpenMP runtime.
+# libspatialindex-c6: rtree (docling dep) dlopens libspatialindex_c.so; the
+#   rtree==1.4.1 wheel only bundles the C++ core, not the C wrapper.
 RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
     --mount=type=cache,target=/var/lib/apt/lists,sharing=locked \
     apt-get update && apt-get install -y --no-install-recommends \
@@ -180,7 +182,8 @@ RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
     gcc \
     libc6-dev \
     libgomp1 \
-    libnuma1
+    libnuma1 \
+    libspatialindex-c6
 
 RUN groupadd -g 1000 sie && useradd -u 1000 -g sie -m sie
 

{sie_server-0.3.2 → sie_server-0.3.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sie-server
-Version: 0.3.2
+Version: 0.3.3
 Summary: Search Inference Engine - GPU inference server for search workloads
 License: Apache-2.0
 License-File: LICENSE

{sie_server-0.3.2 → sie_server-0.3.3}/bundles/default.yaml

@@ -12,6 +12,7 @@ adapters:
 - sie_server.adapters.colbert_rotary_flash
 - sie_server.adapters.colpali
 - sie_server.adapters.colqwen2
+- sie_server.adapters.colqwen3
 - sie_server.adapters.grounding_dino
 - sie_server.adapters.gte_sparse_flash
 - sie_server.adapters.jina_flash_cross_encoder
@@ -44,7 +45,7 @@ adapters:
 - sie_server.adapters.paddleocr_vl
 deps:
   # Most flash adapters; sentence_transformer needs >=4.57
-  transformers: '>=4.57'
+  transformers: '>=4.57,<5'
   # Pin torch — flash-attn's unpinned torch dep can pull in 2.10+ which breaks torchvision
   torch: '>=2.9,<2.10'
   # bge_m3, bge_m3_flash, nemo_colembed, nomic_flash, xlm_roberta_flash
@@ -56,7 +57,7 @@ deps:
   # bge_m3_flag
   FlagEmbedding: '>=1.3'
   # grounding_dino
-  pillow: ''
+  pillow: '>=11,<12'
   # grounding_dino
   requests: ''
   # gliner, gliner_bi

{sie_server-0.3.2 → sie_server-0.3.3}/bundles/transformers5.yaml

@@ -13,5 +13,8 @@ deps:
   transformers: '>=5.0'
   # Pin torch — flash-attn's unpinned torch dep can pull in 2.10+ which breaks torchvision
   torch: '>=2.9,<2.10'
-  huggingface-hub: '>=0.20'
+  # No huggingface-hub override: transformers 5.x's own metadata
+  # pins it to >=1.3,<2.0. An override here was previously
+  # unsatisfiable against that pin and broke `uv run` resolution.
+  # Let transformers drive the version.
   pillow: ''

{sie_server-0.3.2 → sie_server-0.3.3}/models/nvidia__llama-embed-nemotron-8b.yaml

@@ -18,10 +18,19 @@ profiles:
   default:
     max_batch_tokens: 8192
     compute_precision: bfloat16
-    adapter_path: sie_server.adapters.sglang:SGLangEmbeddingAdapter
+    # Was sie_server.adapters.sglang:SGLangEmbeddingAdapter, but
+    # SGLang's generic transformers fallback hits
+    # `assert get_embedding is False` for this Llama-based arch
+    # (no model-specific embedding implementation registered). Route
+    # through PyTorchEmbeddingAdapter like NV-Embed-v2 (also Llama-
+    # based 7B+ embedder); the heavy lane has 96 GiB so non-paged
+    # attention is fine. May be reconciled with a colleague's
+    # in-flight SGLang/arch fix later.
+    adapter_path: sie_server.adapters.pytorch_embedding:PyTorchEmbeddingAdapter
     adapter_options:
       loadtime:
-        mem_fraction_static: 0.85
+        trust_remote_code: true
+        attn_implementation: eager
       runtime:
         pooling: last_token
         normalize: true

sie_server-0.3.3/models/nvidia__nemotron-colembed-vl-4b-v2.yaml

@@ -0,0 +1,43 @@
+sie_id: nvidia/nemotron-colembed-vl-4b-v2
+hf_id: nvidia/nemotron-colembed-vl-4b-v2
+inputs:
+  text: true
+  image: true
+  audio: false
+  video: false
+tasks:
+  encode:
+    dense: null
+    sparse: null
+    multivector:
+      dim: 2560
+  score: null
+  extract: null
+max_sequence_length: 8192
+profiles:
+  default:
+    max_batch_tokens: 4096
+    compute_precision: bfloat16
+    adapter_path: sie_server.adapters.nemo_colembed:NemoColEmbedAdapter
+    adapter_options:
+      loadtime:
+        muvera_config:
+          num_repetitions: 40
+          num_simhash_projections: 6
+          projection_dim: null
+          final_projection_dim: 10240
+        token_dim: 2560
+      runtime:
+        normalize: true
+        batch_size: 4
+  muvera:
+    extends: default
+    adapter_options:
+      runtime:
+        normalize: true
+        batch_size: 4
+        output_types:
+        - dense
+        output_similarity:
+          dense: dot
+        muvera: {}

sie_server-0.3.3/models/tomoroai__tomoro-colqwen3-embed-4b.yaml

@@ -0,0 +1,42 @@
+sie_id: TomoroAI/tomoro-colqwen3-embed-4b
+hf_id: TomoroAI/tomoro-colqwen3-embed-4b
+inputs:
+  text: true
+  image: true
+  audio: false
+  video: false
+tasks:
+  encode:
+    dense: null
+    sparse: null
+    multivector:
+      dim: 320
+  score: null
+  extract: null
+max_sequence_length: 8192
+profiles:
+  default:
+    max_batch_tokens: 4096
+    compute_precision: bfloat16
+    adapter_path: sie_server.adapters.colqwen3:ColQwen3Adapter
+    adapter_options:
+      loadtime:
+        muvera_config:
+          num_repetitions: 40
+          num_simhash_projections: 6
+          projection_dim: null
+          final_projection_dim: 10240
+        token_dim: 320
+        trust_remote_code: true
+      runtime:
+        normalize: true
+  muvera:
+    extends: default
+    adapter_options:
+      runtime:
+        normalize: true
+        output_types:
+        - dense
+        output_similarity:
+          dense: dot
+        muvera: {}

{sie_server-0.3.2 → sie_server-0.3.3}/openapi.json

@@ -3,7 +3,7 @@
   "info": {
     "title": "SIE Server",
     "description": "Search Inference Engine - GPU inference server for search workloads",
-    "version": "0.3.2"
+    "version": "0.3.3"
   },
   "paths": {
     "/": {
@@ -1740,7 +1740,7 @@
               }
             ],
             "default": null,
-            "description": "Adapter-specific options",
+            "description": "Adapter-specific options. Recognized sub-keys include 'overflow_policy' (one of 'default', 'truncate_text', 'error'; default 'default') controlling how inputs exceeding the model's max_sequence_length are handled.",
             "title": "Options"
           }
         },
@@ -1763,6 +1763,23 @@
                 "location"
               ]
             }
+          },
+          {
+            "items": [
+              {
+                "text": "Apple Inc. was founded by Steve Jobs in Cupertino, California."
+              }
+            ],
+            "params": {
+              "labels": [
+                "person",
+                "organization",
+                "location"
+              ],
+              "options": {
+                "overflow_policy": "truncate_text"
+              }
+            }
           }
         ],
         "properties": {

{sie_server-0.3.2 → sie_server-0.3.3}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "sie-server"
-version = "0.3.2"
+version = "0.3.3"
 description = "Search Inference Engine - GPU inference server for search workloads"
 requires-python = ">=3.12,<3.13"
 license = { text = "Apache-2.0" }

sie_server-0.3.3/src/sie_server/adapters/colqwen3/__init__.py

@@ -0,0 +1,337 @@
+from __future__ import annotations
+
+import io
+import logging
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+import numpy as np
+import torch
+from torch.nn import functional as F
+
+from sie_server.adapters._base_adapter import BaseAdapter
+from sie_server.adapters._spec import AdapterSpec
+from sie_server.adapters._types import ComputePrecision
+from sie_server.core.inference_output import EncodeOutput
+
+if TYPE_CHECKING:
+    from PIL import Image as PILImage
+
+    from sie_server.types.inputs import Item
+
+logger = logging.getLogger(__name__)
+
+_ERR_NO_INPUT = "ColQwen3Adapter requires either text or images input"
+
+
+class ColQwen3Adapter(BaseAdapter):
+    """Adapter for ColQwen3-style visual document retrieval models.
+
+    ColQwen3 encodes document page images into multi-vector representations
+    (320-dim per token) for late interaction retrieval. Built on Qwen3-VL,
+    with a custom projection layer wrapper that exposes ``out.embeddings``.
+
+    Target model: ``TomoroAI/tomoro-colqwen3-embed-4b`` (4B params).
+
+    Loaded via ``AutoModel`` + ``AutoProcessor`` with ``trust_remote_code``
+    because the model ships its own ``ColQwen3`` / ``ColQwen3Processor``
+    classes (not in native transformers).
+    """
+
+    spec = AdapterSpec(
+        inputs=("text", "image"),
+        outputs=("multivector", "score"),
+        multivector_dim=320,
+        unload_fields=("_model", "_processor"),
+        default_preprocessor="image",
+    )
+
+    def __init__(
+        self,
+        model_name_or_path: str | Path,
+        *,
+        normalize: bool = True,
+        compute_precision: ComputePrecision = "bfloat16",
+        trust_remote_code: bool = True,
+        max_seq_length: int | None = None,
+        muvera_config: dict[str, Any] | None = None,
+        token_dim: int = 320,
+        max_num_visual_tokens: int = 1280,
+    ) -> None:
+        """Initialize the adapter.
+
+        Args:
+            model_name_or_path: HuggingFace model ID or local path.
+            normalize: Whether to L2-normalize embeddings (the model's wrapper
+                already normalizes; kept for interface parity).
+            compute_precision: Compute precision for inference.
+            trust_remote_code: Required for ColQwen3 (custom processor + model classes).
+            max_seq_length: Ignored — ColQwen3 uses dynamic sequence length.
+            muvera_config: Optional MUVERA configuration (passed to postprocessor).
+            token_dim: Per-token embedding dimension (320 for ColQwen3).
+            max_num_visual_tokens: Cap on visual tokens per image (passed to processor).
+        """
+        self._model_name_or_path = str(model_name_or_path)
+        self._normalize = normalize
+        self._compute_precision = compute_precision
+        self._trust_remote_code = trust_remote_code
+        self._max_num_visual_tokens = max_num_visual_tokens
+
+        self._model: Any = None
+        self._processor: Any = None
+        self._device: str | None = None
+        self._multivector_dim: int = token_dim
+
+    def load(self, device: str) -> None:
+        """Load processor + model onto the specified device."""
+        from transformers import AutoModel, AutoProcessor
+
+        self._device = device
+
+        dtype = self._resolve_dtype()
+        attn_impl = self._resolve_attn_implementation(device)
+
+        logger.info(
+            "Loading ColQwen3 model %s on device=%s with dtype=%s, attn=%s",
+            self._model_name_or_path,
+            device,
+            dtype,
+            attn_impl,
+        )
+
+        self._processor = AutoProcessor.from_pretrained(
+            self._model_name_or_path,
+            trust_remote_code=self._trust_remote_code,
+            max_num_visual_tokens=self._max_num_visual_tokens,
+        )
+
+        load_kwargs: dict[str, Any] = {
+            "trust_remote_code": self._trust_remote_code,
+            "device_map": device,
+            "dtype": dtype,
+        }
+        if attn_impl is not None:
+            load_kwargs["attn_implementation"] = attn_impl
+
+        self._model = AutoModel.from_pretrained(
+            self._model_name_or_path,
+            **load_kwargs,
+        ).eval()
+
+        # Discover token dim from the projection layer when present.
+        proj = getattr(self._model, "embedding_proj_layer", None)
+        out_features = getattr(proj, "out_features", None)
+        if isinstance(out_features, int) and out_features > 0:
+            self._multivector_dim = out_features
+
+    def _resolve_dtype(self) -> torch.dtype:
+        if not self._device or not str(self._device).startswith("cuda"):
+            return torch.float32
+        dtype_map = {
+            "float16": torch.float16,
+            "bfloat16": torch.bfloat16,
+            "float32": torch.float32,
+        }
+        return dtype_map.get(self._compute_precision, torch.bfloat16)
+
+    def _resolve_attn_implementation(self, device: str) -> str | None:
+        if not device.startswith("cuda"):
+            return None
+        try:
+            import flash_attn  # ty: ignore[unresolved-import]
+
+            return "flash_attention_2"
+        except ImportError:
+            logger.info("flash_attn not available, using sdpa attention")
+            return "sdpa"
+
+    # ------------------------------------------------------------------
+    # Encode
+    # ------------------------------------------------------------------
+
+    def encode(
+        self,
+        items: list[Item],
+        output_types: list[str],
+        *,
+        instruction: str | None = None,
+        is_query: bool = False,
+        prepared_items: Any = None,
+        options: dict[str, Any] | None = None,
+    ) -> EncodeOutput:
+        self._check_loaded()
+        self._validate_output_types(output_types)
+
+        if is_query:
+            multivector_list: list[np.ndarray] = []
+            for item in items:
+                if item.text is None:
+                    raise ValueError(_ERR_NO_INPUT)
+                multivector_list.append(self._encode_text(item.text))
+            return EncodeOutput(
+                multivector=multivector_list,
+                batch_size=len(items),
+                is_query=is_query,
+                multivector_token_dim=self._multivector_dim,
+            )
+
+        # Preallocate by index so output order matches input order regardless of
+        # text/image mix, and so multi-image items collapse to one multivector.
+        results: list[np.ndarray | None] = [None] * len(items)
+        all_images: list[PILImage.Image] = []
+        image_slots: list[tuple[int, int]] = []  # (item_idx, image_count)
+        for idx, item in enumerate(items):
+            has_images = item.images is not None and len(item.images) > 0
+            if has_images:
+                images = self._load_images(item)
+                all_images.extend(images)
+                image_slots.append((idx, len(images)))
+            elif item.text is not None:
+                results[idx] = self._encode_text(item.text)
+            else:
+                raise ValueError(_ERR_NO_INPUT)
+
+        if all_images:
+            per_image_mvs = self._encode_images(all_images)
+            cursor = 0
+            for idx, count in image_slots:
+                segment = per_image_mvs[cursor : cursor + count]
+                cursor += count
+                results[idx] = segment[0] if count == 1 else np.concatenate(segment, axis=0)
+
+        multivector_list = [mv for mv in results if mv is not None]
+        assert len(multivector_list) == len(items)
+
+        return EncodeOutput(
+            multivector=multivector_list,
+            batch_size=len(items),
+            is_query=is_query,
+            multivector_token_dim=self._multivector_dim,
+        )
+
+    # ------------------------------------------------------------------
+    # Image encoding
+    # ------------------------------------------------------------------
+
+    def _load_images(self, item: Any) -> list[PILImage.Image]:
+        from PIL import Image
+
+        pil_images: list[PILImage.Image] = []
+        for img_input in item.images or []:
+            pil_img = Image.open(io.BytesIO(img_input["data"]))
+            if pil_img.mode != "RGB":
+                pil_img = pil_img.convert("RGB")
+            pil_images.append(pil_img)
+        return pil_images
+
+    def _encode_images(self, images: list[PILImage.Image]) -> list[np.ndarray]:
+        """Encode a batch of images and return per-image multi-vectors."""
+        assert self._model is not None
+        assert self._processor is not None
+
+        inputs = self._processor(
+            images=images,
+            return_tensors="pt",
+            padding="longest",
+        )
+        inputs = {k: v.to(self._device) for k, v in inputs.items() if hasattr(v, "to")}
+
+        with torch.inference_mode():
+            outputs = self._model(**inputs)
+
+        # ColQwen3 returns a ModelOutput-like object with ``.embeddings``
+        # of shape (batch, seq, token_dim). The wrapper already L2-normalizes
+        # and applies attention-masking; our ``self._normalize`` is a no-op
+        # safety belt for downstream parity.
+        embeddings = outputs.embeddings
+        if self._normalize:
+            embeddings = F.normalize(embeddings, p=2, dim=-1)
+
+        results: list[np.ndarray] = [embeddings[i].float().cpu().numpy() for i in range(embeddings.shape[0])]
+
+        # Free GPU memory between batches to prevent OOM on subsequent calls
+        # (L4 22GB GPUs are tight for VLM models).
+        del outputs, embeddings, inputs
+        if self._device and self._device.startswith("cuda"):
+            torch.cuda.empty_cache()
+
+        return results
+
+    # ------------------------------------------------------------------
+    # Text encoding
+    # ------------------------------------------------------------------
+
+    def _encode_text(self, text: str) -> np.ndarray:
+        """Encode a single text query."""
+        assert self._model is not None
+        assert self._processor is not None
+
+        inputs = self._processor(
+            text=[text],
+            return_tensors="pt",
+            padding="longest",
+        )
+        inputs = {k: v.to(self._device) for k, v in inputs.items() if hasattr(v, "to")}
+
+        with torch.inference_mode():
+            outputs = self._model(**inputs)
+
+        embeddings = outputs.embeddings  # (1, seq, token_dim)
+        if self._normalize:
+            embeddings = F.normalize(embeddings, p=2, dim=-1)
+
+        result = embeddings[0].float().cpu().numpy()
+
+        del outputs, embeddings, inputs
+        if self._device and self._device.startswith("cuda"):
+            torch.cuda.empty_cache()
+
+        return result
+
+    # ------------------------------------------------------------------
+    # Scoring
+    # ------------------------------------------------------------------
+
+    def score(
+        self,
+        query: Any,
+        items: list[Any],
+        *,
+        instruction: str | None = None,
+        options: dict[str, Any] | None = None,
+    ) -> list[float]:
+        """Score documents against a text query using MaxSim."""
+        self._check_loaded()
+
+        query_output = self.encode([query], output_types=["multivector"], is_query=True)
+        if query_output.multivector is None:
+            raise RuntimeError("Failed to encode query: no multivector output")
+        query_vecs = query_output.multivector[0]
+
+        doc_output = self.encode(items, output_types=["multivector"], is_query=False)
+        if doc_output.multivector is None:
+            raise RuntimeError("Failed to encode documents: no multivector output")
+
+        scores: list[float] = []
+        query_tensor = torch.from_numpy(query_vecs).to(self._device)
+        for doc_vecs in doc_output.multivector:
+            doc_tensor = torch.from_numpy(doc_vecs).to(self._device)
+            sim = torch.matmul(query_tensor, doc_tensor.T)
+            scores.append(sim.max(dim=-1).values.sum().item())
+        return scores
+
+    # ------------------------------------------------------------------
+    # Helpers
+    # ------------------------------------------------------------------
+
+    def _validate_output_types(self, output_types: list[str]) -> None:
+        unsupported = set(output_types) - {"multivector"}
+        if unsupported:
+            msg = f"Unsupported output types: {unsupported}. ColQwen3Adapter only supports 'multivector'."
+            raise ValueError(msg)
+
+    def get_preprocessor(self) -> Any | None:
+        # ColQwen3 uses a custom processor that handles both text and images
+        # internally via the ColQwen3Processor; the generic ImagePreprocessor
+        # does not match the (text-only / image-only) call pattern.
+        return None

sie_server-0.3.3/src/sie_server/adapters/errors.py

@@ -0,0 +1,2 @@
+class InputTooLongError(ValueError):
+    pass