wehub-resource-sync
108 lines
3.2 KiB
Python
108 lines
3.2 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""
|
|
Base model utilities for omlx custom model implementations.
|
|
|
|
This module provides common utilities for implementing custom models
|
|
that are not yet supported by mlx-embeddings.
|
|
"""
|
|
|
|
from dataclasses import dataclass
|
|
from typing import Optional
|
|
|
|
import mlx.core as mx
|
|
|
|
|
|
@dataclass
|
|
class BaseModelArgs:
|
|
"""Base class for model configuration arguments."""
|
|
|
|
pass
|
|
|
|
|
|
@dataclass
|
|
class BaseModelOutput:
|
|
"""Base output class for model forward pass."""
|
|
|
|
last_hidden_state: mx.array
|
|
"""Hidden states from the last layer."""
|
|
|
|
text_embeds: Optional[mx.array] = None
|
|
"""Normalized text embeddings."""
|
|
|
|
pooler_output: Optional[mx.array] = None
|
|
"""Pooled output (e.g., CLS token or mean pooling)."""
|
|
|
|
hidden_states: Optional[tuple] = None
|
|
"""All hidden states if output_hidden_states=True."""
|
|
|
|
|
|
def mean_pooling(hidden_states: mx.array, attention_mask: mx.array) -> mx.array:
|
|
"""
|
|
Perform mean pooling over sequence with attention mask.
|
|
|
|
Args:
|
|
hidden_states: Shape (batch_size, seq_len, hidden_size)
|
|
attention_mask: Shape (batch_size, seq_len)
|
|
|
|
Returns:
|
|
Pooled output of shape (batch_size, hidden_size)
|
|
"""
|
|
# Expand mask to match hidden states shape
|
|
mask_expanded = attention_mask[:, :, None].astype(hidden_states.dtype)
|
|
|
|
# Sum embeddings weighted by mask
|
|
sum_embeddings = mx.sum(hidden_states * mask_expanded, axis=1)
|
|
|
|
# Sum mask values (clip to avoid division by zero)
|
|
sum_mask = mx.clip(mx.sum(mask_expanded, axis=1), a_min=1e-9, a_max=None)
|
|
|
|
return sum_embeddings / sum_mask
|
|
|
|
|
|
def last_token_pool(
|
|
hidden_states: mx.array, attention_mask: Optional[mx.array] = None
|
|
) -> mx.array:
|
|
"""
|
|
Pool the last *non-pad* token of each sequence (mask-aware).
|
|
|
|
Decoder embedding models (Qwen2/Qwen3, gte-Qwen2, jina-code) represent a
|
|
sequence by its final token's hidden state. The pool must be
|
|
mask-aware because these tokenizers may pad on either side: a hardcoded
|
|
``[:, -1]`` is only correct under left padding and silently corrupts vectors
|
|
under right padding.
|
|
|
|
Args:
|
|
hidden_states: Shape (batch_size, seq_len, hidden_size)
|
|
attention_mask: Shape (batch_size, seq_len). If None, uses the last
|
|
position.
|
|
|
|
Returns:
|
|
Pooled output of shape (batch_size, hidden_size)
|
|
"""
|
|
if attention_mask is None:
|
|
return hidden_states[:, -1]
|
|
|
|
# Left padding: every sequence ends on a real token, so the last position
|
|
# is always valid and we can index it directly.
|
|
left_padding = attention_mask[:, -1].sum() == attention_mask.shape[0]
|
|
if left_padding:
|
|
return hidden_states[:, -1]
|
|
|
|
# Right (or mixed) padding: the last valid token is at sum(mask) - 1.
|
|
sequence_lengths = attention_mask.sum(axis=1) - 1
|
|
batch_size = hidden_states.shape[0]
|
|
return hidden_states[mx.arange(batch_size), sequence_lengths]
|
|
|
|
|
|
def normalize_embeddings(embeddings: mx.array) -> mx.array:
|
|
"""
|
|
L2 normalize embeddings.
|
|
|
|
Args:
|
|
embeddings: Shape (..., hidden_size)
|
|
|
|
Returns:
|
|
Normalized embeddings with same shape
|
|
"""
|
|
return embeddings / mx.linalg.norm(embeddings, axis=-1, keepdims=True)
|