sie-server 0.1.10__py3-none-any.whl

sie_server/__init__.py

@@ -0,0 +1,3 @@
+"""SIE Server - Search Inference Engine."""
+
+__version__ = "0.1.0"

sie_server/adapters/__init__.py

@@ -0,0 +1,9 @@
+"""SIE Server model adapters."""
+
+from sie_server.adapters.base import ModelAdapter, ModelCapabilities, ModelDims
+
+__all__ = [
+    "ModelAdapter",
+    "ModelCapabilities",
+    "ModelDims",
+]

sie_server/adapters/_flash_base.py

@@ -0,0 +1,160 @@
+from __future__ import annotations
+
+import gc
+import importlib
+import logging
+from typing import TYPE_CHECKING, Any, ClassVar
+
+from sie_server.adapters.base import ModelAdapter
+
+if TYPE_CHECKING:
+    import torch
+
+logger = logging.getLogger(__name__)
+
+
+def _import_adapter_class(adapter_path: str) -> type[ModelAdapter]:
+    """Import an adapter class from a 'module:ClassName' string.
+
+    Uses the same resolution logic as ``loader._import_builtin_adapter`` but
+    accepts the **short form** used in the fallback table (e.g.
+    ``"sentence_transformer:SentenceTransformerDenseAdapter"``).  The
+    ``sie_server.adapters.`` prefix is prepended automatically.
+
+    Args:
+        adapter_path: ``"module:ClassName"`` or
+            ``"sie_server.adapters.module:ClassName"`` string.
+
+    Returns:
+        The adapter class.
+
+    Raises:
+        ImportError: If module or class cannot be found.
+    """
+    if ":" not in adapter_path:
+        msg = f"Invalid adapter_path '{adapter_path}': expected 'module:ClassName'"
+        raise ImportError(msg)
+
+    module_path, class_name = adapter_path.rsplit(":", 1)
+
+    # Allow both short form ("sentence_transformer:Foo") and full form
+    if not module_path.startswith("sie_server."):
+        module_path = f"sie_server.adapters.{module_path}"
+
+    module = importlib.import_module(module_path)
+    if not hasattr(module, class_name):
+        msg = f"Adapter class '{class_name}' not found in module '{module_path}'"
+        raise ImportError(msg)
+
+    return getattr(module, class_name)
+
+
+class FlashBaseAdapter(ModelAdapter):
+    """Thin base class for all flash-attention adapters.
+
+    Provides:
+    - Declarative fallback: set ``fallback_adapter_path`` and optionally
+      ``fallback_kwargs_overrides`` instead of overriding ``create_for_device``.
+    - Common ``unload()`` with gc + cache clearing.
+    - ``_resolve_dtype()`` for compute precision mapping.
+    - ``get_preprocessor()`` returning ``CharCountPreprocessor``.
+
+    Subclasses with custom fallback logic (e.g. SPLADEFlashAdapter) can still
+    override ``create_for_device()`` directly.
+    """
+
+    # -- Declarative fallback ------------------------------------------------
+    # Subclasses set these to enable automatic flash -> non-flash fallback.
+    fallback_adapter_path: ClassVar[str | None] = None
+    fallback_kwargs_overrides: ClassVar[dict[str, Any]] = {}
+
+    @classmethod
+    def create_for_device(cls, device: str, **kwargs: Any) -> ModelAdapter:
+        """Factory method for device-aware adapter instantiation.
+
+        When ``fallback_adapter_path`` is ``None`` the adapter is always
+        returned as-is (no fallback).  Otherwise the standard CUDA +
+        flash-attn check runs and falls back to the declared class on
+        incompatible hardware.
+        """
+        if cls.fallback_adapter_path is None:
+            return cls(**kwargs)
+
+        # Import lazily to avoid circular deps (core.inference -> core.loader -> base)
+        from sie_server.core.inference import is_flash_attention_available
+
+        if device.startswith("cuda") and is_flash_attention_available():
+            return cls(**kwargs)
+
+        # Resolve fallback class from string path
+        fallback_class = _import_adapter_class(cls.fallback_adapter_path)
+        merged = {**kwargs, **cls.fallback_kwargs_overrides}
+
+        if not device.startswith("cuda"):
+            logger.info(
+                "%s requires CUDA. Using %s for device '%s'. "
+                "For optimal performance, use a Linux system with NVIDIA GPU.",
+                cls.__name__,
+                fallback_class.__name__,
+                device,
+            )
+        else:
+            logger.warning(
+                "Flash Attention unavailable (requires Ampere+ GPU and flash-attn package). "
+                "Using %s for device '%s'. "
+                "To install on Linux: pip install sie-server[flash-attn]",
+                fallback_class.__name__,
+                device,
+            )
+
+        return fallback_class(**merged)
+
+    # -- Common unload -------------------------------------------------------
+    def unload(self) -> None:
+        """Unload model weights and free GPU memory."""
+        import torch as _torch
+
+        device = getattr(self, "_device", None)
+
+        # Clear standard fields
+        for attr in ("_model", "_tokenizer"):
+            if getattr(self, attr, None) is not None:
+                setattr(self, attr, None)
+
+        # Clear subclass-specific fields
+        for attr in self._extra_fields_to_clear():
+            if hasattr(self, attr):
+                setattr(self, attr, None)
+
+        self._device = None
+
+        gc.collect()
+        if device and str(device).startswith("cuda"):
+            _torch.cuda.empty_cache()
+
+    def _extra_fields_to_clear(self) -> list[str]:
+        """Override to list additional instance attributes to clear on unload."""
+        return []
+
+    # -- Shared utilities ----------------------------------------------------
+    def _resolve_dtype(self) -> torch.dtype:
+        """Map ``self._compute_precision`` to a ``torch.dtype``."""
+        import torch as _torch
+
+        dtype_map: dict[str, torch.dtype] = {
+            "float16": _torch.float16,
+            "bfloat16": _torch.bfloat16,
+            "float32": _torch.float32,
+        }
+        return dtype_map.get(
+            getattr(self, "_compute_precision", "float16"),
+            _torch.float16,
+        )
+
+    def get_preprocessor(self) -> Any:
+        """Return ``CharCountPreprocessor`` for cost estimation."""
+        from sie_server.core.preprocessor import CharCountPreprocessor
+
+        return CharCountPreprocessor(
+            model_name=getattr(self, "_model_name_or_path", ""),
+        )

sie_server/adapters/_utils.py

@@ -0,0 +1,146 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    import torch
+
+    from sie_server.types.inputs import Item
+
+
+# ---------------------------------------------------------------------------
+# RoPE utilities (eliminates 7 identical copies)
+# ---------------------------------------------------------------------------
+
+
+def rotate_half(x: torch.Tensor) -> torch.Tensor:
+    """Rotate half the hidden dims of the input for RoPE."""
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    import torch as _torch
+
+    return _torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_pos_emb(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Apply Rotary Position Embedding to query and key tensors.
+
+    Args:
+        q: Query tensor ``[total_tokens, num_heads, head_dim]``.
+        k: Key tensor ``[total_tokens, num_kv_heads, head_dim]``.
+        cos: Cosine part ``[total_tokens, head_dim]``.
+        sin: Sine part ``[total_tokens, head_dim]``.
+
+    Returns:
+        Rotated query and key tensors.
+    """
+    cos = cos.unsqueeze(1).to(q.dtype)
+    sin = sin.unsqueeze(1).to(q.dtype)
+    q_embed = (q * cos) + (rotate_half(q) * sin)
+    k_embed = (k * cos) + (rotate_half(k) * sin)
+    return q_embed, k_embed
+
+
+# ---------------------------------------------------------------------------
+# Output type validation (eliminates 9+ copies)
+# ---------------------------------------------------------------------------
+
+
+def validate_output_types(
+    output_types: list[str],
+    supported: set[str],
+    adapter_name: str,
+) -> None:
+    """Raise ``ValueError`` if any requested output type is unsupported."""
+    unsupported = set(output_types) - supported
+    if unsupported:
+        msg = f"Unsupported output types: {unsupported}. {adapter_name} only supports {supported!r}."
+        raise ValueError(msg)
+
+
+# ---------------------------------------------------------------------------
+# Text extraction (eliminates 9+ copies)
+# ---------------------------------------------------------------------------
+
+
+def extract_texts(
+    items: list[Item],
+    instruction: str | None,
+    *,
+    is_query: bool,
+    query_template: str | None = None,
+    doc_template: str | None = None,
+    always_apply_template: bool = False,
+    err_msg: str = "Item must have text",
+) -> list[str]:
+    """Extract text from items, applying query/doc templates.
+
+    Args:
+        items: List of input items.
+        instruction: Optional instruction string.
+        is_query: Whether items are queries (selects template).
+        query_template: Template for queries, e.g. ``"query: {text}"``.
+        doc_template: Template for documents, e.g. ``"passage: {text}"``.
+        always_apply_template: If ``True``, always apply the selected
+            template even when there is no instruction.  This matches
+            the Nomic adapter behavior where task prefixes are mandatory.
+        err_msg: Error message when ``item.text`` is ``None``.
+
+    Returns:
+        List of formatted text strings.
+    """
+    texts: list[str] = []
+    template = query_template if is_query else doc_template
+
+    for item in items:
+        if item.text is None:
+            raise ValueError(err_msg)
+
+        text = item.text
+
+        if template and (instruction or always_apply_template):
+            text = template.format(text=text, instruction=instruction or "")
+        elif instruction:
+            text = f"{instruction} {text}"
+
+        texts.append(text)
+    return texts
+
+
+def extract_text(item: Item, *, err_msg: str = "Item must have text") -> str:
+    """Extract text from a single item (cross-encoder use)."""
+    if item.text is None:
+        raise ValueError(err_msg)
+    return item.text
+
+
+# ---------------------------------------------------------------------------
+# Runtime options resolution (eliminates 5+ copies)
+# ---------------------------------------------------------------------------
+
+
+def resolve_embedding_options(
+    options: dict[str, Any] | None,
+    *,
+    default_normalize: bool,
+    default_pooling: str,
+    default_query_template: str | None,
+    default_doc_template: str | None,
+) -> tuple[bool, str, str | None, str | None]:
+    """Resolve runtime options with adapter defaults as fallback.
+
+    Returns:
+        ``(normalize, pooling, query_template, doc_template)``
+    """
+    opts = options or {}
+    return (
+        opts.get("normalize", default_normalize),
+        opts.get("pooling", default_pooling),
+        opts.get("query_template", default_query_template),
+        opts.get("doc_template", default_doc_template),
+    )