Skip to content

vllm_mlx.embedding

Embedding engine using mlx-embeddings.

View the complete module source at #L1-L131.

API details

Each callable below includes its exact signature, type annotations, inputs, defaults, return contract, documented exceptions, implementation source, and parsed docstring sections when the source provides them.

vllm_mlx.embedding

Embedding engine using mlx-embeddings.

Provides lazy-loaded model management and batch embedding generation for the OpenAI-compatible /v1/embeddings endpoint.

vllm_mlx.embedding.logger module-attribute 

logger = logging.getLogger(__name__)

vllm_mlx.embedding.EmbeddingEngine 

EmbeddingEngine(model_name: str)

Wrapper around mlx-embeddings for text embedding generation.

Supports lazy model loading and batch embedding with proper tokenization and pooling.

Source code in vllm_mlx/embedding.py 
27
28
29
30
31
def __init__(self, model_name: str):
    self.model_name = model_name
    self._model = None
    self._tokenizer = None
    self._max_length: int | None = None

vllm_mlx.embedding.EmbeddingEngine.model_name instance-attribute 

model_name = model_name

vllm_mlx.embedding.EmbeddingEngine._model instance-attribute 

_model = None

vllm_mlx.embedding.EmbeddingEngine._tokenizer instance-attribute 

_tokenizer = None

vllm_mlx.embedding.EmbeddingEngine._max_length instance-attribute 

_max_length: int | None = None

vllm_mlx.embedding.EmbeddingEngine.is_loaded property 

is_loaded: bool

Return whether the embedding model has been loaded.

vllm_mlx.embedding.EmbeddingEngine.load 

load() -> None

Load the embedding model and tokenizer.

Source code in vllm_mlx/embedding.py 
39
40
41
42
43
44
45
46
47
def load(self) -> None:
    """Load the embedding model and tokenizer."""
    from mlx_embeddings import load

    logger.info(f"Loading embedding model: {self.model_name}")
    start = time.perf_counter()
    self._model, self._tokenizer = load(self.model_name)
    elapsed = time.perf_counter() - start
    logger.info(f"Embedding model loaded in {elapsed:.2f}s: {self.model_name}")

vllm_mlx.embedding.EmbeddingEngine._ensure_loaded 

_ensure_loaded() -> None
Source code in vllm_mlx/embedding.py 
49
50
51
def _ensure_loaded(self) -> None:
    if not self.is_loaded:
        self.load()

vllm_mlx.embedding.EmbeddingEngine._resolve_max_length 

_resolve_max_length() -> int

Tokenizer truncation length from the model config (cached).

Source code in vllm_mlx/embedding.py 
53
54
55
56
57
58
59
60
def _resolve_max_length(self) -> int:
    """Tokenizer truncation length from the model config (cached)."""
    if self._max_length is None:
        self._max_length = resolve_max_length(
            getattr(self._model, "config", None),
            self._tokenizer,
        )
    return self._max_length

vllm_mlx.embedding.EmbeddingEngine.embed 

embed(texts: str | list[str]) -> list[list[float]]

Generate embeddings for one or more texts.

Parameters:

  • texts (str | list[str]) –

    A single string or list of strings.

Returns:

  • list[list[float]] –

    List of embedding vectors (one per input text).

Source code in vllm_mlx/embedding.py 
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
def embed(self, texts: str | list[str]) -> list[list[float]]:
    """
    Generate embeddings for one or more texts.

    Args:
        texts: A single string or list of strings.

    Returns:
        List of embedding vectors (one per input text).
    """
    self._ensure_loaded()

    if isinstance(texts, str):
        texts = [texts]

    # Tokenize directly instead of using mlx_embeddings.generate(),
    # which has compatibility issues with newer tokenizers (e.g.
    # GemmaTokenizer lacks batch_encode_plus, and the model's __call__
    # expects positional `inputs` not `input_ids` as a kwarg).
    inner_tok = inner_tokenizer(self._tokenizer)
    encoded = inner_tok(
        texts,
        padding=True,
        truncation=True,
        max_length=self._resolve_max_length(),
        return_tensors="np",
    )

    input_ids = mx.array(encoded["input_ids"])
    attention_mask = mx.array(encoded["attention_mask"])

    output = self._model(input_ids, attention_mask=attention_mask)

    # text_embeds shape: (batch_size, embedding_dim)
    embeds: mx.array = output.text_embeds

    # Convert to Python lists for JSON serialization
    result = embeds.tolist()

    # Release the Metal buffers this pass allocated. MLX keeps freed buffers
    # in its allocator pool, keyed by size, and `padding=True` above makes
    # the sequence length vary from batch to batch — so nearly every request
    # asks for sizes the pool has never seen and cannot reuse. Without this
    # the pool only grows: measured ~70 MB retained per input text, taking a
    # fresh process from 2.3 GB to 24 GB over 320 texts.
    mx.clear_cache()

    return result

vllm_mlx.embedding.EmbeddingEngine.count_tokens 

count_tokens(texts: str | list[str]) -> int

Approximate token count for usage reporting.

Source code in vllm_mlx/embedding.py 
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
def count_tokens(self, texts: str | list[str]) -> int:
    """Approximate token count for usage reporting."""
    self._ensure_loaded()

    if isinstance(texts, str):
        texts = [texts]

    total = 0
    for text in texts:
        try:
            tokens = self._tokenizer.encode(text)
            if isinstance(tokens, list):
                total += len(tokens)
            elif hasattr(tokens, "__len__"):
                total += len(tokens)
            else:
                total += tokens.size
        except Exception:
            # Fallback: rough estimate of ~4 chars per token
            total += max(1, len(text) // 4)
    return total

Complete contract reference

Expand any definition for its exact inputs, annotations, defaults, return contract, directly raised exceptions, source-grounded behavior, and immutable line link. This section includes private and nested definitions that ordinary API generators omit.

vllm_mlx.embedding.EmbeddingEngine · class 
vllm_mlx.embedding.EmbeddingEngine(model_name: str)

Wrapper around mlx-embeddings for text embedding generation.

Parameters

Name Type Required Default Description
model_name str yes none Required positional or keyword input.

Returns

  • Constructs: vllm_mlx.embedding.EmbeddingEngine

Exceptions and behavior

Class EmbeddingEngine declares 7 direct member(s). No direct raise statement appears in this definition.

View source #L19-L131.

vllm_mlx.embedding.EmbeddingEngine.__init__ · method 
vllm_mlx.embedding.EmbeddingEngine.__init__(model_name: str) -> not annotated

Method EmbeddingEngine.__init__ updates self.model_name, self._model, self._tokenizer, self._max_length.

Parameters

Name Type Required Default Description
model_name str yes none Required positional or keyword input.

Returns

  • Type: not annotated

Exceptions and behavior

Method EmbeddingEngine.__init__ updates self.model_name, self._model, self._tokenizer, self._max_length. No direct raise statement appears in this definition.

View source #L27-L31.

vllm_mlx.embedding.EmbeddingEngine.is_loaded · method 
vllm_mlx.embedding.EmbeddingEngine.is_loaded() -> bool

Return whether the embedding model has been loaded.

Parameters

This callable has no explicit inputs.

Returns

  • Type: bool
  • Direct return expressions: self._model is not None

Exceptions and behavior

Method EmbeddingEngine.is_loaded returns self._model is not None. No direct raise statement appears in this definition.

View source #L34-L37.

vllm_mlx.embedding.EmbeddingEngine.load · method 
vllm_mlx.embedding.EmbeddingEngine.load() -> None

Load the embedding model and tokenizer.

Parameters

This callable has no explicit inputs.

Returns

  • Type: None

Exceptions and behavior

Method EmbeddingEngine.load updates self._model, self._tokenizer; calls logger.info, time.perf_counter, load. No direct raise statement appears in this definition.

View source #L39-L47.

vllm_mlx.embedding.EmbeddingEngine._ensure_loaded · method 
vllm_mlx.embedding.EmbeddingEngine._ensure_loaded() -> None

Method EmbeddingEngine._ensure_loaded calls self.load.

Parameters

This callable has no explicit inputs.

Returns

  • Type: None

Exceptions and behavior

Method EmbeddingEngine._ensure_loaded calls self.load. No direct raise statement appears in this definition.

View source #L49-L51.

vllm_mlx.embedding.EmbeddingEngine._resolve_max_length · method 
vllm_mlx.embedding.EmbeddingEngine._resolve_max_length() -> int

Tokenizer truncation length from the model config (cached).

Parameters

This callable has no explicit inputs.

Returns

  • Type: int
  • Direct return expressions: self._max_length

Exceptions and behavior

Method EmbeddingEngine._resolve_max_length updates self._max_length; calls resolve_max_length, getattr; returns self._max_length. No direct raise statement appears in this definition.

View source #L53-L60.

vllm_mlx.embedding.EmbeddingEngine.embed · method 
vllm_mlx.embedding.EmbeddingEngine.embed(texts: str | list[str]) -> list[list[float]]

Generate embeddings for one or more texts.

Parameters

Name Type Required Default Description
texts str \| list[str] yes none A single string or list of strings.

Returns

  • Type: list[list[float]]
  • Direct return expressions: result

Exceptions and behavior

Method EmbeddingEngine.embed calls self._ensure_loaded, isinstance, inner_tokenizer, inner_tok; returns result. No direct raise statement appears in this definition.

View source #L62-L109.

vllm_mlx.embedding.EmbeddingEngine.count_tokens · method 
vllm_mlx.embedding.EmbeddingEngine.count_tokens(texts: str | list[str]) -> int

Approximate token count for usage reporting.

Parameters

Name Type Required Default Description
texts str \| list[str] yes none Required positional or keyword input.

Returns

  • Type: int
  • Direct return expressions: total

Exceptions and behavior

Method EmbeddingEngine.count_tokens calls self._ensure_loaded, isinstance, self._tokenizer.encode, len; returns total. No direct raise statement appears in this definition.

View source #L111-L131.

Complete symbol map

This map also includes private definitions and nested helpers. The signature column exposes every explicit input even when an internal helper has no dedicated parameter prose.

Symbol Kind Signature and inputs What it does Source
EmbeddingEngine class EmbeddingEngine(model_name: str) Wrapper around mlx-embeddings for text embedding generation. #L19-L131
EmbeddingEngine.__init__ method EmbeddingEngine.__init__(model_name: str) -> not annotated Method EmbeddingEngine.__init__ updates self.model_name, self._model, self._tokenizer, self._max_length. #L27-L31
EmbeddingEngine.is_loaded method EmbeddingEngine.is_loaded() -> bool Return whether the embedding model has been loaded. #L34-L37
EmbeddingEngine.load method EmbeddingEngine.load() -> None Load the embedding model and tokenizer. #L39-L47
EmbeddingEngine._ensure_loaded method EmbeddingEngine._ensure_loaded() -> None Method EmbeddingEngine._ensure_loaded calls self.load. #L49-L51
EmbeddingEngine._resolve_max_length method EmbeddingEngine._resolve_max_length() -> int Tokenizer truncation length from the model config (cached). #L53-L60
EmbeddingEngine.embed method EmbeddingEngine.embed(texts: str \| list[str]) -> list[list[float]] Generate embeddings for one or more texts. #L62-L109
EmbeddingEngine.count_tokens method EmbeddingEngine.count_tokens(texts: str \| list[str]) -> int Approximate token count for usage reporting. #L111-L131