wehub-resource-sync
791 lines
29 KiB
Python
791 lines
29 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""
|
|
MLX Embedding Model wrapper.
|
|
|
|
This module provides a wrapper around mlx-embeddings for generating
|
|
text embeddings using Apple's MLX framework, with native fallback
|
|
for XLMRoBERTa, BERT, and Qwen2-decoder embedding models.
|
|
"""
|
|
|
|
import gc
|
|
import inspect
|
|
import json
|
|
import logging
|
|
import os
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Any, Dict, List, Optional, Union
|
|
|
|
import mlx.core as mx
|
|
from mlx.utils import tree_flatten
|
|
|
|
from ..utils.compile_cache import clear_thread_compile_cache
|
|
from ..utils.image import validate_image_data_uri
|
|
from .mlx_embeddings_compat import (
|
|
patch_qwen3_vl_processor_for_torch_free_image_loading,
|
|
)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_DEFAULT_EMBEDDING_MAX_LENGTH = 512
|
|
_TOKENIZER_MAX_LENGTH_SENTINEL = 10**18
|
|
_CONTEXT_LENGTH_ATTRS = (
|
|
"max_position_embeddings",
|
|
"max_seq_len",
|
|
"max_seq_length",
|
|
"seq_length",
|
|
"n_positions",
|
|
)
|
|
_FALSE_ENV_VALUES = {"0", "false", "no", "off"}
|
|
|
|
|
|
@dataclass
|
|
class EmbeddingOutput:
|
|
"""Output from embedding generation."""
|
|
|
|
embeddings: List[List[float]]
|
|
"""List of embedding vectors, one per input text."""
|
|
|
|
total_tokens: int
|
|
"""Total number of tokens in the input."""
|
|
|
|
dimensions: int = 0
|
|
"""Dimension of each embedding vector."""
|
|
|
|
|
|
class MLXEmbeddingModel:
|
|
"""
|
|
Wrapper around mlx-embeddings for generating text embeddings.
|
|
|
|
This class provides a unified interface for loading and running
|
|
embedding models using Apple's MLX framework.
|
|
|
|
Supports:
|
|
- Native XLMRoBERTa embedding (no mlx-embeddings dependency)
|
|
- Native BERT embedding (no mlx-embeddings dependency)
|
|
- Native Qwen2-decoder embedding (last-token + L2; jina-code, gte-Qwen2)
|
|
— mlx-embeddings has no qwen2 module
|
|
- mlx-embeddings fallback for other architectures
|
|
|
|
Example:
|
|
>>> model = MLXEmbeddingModel("mlx-community/all-MiniLM-L6-v2-4bit")
|
|
>>> output = model.embed(["Hello, world!", "How are you?"])
|
|
>>> print(len(output.embeddings)) # 2
|
|
"""
|
|
|
|
def __init__(self, model_name: str, trust_remote_code: bool = False):
|
|
"""
|
|
Initialize the MLX embedding model.
|
|
|
|
Args:
|
|
model_name: HuggingFace model name or local path
|
|
trust_remote_code: Allow execution of custom Python shipped inside
|
|
the model repository. Off by default for security (issue #926).
|
|
"""
|
|
self.model_name = model_name
|
|
self.trust_remote_code = trust_remote_code
|
|
|
|
self.model = None
|
|
self.processor = None
|
|
self._loaded = False
|
|
self._hidden_size: Optional[int] = None
|
|
self._using_native = False
|
|
self._is_compiled = False
|
|
self._compiled_embed = None
|
|
self._remap_input_ids_to_inputs = False
|
|
|
|
def _load_native(self) -> bool:
|
|
"""
|
|
Try to load using native omlx implementations (xlm_roberta, bert).
|
|
|
|
Returns True if native loading succeeded, False otherwise.
|
|
"""
|
|
from transformers import AutoTokenizer
|
|
|
|
model_path = Path(self.model_name)
|
|
config_path = model_path / "config.json"
|
|
if not config_path.exists():
|
|
logger.debug(f"No config.json at {model_path}, native loading skipped")
|
|
return False
|
|
|
|
try:
|
|
with open(config_path) as f:
|
|
config_dict = json.load(f)
|
|
except (json.JSONDecodeError, IOError):
|
|
logger.debug("Failed to read config.json, native loading skipped")
|
|
return False
|
|
|
|
architectures = config_dict.get("architectures", [])
|
|
arch = architectures[0] if architectures else ""
|
|
|
|
native_arch_modules = {
|
|
"XLMRobertaModel": "xlm_roberta",
|
|
"BertModel": "xlm_roberta",
|
|
"BertForMaskedLM": "xlm_roberta",
|
|
"Qwen2ForCausalLM": "qwen2_embedding",
|
|
}
|
|
module_name = native_arch_modules.get(arch)
|
|
if module_name is None:
|
|
logger.debug(
|
|
f"Architecture '{arch}' not natively supported for embedding, "
|
|
"trying mlx-embeddings"
|
|
)
|
|
return False
|
|
|
|
try:
|
|
from importlib import import_module
|
|
|
|
native_module = import_module(f"{__package__}.{module_name}")
|
|
Model = native_module.Model
|
|
ModelArgs = native_module.ModelArgs
|
|
|
|
known_fields = {f.name for f in ModelArgs.__dataclass_fields__.values()}
|
|
model_config = {
|
|
k: v for k, v in config_dict.items() if k in known_fields
|
|
}
|
|
model_config["architectures"] = architectures
|
|
|
|
config = ModelArgs(**model_config)
|
|
model_instance = Model(config)
|
|
|
|
weights = {}
|
|
weight_files = list(model_path.glob("*.safetensors"))
|
|
if not weight_files:
|
|
logger.debug(f"No safetensors files found in {model_path}")
|
|
return False
|
|
|
|
for wf in weight_files:
|
|
weights.update(mx.load(str(wf)))
|
|
|
|
weights = model_instance.sanitize(weights)
|
|
self._validate_native_weights(model_instance, weights)
|
|
model_instance.load_weights(list(weights.items()), strict=False)
|
|
mx.eval(model_instance.parameters())
|
|
# Embedding inference must be deterministic: put the model in eval
|
|
# mode so dropout (p>0 in XLM-RoBERTa/BERT) is disabled. Without this
|
|
# every /v1/embeddings call applies random dropout, producing
|
|
# non-deterministic, corrupted vectors.
|
|
model_instance.train(False)
|
|
|
|
try:
|
|
tokenizer = AutoTokenizer.from_pretrained(
|
|
str(model_path),
|
|
use_fast=False,
|
|
trust_remote_code=self.trust_remote_code,
|
|
)
|
|
except Exception:
|
|
tokenizer = AutoTokenizer.from_pretrained(
|
|
str(model_path),
|
|
trust_remote_code=self.trust_remote_code,
|
|
)
|
|
|
|
self.model = model_instance
|
|
self.processor = tokenizer
|
|
self._hidden_size = config.hidden_size
|
|
self._loaded = True
|
|
self._using_native = True
|
|
self._is_compiled = False
|
|
self._compiled_embed = None
|
|
logger.info(
|
|
f"Embedding model loaded natively: {self.model_name} "
|
|
f"(arch={arch}, hidden_size={config.hidden_size})"
|
|
)
|
|
return True
|
|
|
|
except Exception as e:
|
|
logger.debug(f"Native loading failed for {self.model_name}: {e}")
|
|
return False
|
|
|
|
def load(self) -> None:
|
|
"""Load the model and processor/tokenizer."""
|
|
if self._loaded:
|
|
return
|
|
|
|
# 1. Try native loading first (xlm_roberta, bert)
|
|
if self._load_native():
|
|
return
|
|
|
|
# 2. Fallback to mlx-embeddings
|
|
try:
|
|
patch_qwen3_vl_processor_for_torch_free_image_loading()
|
|
from mlx_embeddings import load
|
|
|
|
logger.info(f"Loading embedding model via mlx-embeddings: {self.model_name}")
|
|
|
|
self.model, self.processor = load(
|
|
self.model_name,
|
|
tokenizer_config={"trust_remote_code": self.trust_remote_code},
|
|
)
|
|
|
|
if hasattr(self.model, "config"):
|
|
config = self.model.config
|
|
self._hidden_size = getattr(config, "hidden_size", None)
|
|
if self._hidden_size is None and hasattr(config, "text_config"):
|
|
self._hidden_size = getattr(config.text_config, "hidden_size", None)
|
|
|
|
self._using_native = False
|
|
self._detect_input_key_remapping()
|
|
self._is_compiled = self._try_compile()
|
|
self._loaded = True
|
|
logger.info(
|
|
f"Embedding model loaded successfully: {self.model_name} "
|
|
f"(hidden_size={self._hidden_size}, compiled={self._is_compiled})"
|
|
)
|
|
|
|
except ImportError:
|
|
raise ImportError(
|
|
"mlx-embeddings is required for embedding generation. "
|
|
"Install with: pip install mlx-embeddings"
|
|
)
|
|
except FileNotFoundError:
|
|
raise FileNotFoundError(
|
|
f"No safetensors weight files found for '{self.model_name}'. "
|
|
f"Embedding models require weights in safetensors format. "
|
|
f"If this is a PyTorch model, use an MLX-converted version "
|
|
f"(e.g., from mlx-community on HuggingFace)."
|
|
)
|
|
except Exception as e:
|
|
logger.error(f"Failed to load embedding model: {e}")
|
|
raise
|
|
|
|
def _extract_embeddings_array(self, outputs):
|
|
"""Extract embedding tensor from model outputs as a 2D (batch, hidden) array."""
|
|
if hasattr(outputs, "text_embeds") and outputs.text_embeds is not None:
|
|
embeddings = outputs.text_embeds
|
|
elif hasattr(outputs, "pooler_output") and outputs.pooler_output is not None:
|
|
embeddings = outputs.pooler_output
|
|
elif (
|
|
hasattr(outputs, "last_hidden_state")
|
|
and outputs.last_hidden_state is not None
|
|
):
|
|
embeddings = outputs.last_hidden_state
|
|
else:
|
|
raise ValueError(
|
|
"Model output does not contain expected embedding fields "
|
|
"(text_embeds, pooler_output, or last_hidden_state)"
|
|
)
|
|
|
|
# Some models (e.g. ModernBERT loaded as MaskedLM) skip pooling and return
|
|
# per-token features with shape (batch, seq_len, hidden). Mean pool so
|
|
# /v1/embeddings always emits one vector per input.
|
|
if embeddings.ndim == 3:
|
|
embeddings = mx.mean(embeddings, axis=1)
|
|
return embeddings
|
|
|
|
def _validate_native_weights(
|
|
self, model_instance, weights: Dict[str, Any]
|
|
) -> None:
|
|
"""Reject native checkpoints with missing or shape-incompatible core weights."""
|
|
expected_weights = dict(tree_flatten(model_instance.parameters()))
|
|
expected_weight_names = set(expected_weights.keys())
|
|
provided_weight_names = set(weights.keys())
|
|
missing_weight_names = expected_weight_names - provided_weight_names
|
|
|
|
optional_missing_prefixes = ("pooler.",)
|
|
required_missing = sorted(
|
|
name
|
|
for name in missing_weight_names
|
|
if not name.startswith(optional_missing_prefixes)
|
|
)
|
|
if required_missing:
|
|
preview = ", ".join(required_missing[:10])
|
|
suffix = "..." if len(required_missing) > 10 else ""
|
|
raise ValueError(
|
|
"Native embedding checkpoint is missing required weights: "
|
|
f"{preview}{suffix}"
|
|
)
|
|
|
|
shape_mismatches = []
|
|
for name in expected_weight_names & provided_weight_names:
|
|
expected_shape = tuple(expected_weights[name].shape)
|
|
provided_shape = tuple(weights[name].shape)
|
|
if expected_shape != provided_shape:
|
|
shape_mismatches.append((name, expected_shape, provided_shape))
|
|
|
|
if shape_mismatches:
|
|
preview = ", ".join(
|
|
f"{name}: expected {expected_shape}, got {provided_shape}"
|
|
for name, expected_shape, provided_shape in shape_mismatches[:5]
|
|
)
|
|
suffix = "..." if len(shape_mismatches) > 5 else ""
|
|
raise ValueError(
|
|
"Native embedding checkpoint has incompatible weight shapes: "
|
|
f"{preview}{suffix}"
|
|
)
|
|
|
|
def _uses_custom_embedding_inputs(self, processor) -> bool:
|
|
"""Return True when processor exposes a custom embedding input API."""
|
|
for attr_name in ("prepare_embedding_inputs", "prepare_model_inputs"):
|
|
try:
|
|
inspect.getattr_static(processor, attr_name)
|
|
return True
|
|
except AttributeError:
|
|
continue
|
|
return False
|
|
|
|
def _normalize_embedding_inputs(
|
|
self,
|
|
inputs: Union[str, Dict[str, str], List[str], List[Dict[str, str]]],
|
|
) -> List[Dict[str, str]]:
|
|
"""Normalize embedding inputs into item dicts."""
|
|
if not inputs:
|
|
return []
|
|
if isinstance(inputs, str):
|
|
return [{"text": inputs}]
|
|
if isinstance(inputs, dict):
|
|
return [dict(inputs)]
|
|
first = inputs[0]
|
|
if isinstance(first, str):
|
|
return [{"text": text} for text in inputs]
|
|
return [dict(item) for item in inputs]
|
|
|
|
@staticmethod
|
|
def _positive_context_length(value: Any) -> Optional[int]:
|
|
"""Return a usable positive context length from config/tokenizer metadata."""
|
|
if isinstance(value, bool) or not isinstance(value, int):
|
|
return None
|
|
if 0 < value < _TOKENIZER_MAX_LENGTH_SENTINEL:
|
|
return value
|
|
return None
|
|
|
|
@classmethod
|
|
def _get_config_value(cls, config: Any, key: str) -> Optional[int]:
|
|
if config is None:
|
|
return None
|
|
if isinstance(config, dict):
|
|
return cls._positive_context_length(config.get(key))
|
|
return cls._positive_context_length(getattr(config, key, None))
|
|
|
|
@classmethod
|
|
def _context_length_from_config(cls, config: Any) -> Optional[int]:
|
|
"""Read context length from model config objects or dictionaries."""
|
|
for key in _CONTEXT_LENGTH_ATTRS:
|
|
value = cls._get_config_value(config, key)
|
|
if value is not None:
|
|
return value
|
|
|
|
for nested_key in ("text_config", "language_config"):
|
|
nested = None
|
|
if isinstance(config, dict):
|
|
nested = config.get(nested_key)
|
|
elif config is not None:
|
|
nested = getattr(config, nested_key, None)
|
|
for key in _CONTEXT_LENGTH_ATTRS:
|
|
value = cls._get_config_value(nested, key)
|
|
if value is not None:
|
|
return value
|
|
|
|
return None
|
|
|
|
def _resolve_max_length(self, max_length: int | None) -> int:
|
|
"""Resolve embedding token length when callers omit an explicit limit."""
|
|
if max_length is not None:
|
|
value = self._positive_context_length(max_length)
|
|
if value is None:
|
|
raise ValueError("max_length must be a positive integer")
|
|
return value
|
|
|
|
for config in (
|
|
getattr(self.model, "config", None),
|
|
getattr(self.processor, "config", None),
|
|
):
|
|
value = self._context_length_from_config(config)
|
|
if value is not None:
|
|
return value
|
|
|
|
processor = self.processor
|
|
tokenizers = [
|
|
processor,
|
|
getattr(processor, "tokenizer", None),
|
|
getattr(processor, "_tokenizer", None),
|
|
]
|
|
for tokenizer in tokenizers:
|
|
for attr_name in ("model_max_length", "max_length"):
|
|
value = self._positive_context_length(
|
|
getattr(tokenizer, attr_name, None)
|
|
)
|
|
if value is not None:
|
|
return value
|
|
|
|
return _DEFAULT_EMBEDDING_MAX_LENGTH
|
|
|
|
def _prepare_embedding_inputs(
|
|
self,
|
|
processor,
|
|
inputs: Union[List[str], List[Dict[str, str]]],
|
|
max_length: int,
|
|
padding: bool,
|
|
truncation: bool,
|
|
):
|
|
"""
|
|
Prepare inputs for embedding inference.
|
|
|
|
Some embedding processors, such as qwen3_vl in mlx-embeddings, expose
|
|
a higher-level embedding API instead of the tokenizer-style
|
|
``processor(texts, ...)`` path. Reuse that official extension point
|
|
when available to avoid positional-argument mismatches.
|
|
"""
|
|
normalized_inputs = self._normalize_embedding_inputs(inputs)
|
|
|
|
if self._uses_custom_embedding_inputs(processor):
|
|
if hasattr(processor, "prepare_embedding_inputs"):
|
|
return processor.prepare_embedding_inputs(
|
|
normalized_inputs, return_tensors="mlx"
|
|
)
|
|
return processor.prepare_model_inputs(
|
|
normalized_inputs, return_tensors="mlx"
|
|
)
|
|
|
|
if any("image" in item for item in normalized_inputs):
|
|
raise ValueError(
|
|
f"Embedding model '{self.model_name}' does not support image inputs"
|
|
)
|
|
|
|
from mlx_embeddings.utils import prepare_inputs
|
|
|
|
return prepare_inputs(
|
|
processor,
|
|
None,
|
|
[item.get("text", "") for item in normalized_inputs],
|
|
max_length,
|
|
padding,
|
|
truncation,
|
|
None,
|
|
)
|
|
|
|
def _detect_input_key_remapping(self) -> None:
|
|
"""Check if the model accepts `inputs` instead of `input_ids` and cache the result."""
|
|
try:
|
|
params = inspect.signature(self.model.__call__).parameters
|
|
self._remap_input_ids_to_inputs = (
|
|
"input_ids" not in params and "inputs" in params
|
|
)
|
|
except (TypeError, ValueError):
|
|
self._remap_input_ids_to_inputs = False
|
|
|
|
def _adapt_model_inputs_for_call(
|
|
self, model_inputs: Dict[str, Any]
|
|
) -> Dict[str, Any]:
|
|
"""Rename prepared inputs to match the embedding model call signature."""
|
|
adapted_inputs = dict(model_inputs)
|
|
if self._remap_input_ids_to_inputs and "input_ids" in adapted_inputs:
|
|
adapted_inputs["inputs"] = adapted_inputs.pop("input_ids")
|
|
return adapted_inputs
|
|
|
|
def _try_compile(self) -> bool:
|
|
"""
|
|
Compile a primitive-output embedding forward function.
|
|
|
|
Root-cause fix:
|
|
- Compiling model.__call__ directly can return arrays without primitives
|
|
for some embedding/reranker models, causing eval() runtime errors.
|
|
- We compile a narrower function that returns only the final embedding array.
|
|
"""
|
|
compile_env = os.getenv("OMLX_EMBEDDING_COMPILE", "1").strip().lower()
|
|
if compile_env in _FALSE_ENV_VALUES:
|
|
logger.info(
|
|
"mx.compile disabled for %s by OMLX_EMBEDDING_COMPILE",
|
|
self.model_name,
|
|
)
|
|
self._compiled_embed = None
|
|
return False
|
|
|
|
base_model = self.model
|
|
|
|
try:
|
|
def _compiled_embed(inputs):
|
|
outputs = base_model(**self._adapt_model_inputs_for_call(inputs))
|
|
return self._extract_embeddings_array(outputs)
|
|
|
|
self._compiled_embed = mx.compile(_compiled_embed)
|
|
|
|
test_inputs = {"input_ids": mx.zeros((1, 4), dtype=mx.int32)}
|
|
_ = self._compiled_embed(test_inputs)
|
|
|
|
logger.info(
|
|
f"mx.compile enabled for {self.model_name} "
|
|
f"(primitive embedding path)"
|
|
)
|
|
return True
|
|
except Exception as e:
|
|
logger.info(f"mx.compile unavailable for {self.model_name}: {e}")
|
|
self._compiled_embed = None
|
|
return False
|
|
|
|
def close(self) -> None:
|
|
"""Release model, processor, and compiled embedding resources."""
|
|
if not self._loaded and self.model is None and self.processor is None:
|
|
self._compiled_embed = None
|
|
self._is_compiled = False
|
|
return
|
|
|
|
logger.info(
|
|
"Releasing embedding model resources: %s "
|
|
"(compiled=%s, native=%s)",
|
|
self.model_name,
|
|
self._is_compiled,
|
|
self._using_native,
|
|
)
|
|
|
|
self._compiled_embed = None
|
|
self._is_compiled = False
|
|
|
|
self.model = None
|
|
self.processor = None
|
|
self._hidden_size = None
|
|
self._loaded = False
|
|
self._using_native = False
|
|
self._remap_input_ids_to_inputs = False
|
|
|
|
gc.collect()
|
|
mx.synchronize()
|
|
mx.clear_cache()
|
|
clear_thread_compile_cache()
|
|
gc.collect()
|
|
|
|
def embed(
|
|
self,
|
|
inputs: Union[str, 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, use
|
|
model/tokenizer metadata when available.
|
|
padding: Whether to pad shorter sequences
|
|
truncation: Whether to truncate longer sequences
|
|
|
|
Returns:
|
|
EmbeddingOutput with embeddings and token count
|
|
"""
|
|
if not self._loaded:
|
|
self.load()
|
|
|
|
max_length = self._resolve_max_length(max_length)
|
|
normalized_inputs = self._normalize_embedding_inputs(inputs)
|
|
for item in normalized_inputs:
|
|
image_ref = item.get("image")
|
|
if isinstance(image_ref, str):
|
|
item["image"] = validate_image_data_uri(
|
|
image_ref,
|
|
field="items[].image",
|
|
)
|
|
input_texts = [item["text"] for item in normalized_inputs if "text" in item]
|
|
has_image_inputs = any("image" in item for item in normalized_inputs)
|
|
|
|
processor = self.processor
|
|
uses_custom_embedding_inputs = self._uses_custom_embedding_inputs(processor)
|
|
if hasattr(processor, "_tokenizer") and not uses_custom_embedding_inputs:
|
|
processor = processor._tokenizer
|
|
|
|
if has_image_inputs and (self._using_native or not uses_custom_embedding_inputs):
|
|
raise ValueError(
|
|
f"Embedding model '{self.model_name}' does not support image inputs"
|
|
)
|
|
|
|
embeddings_array = None
|
|
total_tokens: Optional[int] = None
|
|
|
|
if self._using_native:
|
|
if hasattr(processor, "__call__"):
|
|
encoded = processor(
|
|
input_texts,
|
|
padding=padding,
|
|
truncation=truncation,
|
|
max_length=max_length,
|
|
return_tensors="np",
|
|
)
|
|
input_ids = mx.array(encoded["input_ids"])
|
|
attention_mask = mx.array(encoded["attention_mask"])
|
|
else:
|
|
encoded_ids = []
|
|
masks = []
|
|
for text in input_texts:
|
|
enc = processor.encode(text, add_special_tokens=True)
|
|
ids = list(enc.ids)
|
|
if truncation:
|
|
ids = ids[:max_length]
|
|
encoded_ids.append(ids)
|
|
max_len = max(len(ids) for ids in encoded_ids)
|
|
padded = []
|
|
for ids in encoded_ids:
|
|
pad_len = max_len - len(ids)
|
|
padded.append(ids + [0] * pad_len)
|
|
masks.append([1] * len(ids) + [0] * pad_len)
|
|
input_ids = mx.array(padded)
|
|
attention_mask = mx.array(masks)
|
|
|
|
outputs = self.model(input_ids=input_ids, attention_mask=attention_mask)
|
|
embeddings_array = self._extract_embeddings_array(outputs)
|
|
total_tokens = self._count_prepared_tokens(
|
|
{"attention_mask": attention_mask, "input_ids": input_ids}
|
|
)
|
|
else:
|
|
if self._is_compiled and self._compiled_embed is not None:
|
|
try:
|
|
inputs = self._prepare_embedding_inputs(
|
|
processor,
|
|
normalized_inputs,
|
|
max_length,
|
|
padding,
|
|
truncation,
|
|
)
|
|
if not isinstance(inputs, dict):
|
|
inputs = dict(inputs)
|
|
total_tokens = self._count_prepared_tokens(inputs)
|
|
embeddings_array = self._compiled_embed(inputs)
|
|
except Exception as e:
|
|
logger.warning(
|
|
f"compiled embedding path failed for {self.model_name}: {e}; "
|
|
"disabling compile and falling back to eager generate()"
|
|
)
|
|
self._is_compiled = False
|
|
self._compiled_embed = None
|
|
total_tokens = None
|
|
|
|
if embeddings_array is None:
|
|
if uses_custom_embedding_inputs:
|
|
inputs = self._prepare_embedding_inputs(
|
|
processor,
|
|
normalized_inputs,
|
|
max_length,
|
|
padding,
|
|
truncation,
|
|
)
|
|
if not isinstance(inputs, dict):
|
|
inputs = dict(inputs)
|
|
outputs = self.model(**self._adapt_model_inputs_for_call(inputs))
|
|
total_tokens = self._count_prepared_tokens(inputs)
|
|
else:
|
|
from mlx_embeddings import generate
|
|
|
|
outputs = generate(
|
|
self.model,
|
|
processor,
|
|
input_texts,
|
|
max_length=max_length,
|
|
padding=padding,
|
|
truncation=truncation,
|
|
)
|
|
embeddings_array = self._extract_embeddings_array(outputs)
|
|
|
|
mx.eval(embeddings_array)
|
|
embeddings = embeddings_array.tolist()
|
|
if total_tokens is None:
|
|
total_tokens = self._count_tokens(normalized_inputs)
|
|
dimensions = len(embeddings[0]) if embeddings else 0
|
|
|
|
return EmbeddingOutput(
|
|
embeddings=embeddings,
|
|
total_tokens=total_tokens,
|
|
dimensions=dimensions,
|
|
)
|
|
|
|
def _count_tokens(
|
|
self, inputs: Union[List[str], List[Dict[str, str]]]
|
|
) -> int:
|
|
"""Count total tokens in input texts."""
|
|
total = 0
|
|
processor = self.processor
|
|
|
|
for item in self._normalize_embedding_inputs(inputs):
|
|
text = item.get("text")
|
|
if not text:
|
|
continue
|
|
if hasattr(processor, "encode"):
|
|
tokens = processor.encode(text, add_special_tokens=True)
|
|
if isinstance(tokens, list):
|
|
total += len(tokens)
|
|
elif hasattr(tokens, "shape"):
|
|
total += tokens.shape[-1] if tokens.ndim > 0 else 1
|
|
elif hasattr(tokens, "ids"):
|
|
total += len(tokens.ids)
|
|
else:
|
|
total += len(tokens)
|
|
elif hasattr(processor, "tokenizer"):
|
|
tokens = processor.tokenizer.encode(text, add_special_tokens=True)
|
|
total += len(tokens) if isinstance(tokens, list) else len(list(tokens))
|
|
elif hasattr(processor, "_tokenizer"):
|
|
tokens = processor._tokenizer.encode(text, add_special_tokens=True)
|
|
total += len(tokens) if isinstance(tokens, list) else len(list(tokens))
|
|
else:
|
|
total += len(text.split()) + 2
|
|
|
|
return total
|
|
|
|
def _count_prepared_tokens(self, prepared_inputs: Dict[str, Any]) -> int:
|
|
"""Count tokens from prepared model inputs, including multimodal tokens."""
|
|
attention_mask = prepared_inputs.get("attention_mask")
|
|
if attention_mask is not None:
|
|
try:
|
|
return int(mx.sum(attention_mask).item())
|
|
except (TypeError, ValueError):
|
|
pass
|
|
if isinstance(attention_mask, list):
|
|
return int(sum(sum(row) if isinstance(row, list) else row for row in attention_mask))
|
|
if hasattr(attention_mask, "tolist"):
|
|
values = attention_mask.tolist()
|
|
if values and isinstance(values[0], list):
|
|
return int(sum(sum(row) for row in values))
|
|
return int(sum(values))
|
|
|
|
input_ids = prepared_inputs.get("input_ids")
|
|
if input_ids is None:
|
|
return 0
|
|
if hasattr(input_ids, "shape"):
|
|
if len(input_ids.shape) == 0:
|
|
return 1
|
|
if len(input_ids.shape) == 1:
|
|
return int(input_ids.shape[0])
|
|
return int(input_ids.shape[0] * input_ids.shape[1])
|
|
if isinstance(input_ids, list):
|
|
if input_ids and isinstance(input_ids[0], list):
|
|
return int(sum(len(row) for row in input_ids))
|
|
return int(len(input_ids))
|
|
return 0
|
|
|
|
@property
|
|
def hidden_size(self) -> Optional[int]:
|
|
"""Get the embedding dimension."""
|
|
return self._hidden_size
|
|
|
|
def get_model_info(self) -> dict:
|
|
"""Get information about the loaded model."""
|
|
if not self._loaded:
|
|
return {"loaded": False, "model_name": self.model_name}
|
|
|
|
info = {
|
|
"loaded": True,
|
|
"model_name": self.model_name,
|
|
"hidden_size": self._hidden_size,
|
|
"native_implementation": self._using_native,
|
|
"compiled": self._is_compiled,
|
|
}
|
|
|
|
if hasattr(self.model, "config"):
|
|
config = self.model.config
|
|
info.update(
|
|
{
|
|
"model_type": getattr(config, "model_type", None),
|
|
"vocab_size": getattr(config, "vocab_size", None),
|
|
"max_position_embeddings": getattr(
|
|
config, "max_position_embeddings", None
|
|
),
|
|
}
|
|
)
|
|
|
|
return info
|
|
|
|
def __repr__(self) -> str:
|
|
status = "loaded" if self._loaded else "not loaded"
|
|
impl = "native" if self._using_native else "mlx-embeddings"
|
|
return (
|
|
f"<MLXEmbeddingModel model={self.model_name} "
|
|
f"status={status} impl={impl}>"
|
|
)
|