mnemos-embedkit 0.1.0a1__py3-none-any.whl

embedkit/__init__.py

@@ -0,0 +1,16 @@
+"""mnemos-embedkit — open embedding devkit for heterogeneous silicon.
+
+Quick start:
+
+    import embedkit
+    eng = embedkit.Engine.auto()
+    vec = eng.embed("Hello world")
+
+See docs/DESIGN.md for architecture.
+"""
+from __future__ import annotations
+
+from .engine import Engine
+
+__version__ = "0.1.0-dev"
+__all__ = ["Engine", "__version__"]

embedkit/adapters/__init__.py

@@ -0,0 +1,51 @@
+"""Adapter registry.
+
+Adapters are imported lazily so a missing optional dependency on one
+adapter (e.g. `mlx` not installed on a Linux box) does not break the
+whole kit. Engine.auto() iterates the registry, calls is_available()
+on each class, and picks the fastest among those that report True
+within the highest-populated tier.
+"""
+from __future__ import annotations
+
+from .base import AbstractAdapter
+
+# Adapters registered alphabetically. ORDER HAS NO SEMANTIC MEANING.
+# auto() picks within capability tier by measured throughput, not by
+# this list order.
+_REGISTRY_NAMES: list[tuple[str, str]] = [
+    ("cpu_llamacpp",      "CPULlamaCppAdapter"),
+    ("cpu_sbert",         "CPUSBertAdapter"),
+    ("gpu_amd_rocm",      "AMDROCmAdapter"),
+    ("gpu_apple_mlx",     "AppleMLXAdapter"),
+    ("gpu_intel_igpu",    "IntelIGPUAdapter"),
+    ("gpu_nvidia_cuda",   "NvidiaCUDAAdapter"),
+    ("gpu_nvidia_trt",    "NvidiaTRTAdapter"),
+    ("gpu_vulkan",        "VulkanAdapter"),
+    ("npu_amd_xdna",      "AMDXDNAAdapter"),
+    ("npu_cix_zhouyi",    "CixZhouyiAdapter"),
+    ("npu_intel",         "IntelNPUAdapter"),
+    ("npu_mediatek_apu",  "MediaTekAPUAdapter"),
+    ("npu_rockchip",      "RockchipRKNNAdapter"),
+]
+
+
+def all_adapters() -> list[type[AbstractAdapter]]:
+    """Return all registered adapter classes that can be imported on this host.
+
+    Adapters whose import fails (missing optional deps) are skipped silently
+    and surface only when the caller asks for that adapter explicitly.
+    """
+    out: list[type[AbstractAdapter]] = []
+    for module_name, class_name in _REGISTRY_NAMES:
+        try:
+            module = __import__(f"embedkit.adapters.{module_name}", fromlist=[class_name])
+        except Exception:
+            continue
+        cls = getattr(module, class_name, None)
+        if cls is not None and issubclass(cls, AbstractAdapter):
+            out.append(cls)
+    return out
+
+
+__all__ = ["AbstractAdapter", "all_adapters"]

embedkit/adapters/base.py

@@ -0,0 +1,70 @@
+"""Abstract adapter contract.
+
+Every silicon adapter (CPU / GPU / NPU) must implement this interface.
+The Engine's auto() picks among adapters whose is_available() returns True
+within the highest-populated capability tier (npu > gpu > cpu).
+
+No vendor preference. The kit is silicon-agnostic.
+"""
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import ClassVar
+
+
+class AbstractAdapter(ABC):
+    """Contract every embedkit adapter must satisfy.
+
+    Subclasses set the class-level metadata (name, tier, model_format) at
+    class-definition time so registry probing does not require an
+    instantiation. is_available() must be cheap and side-effect-free.
+    """
+
+    name: ClassVar[str]                    # canonical kit name, e.g. "cix-npu"
+    tier: ClassVar[str]                    # "npu" | "gpu" | "cpu"
+    model_format: ClassVar[str]            # "gguf" | "onnx" | "cix" | "rknn" | "xdna" | "mlx"
+    embed_dim: int = 0                     # populated after model load
+    max_tokens: int = 0                    # populated after model load
+
+    def __init__(self, model_path: str, **kwargs: object) -> None:
+        self.model_path = model_path
+        self._kwargs = kwargs
+
+    @abstractmethod
+    def embed(self, text: str) -> list[float]:
+        """Return a single embedding vector for `text`."""
+
+    @abstractmethod
+    def embed_batch(self, texts: list[str]) -> list[list[float]]:
+        """Return embedding vectors for each entry in `texts`."""
+
+    @abstractmethod
+    def warmup(self) -> None:
+        """Run a couple of dummy embeds to populate caches / load weights."""
+
+    @abstractmethod
+    def close(self) -> None:
+        """Release any device handles, threads, or subprocesses."""
+
+    @classmethod
+    @abstractmethod
+    def is_available(cls) -> tuple[bool, str]:
+        """Return (ok, reason).
+
+        ok=True iff this adapter's runtime + driver are present on the host
+        and a model load is expected to succeed. reason is a short string
+        for the doctor / debug output.
+
+        This MUST be cheap (no model load, no GPU init beyond a probe).
+        """
+
+    def info(self) -> dict[str, object]:
+        """Introspection blob for `Engine.info()`."""
+        return {
+            "adapter": self.name,
+            "tier": self.tier,
+            "model_format": self.model_format,
+            "model_path": self.model_path,
+            "embed_dim": self.embed_dim,
+            "max_tokens": self.max_tokens,
+        }

embedkit/adapters/cpu_llamacpp.py

@@ -0,0 +1,107 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2026 Jason Perlow
+"""llama-cpp-python adapter for GGUF embedding models."""
+from __future__ import annotations
+
+import importlib
+import logging
+import os
+from typing import Any
+
+from .base import AbstractAdapter
+
+log = logging.getLogger("embedkit.adapter.cpu_llamacpp")
+
+
+class CPULlamaCppAdapter(AbstractAdapter):
+    """Embedding adapter backed by llama-cpp-python.
+
+    This path covers CPU execution and any llama.cpp accelerator backend
+    compiled into the local llama-cpp-python wheel, controlled by
+    ``n_gpu_layers`` / ``N_GPU_LAYERS``.
+    """
+
+    name = "cpu-llamacpp"
+    tier = "cpu"
+    model_format = "gguf"
+
+    def __init__(
+        self,
+        model_path: str,
+        *,
+        n_threads: int | None = None,
+        n_ctx: int = 8192,
+        n_gpu_layers: int | None = None,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(model_path, **kwargs)
+        llama_cpp = importlib.import_module("llama_cpp")
+
+        llama_kwargs = dict(kwargs)
+        llama_kwargs.pop("embedding", None)
+        llama_kwargs.pop("verbose", None)
+
+        threads = n_threads or os.cpu_count() or 4
+        gpu_layers = (
+            n_gpu_layers
+            if n_gpu_layers is not None
+            else int(os.environ.get("N_GPU_LAYERS", "0"))
+        )
+
+        self.max_tokens = n_ctx
+        self._llm: Any | None = llama_cpp.Llama(
+            model_path=model_path,
+            n_ctx=n_ctx,
+            n_threads=threads,
+            n_gpu_layers=gpu_layers,
+            embedding=True,
+            verbose=False,
+            **llama_kwargs,
+        )
+
+        sample = self._create_embedding("warmup")
+        self.embed_dim = len(sample)
+
+    @classmethod
+    def is_available(cls) -> tuple[bool, str]:
+        try:
+            importlib.import_module("llama_cpp")
+        except ImportError as exc:
+            return False, f"llama-cpp-python missing: {exc}"
+        except Exception as exc:
+            return False, f"llama-cpp-python import failed: {exc}"
+        return True, "llama-cpp-python installed"
+
+    def embed(self, text: str) -> list[float]:
+        return self._create_embedding(text)
+
+    def embed_batch(self, texts: list[str]) -> list[list[float]]:
+        return [self.embed(text) for text in texts]
+
+    def warmup(self) -> None:
+        for text in ("warmup", "test", "sample"):
+            self._create_embedding(text)
+
+    def close(self) -> None:
+        llm = self._llm
+        self._llm = None
+        if llm is None:
+            return
+        close = getattr(llm, "close", None)
+        if callable(close):
+            try:
+                close()
+            except Exception as exc:
+                log.debug("llama.cpp close failed: %s", exc)
+
+    def _create_embedding(self, text: str) -> list[float]:
+        if self._llm is None:
+            raise RuntimeError("cpu-llamacpp adapter is closed")
+        response = self._llm.create_embedding(text)
+        data = response.get("data", [])
+        if not data:
+            raise RuntimeError("llama-cpp-python returned no embedding data")
+        embedding = data[0].get("embedding")
+        if embedding is None:
+            raise RuntimeError("llama-cpp-python response missing embedding")
+        return [float(value) for value in embedding]

embedkit/adapters/gpu_apple_mlx.py

@@ -0,0 +1,257 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2026 Jason Perlow
+"""Apple Silicon MLX adapter."""
+from __future__ import annotations
+
+import importlib
+import logging
+import platform
+from typing import Any, cast
+
+import numpy as np
+from numpy.typing import NDArray
+
+from .base import AbstractAdapter
+
+log = logging.getLogger("embedkit.adapter.apple_mlx")
+
+
+class AppleMLXAdapter(AbstractAdapter):
+    """Embedding adapter for MLX models on Apple Silicon."""
+
+    name = "apple-mlx"
+    tier = "gpu"
+    model_format = "mlx"
+
+    def __init__(
+        self,
+        model_path: str,
+        *,
+        max_tokens: int = 512,
+        normalize: bool = True,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(model_path, **kwargs)
+        self._mx: Any | None = importlib.import_module("mlx.core")
+        load = self._resolve_mlx_load()
+
+        load_kwargs = dict(kwargs)
+        self._model, self._tokenizer = load(model_path, **load_kwargs)
+        self.max_tokens = max_tokens
+        self._normalize = normalize
+
+        sample = self.embed("warmup")
+        self.embed_dim = len(sample)
+
+    @classmethod
+    def is_available(cls) -> tuple[bool, str]:
+        if platform.system() != "Darwin" or platform.machine() != "arm64":
+            return False, "Apple MLX requires Apple Silicon macOS"
+        try:
+            importlib.import_module("mlx.core")
+            importlib.import_module("mlx_lm")
+        except ImportError as exc:
+            return False, f"mlx/mlx_lm missing: {exc}"
+        except Exception as exc:
+            return False, f"mlx/mlx_lm probe failed: {exc}"
+        return True, "mlx + mlx_lm present on Apple Silicon"
+
+    def embed(self, text: str) -> list[float]:
+        return self.embed_batch([text])[0]
+
+    def embed_batch(self, texts: list[str]) -> list[list[float]]:
+        if not texts:
+            return []
+        if self._mx is None or self._model is None or self._tokenizer is None:
+            raise RuntimeError("apple-mlx adapter is closed")
+
+        input_ids, attention_mask = self._tokenize_batch(texts)
+        mx_input_ids = self._mx.array(input_ids)
+        mx_attention_mask = self._mx.array(attention_mask)
+        outputs = self._call_model(mx_input_ids, mx_attention_mask)
+        vectors = self._select_vectors(outputs, attention_mask)
+        vectors = self._postprocess_vectors(vectors)
+        return [[float(value) for value in row] for row in vectors.tolist()]
+
+    def warmup(self) -> None:
+        self.embed_batch(["warmup", "test", "sample"])
+
+    def close(self) -> None:
+        mx = self._mx
+        self._model = None
+        self._tokenizer = None
+        self._mx = None
+        if mx is None:
+            return
+        try:
+            metal = getattr(mx, "metal", None)
+            clear_cache = getattr(metal, "clear_cache", None)
+            if callable(clear_cache):
+                clear_cache()
+        except Exception as exc:
+            log.debug("MLX cache clear failed: %s", exc)
+
+    @staticmethod
+    def _resolve_mlx_load() -> Any:
+        mlx_lm = importlib.import_module("mlx_lm")
+        load = getattr(mlx_lm, "load", None)
+        if callable(load):
+            return load
+        mlx_lm_utils = importlib.import_module("mlx_lm.utils")
+        load = getattr(mlx_lm_utils, "load", None)
+        if not callable(load):
+            raise RuntimeError("mlx_lm does not expose a load() helper")
+        return load
+
+    def _tokenize_batch(self, texts: list[str]) -> tuple[NDArray[Any], NDArray[Any]]:
+        tokenizer = self._tokenizer
+        try:
+            encoded = tokenizer(
+                texts,
+                padding=True,
+                truncation=True,
+                max_length=self.max_tokens,
+                return_tensors="np",
+            )
+            input_ids = np.asarray(encoded["input_ids"], dtype=np.int32)
+            if "attention_mask" in encoded:
+                attention_mask = np.asarray(encoded["attention_mask"], dtype=np.int32)
+            else:
+                attention_mask = np.ones_like(input_ids, dtype=np.int32)
+            return input_ids, attention_mask
+        except TypeError:
+            pass
+
+        encode = getattr(tokenizer, "encode", None)
+        if not callable(encode):
+            raise RuntimeError("MLX tokenizer does not support batch calls or encode()")
+
+        pad_id = getattr(tokenizer, "pad_token_id", None)
+        if pad_id is None:
+            pad_id = getattr(tokenizer, "eos_token_id", 0)
+        rows: list[list[int]] = []
+        for text in texts:
+            token_ids = list(encode(text))
+            rows.append(token_ids[: self.max_tokens])
+
+        width = max(len(row) for row in rows)
+        input_ids = np.full((len(rows), width), int(pad_id or 0), dtype=np.int32)
+        attention_mask = np.zeros((len(rows), width), dtype=np.int32)
+        for index, row in enumerate(rows):
+            input_ids[index, : len(row)] = row
+            attention_mask[index, : len(row)] = 1
+        return input_ids, attention_mask
+
+    def _call_model(self, input_ids: Any, attention_mask: Any) -> Any:
+        try:
+            outputs = self._model(
+                input_ids,
+                attention_mask=attention_mask,
+                output_hidden_states=True,
+            )
+        except TypeError:
+            try:
+                outputs = self._model(input_ids, attention_mask=attention_mask)
+            except TypeError:
+                outputs = self._model(input_ids)
+
+        eval_fn = getattr(self._mx, "eval", None)
+        if callable(eval_fn):
+            try:
+                eval_fn(outputs)
+            except TypeError:
+                pass
+        return outputs
+
+    def _select_vectors(self, outputs: Any, attention_mask: NDArray[Any]) -> NDArray[Any]:
+        arrays = self._collect_named_arrays(outputs)
+        preferred_names = ("sentence", "embedding", "pooler")
+        for token in preferred_names:
+            for name, array in arrays:
+                if token in name.lower() and array.ndim == 2:
+                    self._reject_logits(array)
+                    return array
+        for _name, array in arrays:
+            if array.ndim == 2:
+                self._reject_logits(array)
+                return array
+        for token in ("last_hidden", "hidden", "token"):
+            for name, array in arrays:
+                if token in name.lower() and array.ndim == 3:
+                    self._reject_logits(array)
+                    return self._mean_pool(array, attention_mask)
+        for _name, array in arrays:
+            if array.ndim == 3:
+                self._reject_logits(array)
+                return self._mean_pool(array, attention_mask)
+
+        shapes = [(name, tuple(array.shape)) for name, array in arrays]
+        raise RuntimeError(f"could not select embedding output from MLX shapes {shapes}")
+
+    def _collect_named_arrays(self, value: Any) -> list[tuple[str, NDArray[Any]]]:
+        names = (
+            "sentence_embedding",
+            "embeddings",
+            "pooler_output",
+            "last_hidden_state",
+            "hidden_states",
+            "logits",
+        )
+        arrays: list[tuple[str, NDArray[Any]]] = []
+
+        if isinstance(value, dict):
+            for name, item in value.items():
+                arrays.extend(
+                    (f"{name}.{child}", array)
+                    for child, array in self._collect_named_arrays(item)
+                )
+            return arrays
+
+        for name in names:
+            if hasattr(value, name):
+                arrays.extend(
+                    (f"{name}.{child}", array)
+                    for child, array in self._collect_named_arrays(getattr(value, name))
+                )
+
+        if isinstance(value, (list, tuple)):
+            for index, item in enumerate(value):
+                arrays.extend(
+                    (f"{index}.{child}", array)
+                    for child, array in self._collect_named_arrays(item)
+                )
+            return arrays
+
+        array = self._as_numpy(value)
+        if array is not None:
+            arrays.append(("output", array))
+        return arrays
+
+    @staticmethod
+    def _as_numpy(value: Any) -> NDArray[Any] | None:
+        if not hasattr(value, "shape"):
+            return None
+        try:
+            return cast(NDArray[Any], np.asarray(value))
+        except Exception:
+            return None
+
+    def _reject_logits(self, array: NDArray[Any]) -> None:
+        vocab_size = getattr(self._tokenizer, "vocab_size", None)
+        if vocab_size is not None and array.shape[-1] == int(vocab_size):
+            raise RuntimeError(
+                "MLX model returned token logits, not embeddings; use an MLX embedding model"
+            )
+
+    @staticmethod
+    def _mean_pool(hidden: NDArray[Any], attention_mask: NDArray[Any]) -> NDArray[Any]:
+        mask = attention_mask.astype(np.float32)[:, :, None]
+        denom = np.clip(mask.sum(axis=1), 1e-9, None)
+        return cast(NDArray[Any], (hidden.astype(np.float32) * mask).sum(axis=1) / denom)
+
+    def _postprocess_vectors(self, vectors: NDArray[Any]) -> NDArray[Any]:
+        out = vectors.astype(np.float32, copy=False)
+        if self._normalize:
+            norms = np.linalg.norm(out, axis=1, keepdims=True)
+            out = out / np.clip(norms, 1e-12, None)
+        return cast(NDArray[Any], out)

embedkit/adapters/gpu_nvidia_cuda.py

@@ -0,0 +1,171 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2026 Jason Perlow
+"""NVIDIA CUDA adapter backed by ONNX Runtime."""
+from __future__ import annotations
+
+import importlib
+import logging
+from pathlib import Path
+from typing import Any, cast
+
+import numpy as np
+from numpy.typing import NDArray
+
+from .base import AbstractAdapter
+
+log = logging.getLogger("embedkit.adapter.nvidia_cuda")
+
+
+class NvidiaCUDAAdapter(AbstractAdapter):
+    """Embedding adapter for ONNX models through ONNX Runtime CUDA EP."""
+
+    name = "nvidia-cuda"
+    tier = "gpu"
+    model_format = "onnx"
+
+    def __init__(
+        self,
+        model_path: str,
+        *,
+        tokenizer_path: str | None = None,
+        max_tokens: int = 512,
+        normalize: bool = True,
+        provider_options: list[dict[str, Any]] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(model_path, **kwargs)
+        ort = importlib.import_module("onnxruntime")
+        transformers = importlib.import_module("transformers")
+
+        providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
+        available = ort.get_available_providers()
+        if "CUDAExecutionProvider" not in available:
+            raise RuntimeError(f"onnxruntime CUDA EP unavailable (providers: {available})")
+
+        session_options = kwargs.get("session_options")
+        self._session: Any | None = ort.InferenceSession(
+            model_path,
+            sess_options=session_options,
+            providers=providers,
+            provider_options=provider_options,
+        )
+        self._tokenizer: Any | None = transformers.AutoTokenizer.from_pretrained(
+            str(self._resolve_tokenizer_path(model_path, tokenizer_path))
+        )
+        self.max_tokens = max_tokens
+        self._normalize = normalize
+        self._input_meta = list(self._session.get_inputs())
+        self._output_names = [output.name for output in self._session.get_outputs()]
+
+        sample = self.embed("warmup")
+        self.embed_dim = len(sample)
+
+    @classmethod
+    def is_available(cls) -> tuple[bool, str]:
+        try:
+            ort = importlib.import_module("onnxruntime")
+            providers = ort.get_available_providers()
+        except ImportError as exc:
+            return False, f"onnxruntime not installed: {exc}"
+        except Exception as exc:
+            return False, f"onnxruntime probe failed: {exc}"
+
+        if "CUDAExecutionProvider" in providers:
+            return True, "onnxruntime CUDA EP present"
+        return False, f"onnxruntime found but no CUDA EP (providers: {providers})"
+
+    def embed(self, text: str) -> list[float]:
+        return self.embed_batch([text])[0]
+
+    def embed_batch(self, texts: list[str]) -> list[list[float]]:
+        if not texts:
+            return []
+        if self._session is None or self._tokenizer is None:
+            raise RuntimeError("nvidia-cuda adapter is closed")
+
+        encoded = self._tokenizer(
+            texts,
+            padding=True,
+            truncation=True,
+            max_length=self.max_tokens,
+            return_tensors="np",
+        )
+        feeds = self._build_feeds(encoded)
+        outputs = self._session.run(None, feeds)
+        vectors = self._select_vectors(outputs, encoded.get("attention_mask"))
+        vectors = self._postprocess_vectors(vectors)
+        return [[float(value) for value in row] for row in vectors.tolist()]
+
+    def warmup(self) -> None:
+        self.embed_batch(["warmup", "test", "sample"])
+
+    def close(self) -> None:
+        self._session = None
+        self._tokenizer = None
+
+    @staticmethod
+    def _resolve_tokenizer_path(model_path: str, tokenizer_path: str | None) -> Path:
+        if tokenizer_path is not None:
+            return Path(tokenizer_path)
+        path = Path(model_path)
+        return path if path.is_dir() else path.parent
+
+    def _build_feeds(self, encoded: dict[str, Any]) -> dict[str, NDArray[Any]]:
+        batch = int(np.asarray(encoded["input_ids"]).shape[0])
+        seq_len = int(np.asarray(encoded["input_ids"]).shape[1])
+        feeds: dict[str, NDArray[Any]] = {}
+
+        for meta in self._input_meta:
+            name = meta.name
+            dtype = np.int32 if "int32" in str(meta.type) else np.int64
+            if name in encoded:
+                feeds[name] = np.asarray(encoded[name], dtype=dtype)
+            elif name == "token_type_ids":
+                feeds[name] = np.zeros((batch, seq_len), dtype=dtype)
+            elif name == "position_ids":
+                positions = np.arange(seq_len, dtype=dtype)
+                feeds[name] = np.broadcast_to(positions, (batch, seq_len)).copy()
+            else:
+                raise RuntimeError(f"ONNX model requires unsupported input '{name}'")
+
+        return feeds
+
+    def _select_vectors(self, outputs: list[Any], attention_mask: Any) -> NDArray[Any]:
+        named = [
+            (name, np.asarray(value))
+            for name, value in zip(self._output_names, outputs, strict=False)
+        ]
+        preferred_names = ("sentence", "embedding", "pooler")
+
+        for token in preferred_names:
+            for name, array in named:
+                if token in name.lower() and array.ndim == 2:
+                    return array
+        for _name, array in named:
+            if array.ndim == 2:
+                return array
+        for token in ("last_hidden", "hidden", "token"):
+            for name, array in named:
+                if token in name.lower() and array.ndim == 3:
+                    return self._mean_pool(array, attention_mask)
+        for _name, array in named:
+            if array.ndim == 3:
+                return self._mean_pool(array, attention_mask)
+
+        shapes = [tuple(array.shape) for _name, array in named]
+        raise RuntimeError(f"could not select embedding output from ONNX shapes {shapes}")
+
+    @staticmethod
+    def _mean_pool(hidden: NDArray[Any], attention_mask: Any) -> NDArray[Any]:
+        if attention_mask is None:
+            return cast(NDArray[Any], hidden.mean(axis=1))
+        mask = np.asarray(attention_mask, dtype=np.float32)[:, :, None]
+        denom = np.clip(mask.sum(axis=1), 1e-9, None)
+        return cast(NDArray[Any], (hidden.astype(np.float32) * mask).sum(axis=1) / denom)
+
+    def _postprocess_vectors(self, vectors: NDArray[Any]) -> NDArray[Any]:
+        out = vectors.astype(np.float32, copy=False)
+        if self._normalize:
+            norms = np.linalg.norm(out, axis=1, keepdims=True)
+            out = out / np.clip(norms, 1e-12, None)
+        return cast(NDArray[Any], out)