# SPDX-License-Identifier: Apache-2.0 """ Embedding engine for oMLX. This module provides an engine for generating text embeddings using mlx-embeddings. Unlike LLM engines, embedding engines don't support streaming or chat completion. """ import asyncio import gc import logging from typing import Any, Dict, List, Optional, Union import mlx.core as mx from ..engine_core import get_mlx_executor from ..models.embedding import EmbeddingOutput, MLXEmbeddingModel from .base import BaseNonStreamingEngine logger = logging.getLogger(__name__) class EmbeddingEngine(BaseNonStreamingEngine): """ Engine for generating text embeddings. This engine wraps MLXEmbeddingModel and provides async methods for integration with the oMLX server. Unlike BaseEngine, this doesn't support streaming or chat since embeddings are computed in a single forward pass. """ def __init__( self, model_name: str, trust_remote_code: bool = False, batch_size: int | None = None, *, scheduler_config: Any | None = None, ): """ Initialize the embedding engine. Args: model_name: HuggingFace model name or local path trust_remote_code: Allow loaders to execute custom Python shipped with the model repo. Off by default for security (issue #926). batch_size: Explicit per-forward input chunk size override. scheduler_config: Shared scheduler configuration. Embedding uses embedding_batch_size as its per-forward input chunk size. """ super().__init__() self._model_name = model_name self._trust_remote_code = trust_remote_code if batch_size is None: batch_size = ( getattr(scheduler_config, "embedding_batch_size", 32) if scheduler_config is not None else 32 ) self._batch_size = max(1, int(batch_size)) self._model: Optional[MLXEmbeddingModel] = None @property def model_name(self) -> str: """Get the model name.""" return self._model_name @property def processor(self) -> Any: """Get the processor/tokenizer.""" return self._model.processor if self._model else None @property def hidden_size(self) -> Optional[int]: """Get the embedding dimension.""" return self._model.hidden_size if self._model else None async def start(self) -> None: """Start the engine (load model if not loaded). Model loading runs on the global MLX executor to avoid Metal command buffer races with concurrent BatchGenerator steps. """ if self._model is not None: return logger.info(f"Starting embedding engine: {self._model_name}") self._model = MLXEmbeddingModel( self._model_name, trust_remote_code=self._trust_remote_code ) loop = asyncio.get_running_loop() await loop.run_in_executor(get_mlx_executor(), self._model.load) logger.info(f"Embedding engine started: {self._model_name}") async def stop(self) -> None: """Stop the engine and cleanup resources.""" if self._model is None: return logger.info(f"Stopping embedding engine: {self._model_name}") model = self._model loop = asyncio.get_running_loop() await loop.run_in_executor(get_mlx_executor(), model.close) self._model = None gc.collect() logger.info(f"Embedding engine stopped: {self._model_name}") async def embed( self, texts: Union[List[str], List[Dict[str, str]]], max_length: int | None = None, padding: bool = True, truncation: bool = True, ) -> EmbeddingOutput: """ Generate embeddings for input texts. Args: texts: List of input texts max_length: Maximum token length for each text. If omitted, the model resolves its configured limit. padding: Whether to pad shorter sequences truncation: Whether to truncate longer sequences Returns: EmbeddingOutput with embeddings and token count """ if self._model is None: raise RuntimeError("Engine not started. Call start() first.") model = self._model input_items = [texts] if isinstance(texts, str) else list(texts) if not input_items: return EmbeddingOutput(embeddings=[], total_tokens=0, dimensions=0) batch_size = self._batch_size activity_id = self._begin_activity( "embedding", detail="Embedding", total_items=len(input_items), metadata={"input_count": len(input_items), "batch_size": batch_size}, ) try: loop = asyncio.get_running_loop() embeddings: List[List[float]] = [] total_tokens = 0 dimensions = 0 for start in range(0, len(input_items), batch_size): batch = input_items[start:start + batch_size] def _embed_sync(): try: return model.embed( inputs=batch, max_length=max_length, padding=padding, truncation=truncation, ) finally: mx.synchronize() mx.clear_cache() output = await loop.run_in_executor(get_mlx_executor(), _embed_sync) embeddings.extend(output.embeddings) total_tokens += output.total_tokens if output.dimensions: dimensions = output.dimensions self._update_activity( activity_id, completed_items=min(start + len(batch), len(input_items)), token_count=total_tokens, dimensions=dimensions, ) output = EmbeddingOutput( embeddings=embeddings, total_tokens=total_tokens, dimensions=dimensions, ) self._update_activity( activity_id, token_count=output.total_tokens, dimensions=output.dimensions, ) return output finally: self._end_activity(activity_id) def get_stats(self) -> Dict[str, Any]: """Get engine statistics.""" return { "model_name": self._model_name, "loaded": self._model is not None, "hidden_size": self.hidden_size, "batch_size": self._batch_size, } def get_model_info(self) -> Dict[str, Any]: """Get information about the loaded model.""" if self._model is None: return {"loaded": False, "model_name": self._model_name} return self._model.get_model_info() def __repr__(self) -> str: status = "running" if self._model is not None else "stopped" return f""