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

{sie_server-0.3.1 → sie_server-0.3.2}/.gitignore

@@ -16,6 +16,9 @@ eggs/
 .eggs/
 lib/
 lib64/
+# JS/TS projects under tools/ legitimately use a `lib/` directory.
+!tools/*/lib/
+!tools/*/lib/**
 parts/
 sdist/
 var/

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

@@ -39,7 +39,10 @@ COPY packages/sie_sdk/pyproject.toml /tmp/sie_sdk/pyproject.toml
 COPY packages/sie_server/pyproject.toml ./pyproject.toml
 
 # Stub source trees so pip accepts the editable installs during dep resolution.
-RUN mkdir -p /tmp/sie_sdk/src/sie_sdk src/sie_server \
+# Also create empty bundles/ and models/ — referenced by force-include in
+# pyproject.toml; hatchling resolves them at editable-metadata time even though
+# real contents only land in the `base` stage.
+RUN mkdir -p /tmp/sie_sdk/src/sie_sdk src/sie_server bundles models \
     && touch /tmp/sie_sdk/src/sie_sdk/__init__.py \
     && touch src/sie_server/__init__.py
 

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

@@ -41,7 +41,10 @@ WORKDIR /app
 COPY packages/sie_sdk/pyproject.toml /tmp/sie_sdk/pyproject.toml
 COPY packages/sie_server/pyproject.toml ./pyproject.toml
 
-RUN mkdir -p /tmp/sie_sdk/src/sie_sdk src/sie_server \
+# Stubs for editable metadata generation — bundles/ and models/ are referenced
+# by force-include in pyproject.toml and must exist; real contents are copied
+# in the base stage.
+RUN mkdir -p /tmp/sie_sdk/src/sie_sdk src/sie_server bundles models \
     && touch /tmp/sie_sdk/src/sie_sdk/__init__.py \
     && touch src/sie_server/__init__.py
 

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

@@ -44,7 +44,10 @@ WORKDIR /app
 COPY packages/sie_sdk/pyproject.toml /tmp/sie_sdk/pyproject.toml
 COPY packages/sie_server/pyproject.toml ./pyproject.toml
 
-RUN mkdir -p /tmp/sie_sdk/src/sie_sdk src/sie_server \
+# Stubs for editable metadata generation — bundles/ and models/ are referenced
+# by force-include in pyproject.toml and must exist; real contents are copied
+# in the base stage.
+RUN mkdir -p /tmp/sie_sdk/src/sie_sdk src/sie_server bundles models \
     && touch /tmp/sie_sdk/src/sie_sdk/__init__.py \
     && touch src/sie_server/__init__.py
 

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

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

{sie_server-0.3.1 → sie_server-0.3.2}/models/answerdotai__answerai-colbert-small-v1.yaml

@@ -11,7 +11,7 @@ tasks:
     sparse: null
     multivector:
       dim: 96
-  score: null
+  score: {}
   extract: null
 max_sequence_length: 512
 profiles:

{sie_server-0.3.1 → sie_server-0.3.2}/models/colbert-ir__colbertv2.0.yaml

@@ -11,7 +11,7 @@ tasks:
     sparse: null
     multivector:
       dim: 128
-  score: null
+  score: {}
   extract: null
 max_sequence_length: 512
 profiles:

{sie_server-0.3.1 → sie_server-0.3.2}/models/docling.yaml

@@ -18,3 +18,11 @@ profiles:
     adapter_options:
       loadtime: {}
       runtime: {}
+  ocr:
+    max_batch_tokens: 1
+    compute_precision: null
+    adapter_path: sie_server.adapters.docling:DoclingAdapter
+    adapter_options:
+      loadtime: {}
+      runtime:
+        ocr: true

{sie_server-0.3.1 → sie_server-0.3.2}/models/jinaai__jina-colbert-v2.yaml

@@ -11,7 +11,7 @@ tasks:
     sparse: null
     multivector:
       dim: 128
-  score: null
+  score: {}
   extract: null
 max_sequence_length: 8192
 profiles:

{sie_server-0.3.1 → sie_server-0.3.2}/models/lightonai__GTE-ModernColBERT-v1.yaml

@@ -11,7 +11,7 @@ tasks:
     sparse: null
     multivector:
       dim: 128
-  score: null
+  score: {}
   extract: null
 max_sequence_length: 8192
 profiles:

{sie_server-0.3.1 → sie_server-0.3.2}/models/lightonai__Reason-ModernColBERT.yaml

@@ -11,7 +11,7 @@ tasks:
     sparse: null
     multivector:
       dim: 128
-  score: null
+  score: {}
   extract: null
 max_sequence_length: 8192
 profiles:

{sie_server-0.3.1 → sie_server-0.3.2}/models/mixedbread-ai__mxbai-colbert-large-v1.yaml

@@ -11,7 +11,7 @@ tasks:
     sparse: null
     multivector:
       dim: 128
-  score: null
+  score: {}
   extract: null
 max_sequence_length: 512
 profiles:

{sie_server-0.3.1 → sie_server-0.3.2}/models/mixedbread-ai__mxbai-edge-colbert-v0-32m.yaml

@@ -11,7 +11,7 @@ tasks:
     sparse: null
     multivector:
       dim: 64
-  score: null
+  score: {}
   extract: null
 max_sequence_length: 8192
 profiles:

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

@@ -3,7 +3,7 @@
   "info": {
     "title": "SIE Server",
     "description": "Search Inference Engine - GPU inference server for search workloads",
-    "version": "0.3.1"
+    "version": "0.3.2"
   },
   "paths": {
     "/": {

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

@@ -1,6 +1,6 @@
 [project]
 name = "sie-server"
-version = "0.3.1"
+version = "0.3.2"
 description = "Search Inference Engine - GPU inference server for search workloads"
 requires-python = ">=3.12,<3.13"
 license = { text = "Apache-2.0" }
@@ -80,6 +80,10 @@ build-backend = "hatchling.build"
 [tool.hatch.build.targets.wheel]
 packages = ["src/sie_server"]
 
+[tool.hatch.build.targets.wheel.force-include]
+"models" = "sie_server/models"
+"bundles" = "sie_server/bundles"
+
 [tool.uv.sources]
 # Prebuilt flash-attn wheel for torch 2.9 + cu128 (official wheels only go up to torch 2.8)
 # Platform-specific: Linux x86_64 only. Non-Linux users should not install the flash-attn extra.

{sie_server-0.3.1 → sie_server-0.3.2}/src/sie_server/adapters/_base_adapter.py

@@ -6,11 +6,15 @@ from typing import TYPE_CHECKING, Any, ClassVar, cast
 
 from sie_server.adapters._spec import AdapterSpec
 from sie_server.adapters._types import ERR_NOT_LOADED
+from sie_server.adapters._utils import grouped_score_pairs
 from sie_server.adapters.base import ModelAdapter, ModelCapabilities, ModelDims
 
 if TYPE_CHECKING:
     import torch
 
+    from sie_server.core.inference_output import ScoreOutput
+    from sie_server.types.inputs import Item
+
 logger = logging.getLogger(__name__)
 
 
@@ -57,7 +61,16 @@ class BaseAdapter(ModelAdapter):
             raise TypeError(msg)
 
         if "score" in spec.outputs:
-            if cls.score is ModelAdapter.score and cls.score_pairs is ModelAdapter.score_pairs:
+            # BaseAdapter ships a default score_pairs() that delegates to score().
+            # Treat that default as "not implemented" for validation purposes:
+            # subclasses must override either score() or score_pairs() so the
+            # default delegate doesn't bottom out in ModelAdapter.score().
+            score_overridden = cls.score is not ModelAdapter.score
+            score_pairs_overridden = cls.score_pairs not in (
+                ModelAdapter.score_pairs,
+                BaseAdapter.score_pairs,
+            )
+            if not score_overridden and not score_pairs_overridden:
                 msg = f"{cls.__name__} declares 'score' in outputs but does not implement score() or score_pairs()"
                 raise TypeError(msg)
 
@@ -117,6 +130,41 @@ class BaseAdapter(ModelAdapter):
             model_name=getattr(self, "_model_name_or_path", ""),
         )
 
+    # -- Default batched scoring ---------------------------------------------
+
+    def score_pairs(
+        self,
+        queries: list[Item],
+        docs: list[Item],
+        *,
+        instruction: str | None = None,
+        options: dict[str, Any] | None = None,
+    ) -> ScoreOutput:
+        """Default ``score_pairs()`` that batches via per-query grouping.
+
+        Groups parallel ``(query, doc)`` pairs by ``(text, id, instruction)``
+        so each unique query is encoded once and its docs are scored as a
+        single ``score()`` call. Subclasses with a more efficient native
+        cross-batch path (e.g. cross-encoders that pack queries and docs
+        into one transformer pass) should override this.
+
+        Per-call ``options`` are not supported by this default delegate
+        (it dispatches per-query and cannot route options into ``score()``
+        without subclass-specific knowledge). If ``options`` is a non-empty
+        mapping, this raises ``NotImplementedError`` to surface the
+        unsupported configuration; pass ``options=None`` (or ``{}``) or
+        override ``score_pairs()`` with an options-aware implementation.
+        """
+        if options:
+            msg = (
+                f"{type(self).__name__}.score_pairs(): per-call options are "
+                f"not supported by the default batching path "
+                f"(got options={options!r}). Override score_pairs() with an "
+                f"options-aware implementation."
+            )
+            raise NotImplementedError(msg)
+        return grouped_score_pairs(self.score, queries, docs, instruction=instruction)
+
     # -- Shared helpers ------------------------------------------------------
 
     def _check_loaded(self) -> None:

{sie_server-0.3.1 → sie_server-0.3.2}/src/sie_server/adapters/_utils.py

@@ -1,6 +1,10 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, Protocol
+
+import numpy as np
+
+from sie_server.core.inference_output import ScoreOutput
 
 if TYPE_CHECKING:
     import torch
@@ -8,6 +12,16 @@ if TYPE_CHECKING:
     from sie_server.types.inputs import Item
 
 
+class _ScoreFn(Protocol):
+    def __call__(
+        self,
+        query: Item,
+        items: list[Item],
+        *,
+        instruction: str | None = ...,
+    ) -> list[float]: ...
+
+
 # ---------------------------------------------------------------------------
 # RoPE utilities (eliminates 7 identical copies)
 # ---------------------------------------------------------------------------
@@ -140,3 +154,67 @@ def resolve_embedding_options(
         opts.get("query_template", default_query_template),
         opts.get("doc_template", default_doc_template),
     )
+
+
+# ---------------------------------------------------------------------------
+# Score-pair grouping (shared by ColBERT-family adapters)
+# ---------------------------------------------------------------------------
+
+
+def grouped_score_pairs(
+    score_fn: _ScoreFn,
+    queries: list[Item],
+    docs: list[Item],
+    *,
+    instruction: str | None = None,
+) -> ScoreOutput:
+    """Run a per-query ``score()`` callable over parallel (query, doc) pairs.
+
+    Groups pairs by ``(query.text, query.id, instruction)`` so each unique
+    query is encoded once and its docs are scored as one batch. Used by
+    ColBERT-family adapters to satisfy the worker's ``score_pairs()``
+    contract while reusing the optimized batched ``score()``.
+
+    Queries with ``text is None`` are not supported and raise ``ValueError``
+    (ColBERT scoring requires text). The grouping key is
+    ``(query.text, query.id or "", instruction or "")`` — two distinct
+    ``Item`` objects with identical text/id/instruction collapse to one
+    encoding pass.
+
+    Args:
+        score_fn: Bound ``adapter.score(query, items, *, instruction=None)``.
+        queries: Query items (parallel to docs).
+        docs: Document items to score.
+        instruction: Optional instruction passed through to ``score_fn``.
+
+    Returns:
+        ``ScoreOutput`` with one float per pair, in the original input order.
+
+    Raises:
+        ValueError: If ``queries`` and ``docs`` lengths differ, or any query
+            lacks text.
+    """
+    if len(queries) != len(docs):
+        msg = f"queries and docs must be parallel; got {len(queries)} vs {len(docs)}"
+        raise ValueError(msg)
+
+    if not docs:
+        return ScoreOutput(scores=np.zeros(0, dtype=np.float32), batch_size=0)
+
+    groups: dict[tuple[str, str, str], list[int]] = {}
+    for i, q in enumerate(queries):
+        if q.text is None:
+            msg = f"grouped_score_pairs requires queries[{i}].text; got None"
+            raise ValueError(msg)
+        key = (q.text, q.id or "", instruction or "")
+        groups.setdefault(key, []).append(i)
+
+    scores = np.zeros(len(docs), dtype=np.float32)
+    for indices in groups.values():
+        q = queries[indices[0]]
+        group_docs = [docs[i] for i in indices]
+        group_scores = score_fn(q, group_docs, instruction=instruction)
+        for idx, s in zip(indices, group_scores, strict=True):
+            scores[idx] = float(s)
+
+    return ScoreOutput(scores=scores, batch_size=len(docs))

{sie_server-0.3.1 → sie_server-0.3.2}/src/sie_server/adapters/base.py

@@ -127,6 +127,20 @@ class ModelAdapter(ABC):
             device: Device string (e.g., "cuda:0", "cpu").
         """
 
+    def warmup(self) -> None:
+        """Run a warmup forward pass on the loaded model.
+
+        Called by the model loader after ``load()`` has completed. The default
+        implementation is a no-op for adapters that do not need warmup. Adapters
+        that compile kernels on first call (e.g. flash-attention) or otherwise
+        benefit from a priming pass should override this and run a single
+        inference pass against a tiny synthetic input.
+
+        Splitting this from ``load()`` lets the cold-start instrumentation
+        attribute deserialize and warmup time separately.
+        """
+        return
+
     @abstractmethod
     def unload(self) -> None:
         """Unload the model and free resources.

{sie_server-0.3.1 → sie_server-0.3.2}/src/sie_server/adapters/bert_flash/__init__.py

@@ -172,6 +172,7 @@ class BertFlashAdapter(PEFTLoRAMixin, FlashBaseAdapter):
             self._max_seq_length,
         )
 
+    def warmup(self) -> None:
         # Warmup flash attention kernels
         logger.info("Warming up CUDA kernels...")
         warmup_items = [Item(text="warmup")]

{sie_server-0.3.1 → sie_server-0.3.2}/src/sie_server/adapters/docling/__init__.py

@@ -2,7 +2,6 @@ from __future__ import annotations
 
 import io
 import logging
-from concurrent.futures import ThreadPoolExecutor
 from typing import TYPE_CHECKING, Any, ClassVar
 
 from sie_server.adapters._base_adapter import BaseAdapter
@@ -52,11 +51,17 @@ class DoclingAdapter(BaseAdapter):
     OCR is disabled by default for speed and predictability. Pass
     ``options={"ocr": True}`` per request to enable it.
 
-    Concurrency: a fresh ``DocumentConverter`` is built per item rather than
-    sharing one across threads. Construction is ~10 ms once Docling's
-    layout/table models have been pre-warmed (they cache globally), and this
-    sidesteps thread-safety concerns reported upstream
-    (https://github.com/docling-project/docling/issues/115).
+    Concurrency: one ``DocumentConverter`` is cached per ``ocr_enabled`` value
+    on the adapter instance. ``self._device`` is set once in ``load()`` and is
+    stable for the adapter's lifetime, so the effective cache key is
+    ``(self._device, ocr_enabled)`` and at most two converters ever exist per
+    adapter instance. Cross-request serialization is provided by
+    ``ModelWorker._inference_executor`` (max_workers=1), so the cache itself
+    does not need a lock. Items within one batch are processed serially
+    (rather than via a per-item thread pool) to sidestep the converter's known
+    thread-safety issue (https://github.com/docling-project/docling/issues/115);
+    at GPU-bound concurrency the upstream worker is already saturating the
+    device, so intra-batch parallelism does not buy real throughput.
     """
 
     spec: ClassVar[AdapterSpec] = AdapterSpec(
@@ -69,12 +74,13 @@ class DoclingAdapter(BaseAdapter):
         self,
         model_name_or_path: str | None = None,  # unused; Docling is package-backed
         *,
-        compute_precision: str | None = None,  # unused; Docling runs on CPU
+        compute_precision: str | None = None,  # unused; device is threaded via load()
         **kwargs: Any,
     ) -> None:
         _ = (model_name_or_path, compute_precision, kwargs)
         self._loaded = False
         self._device: str | None = None
+        self._converters: dict[bool, Any] = {}
 
     def load(self, device: str) -> None:
         self._device = device
@@ -82,12 +88,20 @@ class DoclingAdapter(BaseAdapter):
         # the first real request doesn't block on a multi-hundred-MB pull.
         # Models cache globally, so subsequent per-task converters are cheap.
         try:
-            warm_converter = self._make_converter(ocr_enabled=False)
+            warm_converter = self._get_converter(ocr_enabled=False)
             self._convert_bytes(warm_converter, _TINY_PDF_BYTES, format_hint="pdf")
+            # Also build the OCR-enabled converter so the first ocr-profile
+            # request doesn't pay layout+OCR model-init latency.
+            self._get_converter(ocr_enabled=True)
         except Exception:
             logger.exception("Docling pre-warm failed; first real request may be slow")
         self._loaded = True
 
+    def unload(self) -> None:
+        self._converters.clear()
+        self._loaded = False
+        super().unload()
+
     def extract(
         self,
         items: list[Item],
@@ -113,28 +127,35 @@ class DoclingAdapter(BaseAdapter):
         )
 
     def _run_extract(self, items: list[Item], *, ocr_enabled: bool) -> list[dict[str, Any]]:
-        """Run extract per-item, parallelized across the batch.
+        """Run extract per-item, serially.
 
-        Each task gets its own DocumentConverter (see class docstring).
+        Items are processed one at a time so we can share a single cached
+        DocumentConverter (see class docstring). At GPU-bound concurrency the
+        worker-level inference executor is already saturating the device, so
+        intra-batch parallelism does not buy real throughput.
         """
-        if len(items) <= 1:
-            return [self._extract_one(item, ocr_enabled=ocr_enabled) for item in items]
-
-        with ThreadPoolExecutor(max_workers=min(len(items), 4)) as pool:
-            futures = [pool.submit(self._extract_one, item, ocr_enabled=ocr_enabled) for item in items]
-            return [f.result() for f in futures]
+        return [self._extract_one(item, ocr_enabled=ocr_enabled) for item in items]
 
     def _extract_one(self, item: Item, *, ocr_enabled: bool) -> dict[str, Any]:
         document = item.document
         if not is_document_input(document):
             return {"error": _ERR_REQUIRES_DOCUMENT}
         try:
-            converter = self._make_converter(ocr_enabled=ocr_enabled)
+            converter = self._get_converter(ocr_enabled=ocr_enabled)
             return self._convert_bytes(converter, document["data"], format_hint=document.get("format"))
         except Exception as e:  # noqa: BLE001 - per-item failure must not poison the batch
             logger.warning("Docling extract failed for item id=%s: %s", item.id, e)
             return {"error": str(e)}
 
+    def _get_converter(self, *, ocr_enabled: bool) -> Any:
+        """Return the cached DocumentConverter for this ocr_enabled value, building lazily on first use."""
+        cached = self._converters.get(ocr_enabled)
+        if cached is not None:
+            return cached
+        converter = self._make_converter(ocr_enabled=ocr_enabled)
+        self._converters[ocr_enabled] = converter
+        return converter
+
     def _convert_bytes(self, converter: Any, data: bytes, *, format_hint: str | None) -> dict[str, Any]:
         from docling.datamodel.base_models import DocumentStream  # ty: ignore[unresolved-import]
 
@@ -151,15 +172,45 @@ class DoclingAdapter(BaseAdapter):
         }
 
     def _make_converter(self, *, ocr_enabled: bool) -> Any:
-        """Build a fresh DocumentConverter. One per task — see class docstring."""
+        """Build a fresh DocumentConverter. Callers should usually go through _get_converter() for caching.
+
+        Threads self._device through Docling's AcceleratorOptions so layout, table,
+        and OCR models actually run on the configured device. Without this, Docling
+        silently defaults to CPU regardless of how SIE was launched.
+        """
         from docling.document_converter import DocumentConverter  # ty: ignore[unresolved-import]
 
-        if not ocr_enabled:
-            return DocumentConverter()
+        accelerator_options = self._build_accelerator_options()
 
         from docling.datamodel.base_models import InputFormat  # ty: ignore[unresolved-import]
         from docling.datamodel.pipeline_options import PdfPipelineOptions  # ty: ignore[unresolved-import]
         from docling.document_converter import PdfFormatOption  # ty: ignore[unresolved-import]
 
-        pdf_opts = PdfPipelineOptions(do_ocr=True)
+        # Pass do_ocr explicitly on both paths. Docling's PdfPipelineOptions defaults
+        # do_ocr=True, so an unset default would silently OCR every PDF and make the
+        # `ocr` profile a no-op vs. the default profile.
+        pdf_kwargs: dict[str, Any] = {"do_ocr": ocr_enabled}
+        if accelerator_options is not None:
+            pdf_kwargs["accelerator_options"] = accelerator_options
+        pdf_opts = PdfPipelineOptions(**pdf_kwargs)
         return DocumentConverter(format_options={InputFormat.PDF: PdfFormatOption(pipeline_options=pdf_opts)})
+
+    def _build_accelerator_options(self) -> Any:
+        """Translate self._device into a Docling AcceleratorOptions, or None."""
+        if not self._device:
+            return None
+        from docling.datamodel.accelerator_options import AcceleratorOptions  # ty: ignore[unresolved-import]
+
+        try:
+            return AcceleratorOptions(device=str(self._device))
+        except Exception as e:  # noqa: BLE001 - pydantic validation; fall back to auto
+            logger.warning(
+                "Docling: invalid device %r, falling back to 'auto' (%s)",
+                self._device,
+                e,
+            )
+            try:
+                return AcceleratorOptions(device="auto")
+            except Exception:
+                logger.exception("Docling: failed to build AcceleratorOptions even with 'auto'")
+                return None

{sie_server-0.3.1 → sie_server-0.3.2}/src/sie_server/adapters/modernbert_flash/__init__.py

@@ -142,6 +142,7 @@ class ModernBERTFlashAdapter(PEFTLoRAMixin, FlashBaseAdapter):
             self._max_seq_length,
         )
 
+    def warmup(self) -> None:
         # Warmup flash attention kernels
         logger.info("Warming up CUDA kernels...")
         warmup_items = [Item(text="warmup")]

{sie_server-0.3.1 → sie_server-0.3.2}/src/sie_server/adapters/paddleocr_vl/__init__.py

@@ -293,6 +293,7 @@ class PaddleOCRVLAdapter(BaseAdapter):
                     max_new_tokens=max_new_tokens,
                     do_sample=False,
                     num_beams=num_beams,
+                    use_cache=True,
                 )
 
             generated_ids = output_ids[0, prompt_len:]
@@ -344,6 +345,7 @@ class PaddleOCRVLAdapter(BaseAdapter):
                 max_new_tokens=max_new_tokens,
                 do_sample=False,
                 num_beams=num_beams,
+                use_cache=True,
             )
 
         generated_ids = output_ids[0, prompt_len:]

{sie_server-0.3.1 → sie_server-0.3.2}/src/sie_server/adapters/sentence_transformer/__init__.py

@@ -82,13 +82,17 @@ class SentenceTransformerDenseAdapter(BaseAdapter):
             trust_remote_code=self._trust_remote_code,
             config_kwargs=self._config_kwargs,
         )
-        _ = self._model.encode(["warmup"], convert_to_numpy=True, show_progress_bar=False)
 
         if self._max_seq_length is not None:
             self._model.max_seq_length = self._max_seq_length
 
         self._dense_dim = self._model.get_embedding_dimension()
 
+    def warmup(self) -> None:
+        if self._model is None:
+            return
+        _ = self._model.encode(["warmup"], convert_to_numpy=True, show_progress_bar=False)
+
     def encode(
         self,
         items: list[Item],