jundot--omlx/omlx/engine/embedding.py

# 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"<EmbeddingEngine model={self._model_name} status={status}>"