wehub-resource-sync
1427 lines
54 KiB
Python
1427 lines
54 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""
|
|
MLX Reranker Model wrapper.
|
|
|
|
This module provides a wrapper for document reranking using SequenceClassification
|
|
and CausalLM-based reranker models on Apple's MLX framework.
|
|
|
|
Supports:
|
|
- ModernBertForSequenceClassification (via mlx-embeddings)
|
|
- XLMRobertaForSequenceClassification (omlx native implementation)
|
|
- CausalLM-based rerankers (e.g., Qwen3-Reranker) via yes/no logit scoring
|
|
"""
|
|
|
|
import gc
|
|
import json
|
|
import logging
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Any, Dict, Tuple
|
|
|
|
import mlx.core as mx
|
|
|
|
from ..model_discovery import (
|
|
CAUSAL_LM_RERANKER_ARCHITECTURES,
|
|
MULTIMODAL_RERANKER_ARCHITECTURES,
|
|
SUPPORTED_RERANKER_ARCHITECTURES,
|
|
_is_causal_lm_reranker,
|
|
)
|
|
from ..utils.compile_cache import clear_thread_compile_cache
|
|
from ..utils.image import load_image
|
|
from .mlx_embeddings_compat import (
|
|
patch_qwen3_vl_processor_for_torch_free_image_loading,
|
|
)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def _coerce_item_to_text(item: Any) -> str:
|
|
"""Reduce a rerank input (str or dict with 'text') to plain text.
|
|
|
|
Used by text-only reranker paths so dict inputs stay compatible with
|
|
callers that previously passed bare strings.
|
|
"""
|
|
if isinstance(item, str):
|
|
return item
|
|
if isinstance(item, dict):
|
|
return item.get("text", "") or ""
|
|
return str(item)
|
|
|
|
|
|
@dataclass
|
|
class RerankOutput:
|
|
"""Output from rerank operation."""
|
|
|
|
scores: list[float]
|
|
"""Relevance scores for each document (0 to 1)."""
|
|
|
|
indices: list[int]
|
|
"""Document indices sorted by score (descending)."""
|
|
|
|
total_tokens: int
|
|
"""Total number of tokens processed."""
|
|
|
|
|
|
class MLXRerankerModel:
|
|
"""
|
|
Wrapper for document reranking on Apple's MLX framework.
|
|
|
|
Supports two reranking paradigms:
|
|
|
|
1. SequenceClassification models (encoder-based):
|
|
- ModernBertForSequenceClassification (via mlx-embeddings)
|
|
- XLMRobertaForSequenceClassification (omlx native implementation)
|
|
|
|
2. CausalLM-based rerankers (decoder-based):
|
|
- Qwen3-Reranker and similar models that use yes/no logit scoring
|
|
- Uses instruction prompts and extracts relevance from token logits
|
|
|
|
Example:
|
|
>>> model = MLXRerankerModel("BAAI/bge-reranker-v2-m3")
|
|
>>> model.load()
|
|
>>> output = model.rerank("What is ML?", ["ML is...", "Weather is..."])
|
|
>>> print(output.scores) # [0.95, 0.12]
|
|
"""
|
|
|
|
# CausalLM reranker prompt template (Qwen3-Reranker format)
|
|
_CAUSAL_LM_SYSTEM_PROMPT = (
|
|
"Judge whether the Document meets the requirements based on the "
|
|
"Query and the Instruct provided. Note that the answer can only be "
|
|
'"yes" or "no".'
|
|
)
|
|
_CAUSAL_LM_DEFAULT_INSTRUCTION = (
|
|
"Given a web search query, retrieve relevant passages that answer the query"
|
|
)
|
|
|
|
def __init__(self, model_name: str, trust_remote_code: bool = False):
|
|
"""
|
|
Initialize the MLX reranker 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._num_labels: int | None = None
|
|
self._is_causal_lm = False
|
|
self._is_jina_reranker = False
|
|
self._is_vl_reranker = False
|
|
self._token_true_id: int | None = None
|
|
self._token_false_id: int | None = None
|
|
self._doc_embed_token_id: int | None = None
|
|
self._query_embed_token_id: int | None = None
|
|
self._jina_projector = None
|
|
self._prefix_tokens: list[int] | None = None
|
|
self._suffix_tokens: list[int] | None = None
|
|
self._is_compiled = False
|
|
self._compiled_seq_logits = None
|
|
|
|
def _get_architecture(self) -> str | None:
|
|
"""Get the model architecture from config.json."""
|
|
config_path = Path(self.model_name) / "config.json"
|
|
if not config_path.exists():
|
|
return None
|
|
|
|
try:
|
|
with open(config_path) as f:
|
|
config = json.load(f)
|
|
architectures = config.get("architectures", [])
|
|
return architectures[0] if architectures else None
|
|
except (json.JSONDecodeError, IOError):
|
|
return None
|
|
|
|
def _load_xlm_roberta(self) -> Tuple[Any, Any]:
|
|
"""Load XLMRoberta model using omlx native implementation."""
|
|
import mlx.core as mx
|
|
from mlx.utils import tree_unflatten
|
|
from transformers import AutoTokenizer
|
|
|
|
from .xlm_roberta import Model, ModelArgs
|
|
|
|
model_path = Path(self.model_name)
|
|
|
|
# Load config
|
|
with open(model_path / "config.json") as f:
|
|
config_dict = json.load(f)
|
|
|
|
config = ModelArgs(
|
|
**{
|
|
k: v
|
|
for k, v in config_dict.items()
|
|
if k in ModelArgs.__dataclass_fields__
|
|
}
|
|
)
|
|
|
|
# Create model
|
|
model = Model(config)
|
|
|
|
# Load weights. Use mx.load (not safetensors.safe_open + get_tensor),
|
|
# which reads safetensors directly into MLX arrays and supports the
|
|
# bfloat16 dtype. safe_open(framework="mlx").get_tensor() routes bf16
|
|
# through numpy, which has no bfloat16 dtype and raises
|
|
# "TypeError: data type 'bfloat16' not understood".
|
|
weights = {}
|
|
weight_files = list(model_path.glob("*.safetensors"))
|
|
for wf in weight_files:
|
|
weights.update(mx.load(str(wf)))
|
|
|
|
# Sanitize weights (remove "roberta." prefix, etc.)
|
|
weights = model.sanitize(weights)
|
|
|
|
# Load weights into model
|
|
model.load_weights(list(weights.items()))
|
|
mx.eval(model.parameters())
|
|
# Reranker inference must be deterministic: disable dropout in the
|
|
# native XLM-RoBERTa path just like the native embedding loader does.
|
|
model.train(False)
|
|
|
|
# Load tokenizer
|
|
tokenizer = AutoTokenizer.from_pretrained(
|
|
str(model_path), trust_remote_code=self.trust_remote_code
|
|
)
|
|
|
|
return model, tokenizer
|
|
|
|
def _load_vl_reranker(self) -> Tuple[Any, Any]:
|
|
"""Load a multimodal reranker (e.g., Qwen3-VL-Reranker) via mlx-embeddings.
|
|
|
|
mlx-embeddings exposes a unified `load()` + `model.process()` API that
|
|
handles both embedding and reranking variants of Qwen3-VL. Reranker vs
|
|
embedder is decided by the input dict shape at inference time.
|
|
"""
|
|
patch_qwen3_vl_processor_for_torch_free_image_loading()
|
|
from mlx_embeddings import load as mlx_emb_load
|
|
|
|
return mlx_emb_load(
|
|
str(self.model_name),
|
|
tokenizer_config={"trust_remote_code": self.trust_remote_code},
|
|
)
|
|
|
|
def _build_vl_item(self, item: "str | dict[str, Any]") -> Dict[str, Any]:
|
|
"""Normalize a rerank input into the mlx-embeddings VL item format.
|
|
|
|
Accepts either a bare string (text) or a dict with 'text' and/or
|
|
'image' keys. Request-facing image strings must be base64 data URIs and
|
|
get loaded via omlx's shared image loader.
|
|
"""
|
|
if isinstance(item, str):
|
|
return {"text": item}
|
|
if not isinstance(item, dict):
|
|
return {"text": str(item)}
|
|
|
|
result: Dict[str, Any] = {}
|
|
text = item.get("text")
|
|
if text:
|
|
result["text"] = text
|
|
image_ref = item.get("image")
|
|
if image_ref:
|
|
if isinstance(image_ref, str):
|
|
result["image"] = load_image(image_ref, field="image")
|
|
else:
|
|
# Already a PIL image or similar — pass through
|
|
result["image"] = image_ref
|
|
if not result:
|
|
raise ValueError("VL reranker item must have at least 'text' or 'image'.")
|
|
return result
|
|
|
|
def _rerank_vl(
|
|
self,
|
|
query: "str | dict[str, Any]",
|
|
documents: "list[str] | list[dict[str, Any]]",
|
|
max_length: int,
|
|
) -> RerankOutput:
|
|
"""Rerank using mlx-embeddings' multimodal model.process() API."""
|
|
query_item = self._build_vl_item(query)
|
|
doc_items = [self._build_vl_item(d) for d in documents]
|
|
|
|
inputs = {
|
|
"instruction": self._CAUSAL_LM_DEFAULT_INSTRUCTION,
|
|
"query": query_item,
|
|
"documents": doc_items,
|
|
}
|
|
|
|
scores = self.model.process(inputs, processor=self.processor)
|
|
mx.eval(scores)
|
|
scores_list = [float(s) for s in scores.tolist()]
|
|
indices = sorted(
|
|
range(len(scores_list)),
|
|
key=lambda i: scores_list[i],
|
|
reverse=True,
|
|
)
|
|
|
|
return RerankOutput(
|
|
scores=scores_list,
|
|
indices=indices,
|
|
total_tokens=0,
|
|
)
|
|
|
|
def _load_causal_lm(self) -> Tuple[Any, Any]:
|
|
"""Load a CausalLM-based reranker model using mlx-lm."""
|
|
from ..utils.model_loading import (
|
|
lm_load_compat as mlx_lm_load,
|
|
maybe_load_custom_quantization,
|
|
)
|
|
|
|
model_path = str(self.model_name)
|
|
tokenizer_config = {"trust_remote_code": self.trust_remote_code}
|
|
custom_loaded = maybe_load_custom_quantization(
|
|
model_path,
|
|
is_vlm=False,
|
|
)
|
|
if custom_loaded is not None:
|
|
model, tokenizer_wrapper = custom_loaded
|
|
else:
|
|
loaded = mlx_lm_load(
|
|
model_path,
|
|
tokenizer_config=tokenizer_config,
|
|
trust_remote_code=self.trust_remote_code,
|
|
)
|
|
model = loaded[0]
|
|
tokenizer_wrapper = loaded[1]
|
|
|
|
# mlx-lm returns a TokenizerWrapper; unwrap to get the underlying
|
|
# transformers tokenizer which supports __call__ for batch encoding.
|
|
tokenizer = getattr(tokenizer_wrapper, "_tokenizer", tokenizer_wrapper)
|
|
|
|
# Resolve yes/no token IDs from tokenizer
|
|
self._token_true_id = tokenizer.convert_tokens_to_ids("yes")
|
|
self._token_false_id = tokenizer.convert_tokens_to_ids("no")
|
|
|
|
if self._token_true_id is None or self._token_false_id is None:
|
|
raise ValueError(
|
|
"Could not find 'yes'/'no' token IDs in tokenizer. "
|
|
"This model may not be a compatible CausalLM reranker."
|
|
)
|
|
|
|
# Pre-compute prefix and suffix tokens for the prompt template.
|
|
prefix, suffix = self._extract_causal_lm_affixes(tokenizer)
|
|
|
|
self._prefix_tokens = tokenizer.encode(prefix, add_special_tokens=False)
|
|
self._suffix_tokens = tokenizer.encode(suffix, add_special_tokens=False)
|
|
|
|
logger.info(
|
|
f"CausalLM reranker tokens: yes={self._token_true_id}, "
|
|
f"no={self._token_false_id}, "
|
|
f"prefix_len={len(self._prefix_tokens)}, "
|
|
f"suffix_len={len(self._suffix_tokens)}"
|
|
)
|
|
|
|
return model, tokenizer
|
|
|
|
def _extract_causal_lm_affixes(self, tokenizer: Any) -> Tuple[str, str]:
|
|
"""Extract the static prompt prefix/suffix around the rerank content.
|
|
|
|
Handles two chat template shapes:
|
|
|
|
1. Standard chat template (system/user roles): render with a sentinel
|
|
as the user content and split around it.
|
|
2. Reranker-native template (Qwen/Qwen3-Reranker ships one as
|
|
chat_template.jinja since its 2026-04 sentence-transformers
|
|
update, and MLX conversions made after that inherit it): the
|
|
template only understands system/query/document roles and silently
|
|
drops user messages, so the sentinel never appears in the output.
|
|
Render with per-slot sentinels instead and split around the
|
|
combined "<Instruct>/<Query>/<Document>" block — the exact content
|
|
that _rerank_causal_lm reconstructs at scoring time.
|
|
"""
|
|
# Fail fast with a clear error when the tokenizer has no chat template
|
|
# at all — rendering would only produce an opaque downstream failure.
|
|
if hasattr(tokenizer, "chat_template") and tokenizer.chat_template is None:
|
|
raise ValueError(
|
|
f"Tokenizer for {self.model_name} has no chat template; "
|
|
f"cannot derive the CausalLM reranker prompt prefix/suffix."
|
|
)
|
|
|
|
_SENTINEL = "<<__CONTENT_SENTINEL__>>"
|
|
messages = [
|
|
{"role": "system", "content": self._CAUSAL_LM_SYSTEM_PROMPT},
|
|
{"role": "user", "content": _SENTINEL},
|
|
]
|
|
standard_rendered = ""
|
|
standard_error: Exception | None = None
|
|
try:
|
|
standard_rendered = tokenizer.apply_chat_template(
|
|
messages, tokenize=False, add_generation_prompt=True
|
|
)
|
|
except Exception as e:
|
|
standard_error = e
|
|
logger.warning(
|
|
f"system/user chat template rendering failed for "
|
|
f"{self.model_name}: {e}"
|
|
)
|
|
|
|
parts = standard_rendered.split(_SENTINEL)
|
|
if len(parts) == 2:
|
|
suffix = parts[1]
|
|
# Append <think> block for models that use the
|
|
# thinking-then-answering format, unless the template already
|
|
# emitted a think prefill of its own.
|
|
if "<think>" not in suffix:
|
|
suffix += "<think>\n\n</think>\n\n"
|
|
return parts[0], suffix
|
|
|
|
native_affixes, native_rendered, native_error = (
|
|
self._extract_reranker_native_affixes(tokenizer)
|
|
)
|
|
if native_affixes is not None:
|
|
logger.info(
|
|
"Using reranker-native chat template (query/document roles) "
|
|
f"for {self.model_name}"
|
|
)
|
|
return native_affixes
|
|
|
|
raise ValueError(
|
|
f"Could not extract CausalLM reranker prompt affixes for "
|
|
f"{self.model_name}. "
|
|
f"Standard system/user attempt: {standard_rendered!r} "
|
|
f"(error: {standard_error!r}). "
|
|
f"Reranker-native query/document attempt: {native_rendered!r} "
|
|
f"(error: {native_error!r})."
|
|
) from (standard_error or native_error)
|
|
|
|
def _extract_reranker_native_affixes(
|
|
self, tokenizer: Any
|
|
) -> "Tuple[Tuple[str, str] | None, str, Exception | None]":
|
|
"""Extract affixes from a reranker-native chat template, if present.
|
|
|
|
Returns (affixes, rendered, error). Affixes is None when the template
|
|
raises or does not render the expected "<Instruct>/<Query>/<Document>"
|
|
content block; the rendered string and the exception (if any) are
|
|
returned for diagnostics.
|
|
"""
|
|
instruct_sentinel = "<<__INSTRUCT_SENTINEL__>>"
|
|
query_sentinel = "<<__QUERY_SENTINEL__>>"
|
|
document_sentinel = "<<__DOCUMENT_SENTINEL__>>"
|
|
# The native template maps the system role to the <Instruct> slot; its
|
|
# judge system prompt is hardcoded inside the template itself.
|
|
messages = [
|
|
{"role": "system", "content": instruct_sentinel},
|
|
{"role": "query", "content": query_sentinel},
|
|
{"role": "document", "content": document_sentinel},
|
|
]
|
|
try:
|
|
rendered = tokenizer.apply_chat_template(
|
|
messages, tokenize=False, add_generation_prompt=True
|
|
)
|
|
except Exception as e:
|
|
logger.warning(
|
|
f"reranker-native chat template rendering failed for "
|
|
f"{self.model_name}: {e}"
|
|
)
|
|
return None, "", e
|
|
|
|
# Intentionally strict, byte-exact match against the content block the
|
|
# upstream Qwen/Qwen3-Reranker template renders. _rerank_causal_lm
|
|
# reconstructs this exact block at scoring time, so tolerating
|
|
# formatting drift here would silently produce prompts that differ
|
|
# from what the template intends; failing detection loudly is safer.
|
|
content_block = (
|
|
f"<Instruct>: {instruct_sentinel}\n"
|
|
f"<Query>: {query_sentinel}\n"
|
|
f"<Document>: {document_sentinel}"
|
|
)
|
|
parts = rendered.split(content_block)
|
|
if len(parts) != 2:
|
|
return None, rendered, None
|
|
|
|
# The native template already emits the trailing <think> block, so the
|
|
# suffix is used as-is.
|
|
return (parts[0], parts[1]), rendered, None
|
|
|
|
def _load_jina_reranker(self) -> Tuple[Any, Any]:
|
|
"""
|
|
Load a Jina v3 reranker model using mlx-lm.
|
|
|
|
Jina v3 reranker uses special-token hidden states + projector + cosine
|
|
similarity for listwise scoring.
|
|
"""
|
|
from ..utils.model_loading import (
|
|
lm_load_compat as mlx_lm_load,
|
|
maybe_load_custom_quantization,
|
|
)
|
|
|
|
model_path = str(self.model_name)
|
|
tokenizer_config = {"trust_remote_code": self.trust_remote_code}
|
|
custom_loaded = maybe_load_custom_quantization(
|
|
model_path,
|
|
is_vlm=False,
|
|
)
|
|
if custom_loaded is not None:
|
|
model, tokenizer_wrapper = custom_loaded
|
|
else:
|
|
loaded = mlx_lm_load(
|
|
model_path,
|
|
tokenizer_config=tokenizer_config,
|
|
trust_remote_code=self.trust_remote_code,
|
|
)
|
|
model = loaded[0]
|
|
tokenizer_wrapper = loaded[1]
|
|
|
|
# mlx-lm returns a TokenizerWrapper; unwrap to get the underlying
|
|
# transformers tokenizer which supports __call__ for batch encoding.
|
|
tokenizer = getattr(tokenizer_wrapper, "_tokenizer", tokenizer_wrapper)
|
|
|
|
doc_embed_token_id = self._resolve_token_id(tokenizer, "<|embed_token|>")
|
|
query_embed_token_id = self._resolve_token_id(tokenizer, "<|rerank_token|>")
|
|
|
|
if doc_embed_token_id is None or query_embed_token_id is None:
|
|
raise ValueError(
|
|
"Could not resolve required Jina special tokens "
|
|
"('<|embed_token|>', '<|rerank_token|>'). "
|
|
"This model may not be a compatible Jina v3 reranker."
|
|
)
|
|
|
|
self._doc_embed_token_id = doc_embed_token_id
|
|
self._query_embed_token_id = query_embed_token_id
|
|
self._jina_projector = self._load_jina_projector(self.model_name)
|
|
|
|
logger.info(
|
|
f"Jina reranker tokens: embed_token={doc_embed_token_id}, "
|
|
f"rerank_token={query_embed_token_id}"
|
|
)
|
|
|
|
return model, tokenizer
|
|
|
|
def _resolve_token_id(self, tokenizer: Any, token_text: str) -> int | None:
|
|
"""Resolve token IDs across tokenizer implementations."""
|
|
added_tokens = getattr(tokenizer, "added_tokens_decoder", {}) or {}
|
|
for tid, tinfo in added_tokens.items():
|
|
content = ""
|
|
if isinstance(tinfo, str):
|
|
content = tinfo
|
|
elif hasattr(tinfo, "content"):
|
|
content = tinfo.content
|
|
elif isinstance(tinfo, dict):
|
|
content = tinfo.get("content", "")
|
|
|
|
if content == token_text:
|
|
return int(tid)
|
|
|
|
convert_tokens_to_ids = getattr(tokenizer, "convert_tokens_to_ids", None)
|
|
if callable(convert_tokens_to_ids):
|
|
try:
|
|
token_id = convert_tokens_to_ids(token_text)
|
|
except Exception:
|
|
token_id = None
|
|
|
|
if isinstance(token_id, int) and token_id >= 0:
|
|
unk_token_id = getattr(tokenizer, "unk_token_id", None)
|
|
if unk_token_id is None or token_id != unk_token_id:
|
|
return token_id
|
|
|
|
get_added_vocab = getattr(tokenizer, "get_added_vocab", None)
|
|
if callable(get_added_vocab):
|
|
try:
|
|
added_vocab = get_added_vocab() or {}
|
|
except Exception:
|
|
added_vocab = {}
|
|
|
|
token_id = added_vocab.get(token_text)
|
|
if isinstance(token_id, int):
|
|
return token_id
|
|
|
|
get_vocab = getattr(tokenizer, "get_vocab", None)
|
|
if callable(get_vocab):
|
|
try:
|
|
vocab = get_vocab() or {}
|
|
except Exception:
|
|
vocab = {}
|
|
|
|
token_id = vocab.get(token_text)
|
|
if isinstance(token_id, int):
|
|
return token_id
|
|
|
|
encode = getattr(tokenizer, "encode", None)
|
|
if callable(encode):
|
|
try:
|
|
encoded = encode(token_text, add_special_tokens=False)
|
|
except TypeError:
|
|
encoded = encode(token_text)
|
|
except Exception:
|
|
encoded = None
|
|
|
|
if hasattr(encoded, "ids"):
|
|
encoded = encoded.ids
|
|
|
|
if (
|
|
isinstance(encoded, list)
|
|
and len(encoded) == 1
|
|
and isinstance(encoded[0], int)
|
|
):
|
|
return encoded[0]
|
|
|
|
return None
|
|
|
|
def _load_jina_projector(self, model_dir: str | Path):
|
|
"""Load Jina projector weights and return a projection callable."""
|
|
model_path = Path(model_dir)
|
|
projector_path = model_path / "projector.safetensors"
|
|
if not projector_path.exists():
|
|
raise FileNotFoundError(
|
|
f"Missing Jina projector file: {projector_path}. "
|
|
"Expected projector.safetensors for JinaForRanking models."
|
|
)
|
|
|
|
# mx.load reads safetensors into MLX arrays with bfloat16 support;
|
|
# safe_open(framework="mlx").get_tensor() routes bf16 through numpy and
|
|
# raises "TypeError: data type 'bfloat16' not understood".
|
|
weights = mx.load(str(projector_path))
|
|
|
|
required_keys = ("linear1.weight", "linear2.weight")
|
|
missing_keys = [key for key in required_keys if key not in weights]
|
|
if missing_keys:
|
|
raise ValueError(
|
|
f"Jina projector is malformed: missing keys {missing_keys} in "
|
|
f"{projector_path}. "
|
|
f"Available keys: {sorted(weights.keys())}"
|
|
)
|
|
|
|
linear1_weight = weights["linear1.weight"]
|
|
linear2_weight = weights["linear2.weight"]
|
|
|
|
if len(linear1_weight.shape) != 2 or len(linear2_weight.shape) != 2:
|
|
raise ValueError(
|
|
"Jina projector weights must be 2D matrices: "
|
|
f"linear1.weight={linear1_weight.shape}, "
|
|
f"linear2.weight={linear2_weight.shape}."
|
|
)
|
|
|
|
if linear1_weight.shape != (512, 1024) or linear2_weight.shape != (512, 512):
|
|
raise ValueError(
|
|
"Unexpected Jina projector shapes. Expected "
|
|
"linear1.weight=(512, 1024) and linear2.weight=(512, 512), "
|
|
f"got linear1.weight={linear1_weight.shape}, "
|
|
f"linear2.weight={linear2_weight.shape}."
|
|
)
|
|
|
|
def _project(x):
|
|
if x.shape[-1] != linear1_weight.shape[1]:
|
|
raise ValueError(
|
|
"Jina projector input dim mismatch for linear1: "
|
|
f"input={x.shape[-1]}, expected={linear1_weight.shape[1]}."
|
|
)
|
|
hidden = x @ mx.transpose(linear1_weight)
|
|
hidden = mx.maximum(hidden, 0)
|
|
return hidden @ mx.transpose(linear2_weight)
|
|
|
|
return _project
|
|
|
|
def _sanitize_jina_text(self, text: str) -> str:
|
|
"""Strip conflicting special tokens from user-provided text."""
|
|
sanitized = str(text)
|
|
sanitized = sanitized.replace("<|embed_token|>", " ")
|
|
sanitized = sanitized.replace("<|rerank_token|>", " ")
|
|
sanitized = sanitized.replace("<|score_token|>", " ")
|
|
sanitized = sanitized.replace("<|im_start|>", " ")
|
|
sanitized = sanitized.replace("<|im_end|>", " ")
|
|
return sanitized.strip()
|
|
|
|
def _format_jina_prompt(
|
|
self,
|
|
query: str,
|
|
documents: list[str],
|
|
instruction: str | None = None,
|
|
) -> str:
|
|
"""Format a listwise Jina reranking prompt."""
|
|
sanitized_query = self._sanitize_jina_text(query)
|
|
sanitized_docs = [self._sanitize_jina_text(doc) for doc in documents]
|
|
sanitized_instruction = (
|
|
self._sanitize_jina_text(instruction) if instruction is not None else None
|
|
)
|
|
|
|
user_content = (
|
|
f"I will provide you with {len(sanitized_docs)} passages, each indicated "
|
|
f"by a numerical identifier. Rank the passages based on their relevance "
|
|
f"to query: {sanitized_query}\n"
|
|
)
|
|
if sanitized_instruction:
|
|
user_content += f"<instruct>\n{sanitized_instruction}\n</instruct>\n"
|
|
|
|
doc_prompts = [
|
|
f'<passage id="{idx}">\n{doc}<|embed_token|>\n</passage>'
|
|
for idx, doc in enumerate(sanitized_docs)
|
|
]
|
|
user_content += "\n".join(doc_prompts) + "\n"
|
|
user_content += f"<query>\n{sanitized_query}<|rerank_token|>\n</query>"
|
|
|
|
system_prompt = (
|
|
"You are a search relevance expert who can determine a ranking of the "
|
|
"passages based on how relevant they are to the query. If the query is "
|
|
"a question, how relevant a passage is depends on how well it answers "
|
|
"the question. If not, try to analyze the intent of the query and "
|
|
"assess how well each passage satisfies the intent. If an instruction "
|
|
"is provided, you should follow the instruction when determining the "
|
|
"ranking."
|
|
)
|
|
|
|
return (
|
|
"<|im_start|>system\n"
|
|
f"{system_prompt}"
|
|
"<|im_end|>\n"
|
|
"<|im_start|>user\n"
|
|
f"{user_content}"
|
|
"<|im_end|>\n"
|
|
"<|im_start|>assistant\n"
|
|
"<think>\n\n</think>\n\n"
|
|
)
|
|
|
|
def _get_jina_hidden_states(self, input_ids):
|
|
"""Extract final hidden states from the Jina mlx-lm backbone."""
|
|
|
|
backbone = getattr(self.model, "model", None)
|
|
if backbone is None or not callable(backbone):
|
|
model_type = type(self.model).__name__ if self.model is not None else "None"
|
|
raise ValueError(
|
|
"Could not find Jina model backbone (model.model). "
|
|
f"The mlx-lm model wrapper may have changed: {model_type}."
|
|
)
|
|
|
|
hidden_states = backbone(input_ids)
|
|
|
|
if not hasattr(hidden_states, "shape"):
|
|
raise ValueError("Jina backbone did not return hidden states as a tensor.")
|
|
|
|
if len(hidden_states.shape) == 2:
|
|
return mx.expand_dims(hidden_states, axis=0)
|
|
|
|
if len(hidden_states.shape) != 3:
|
|
raise ValueError(
|
|
"Jina hidden states must be rank 2 or 3. "
|
|
f"Got shape: {hidden_states.shape}"
|
|
)
|
|
|
|
return hidden_states
|
|
|
|
def _cosine_similarity(self, query_vec, doc_vecs, eps: float = 1e-8):
|
|
"""Compute cosine similarity between one query vector and many docs."""
|
|
if len(query_vec.shape) == 2:
|
|
query_vec = query_vec[0]
|
|
if len(doc_vecs.shape) == 1:
|
|
doc_vecs = mx.expand_dims(doc_vecs, axis=0)
|
|
|
|
query_norm = mx.linalg.norm(query_vec)
|
|
doc_norms = mx.linalg.norm(doc_vecs, axis=-1)
|
|
denom = mx.maximum(doc_norms * query_norm, eps)
|
|
numer = mx.sum(doc_vecs * query_vec, axis=-1)
|
|
return numer / denom
|
|
|
|
def load(self) -> None:
|
|
"""Load the model and processor/tokenizer."""
|
|
if self._loaded:
|
|
return
|
|
|
|
# Check architecture before loading
|
|
self._validate_architecture()
|
|
|
|
arch = self._get_architecture()
|
|
logger.info(f"Loading reranker model: {self.model_name} (arch={arch})")
|
|
|
|
try:
|
|
if arch in MULTIMODAL_RERANKER_ARCHITECTURES:
|
|
# Multimodal reranker (e.g., Qwen3-VL-Reranker) via mlx-embeddings
|
|
self.model, self.processor = self._load_vl_reranker()
|
|
self._is_vl_reranker = True
|
|
self._num_labels = 1
|
|
elif arch == "JinaForRanking":
|
|
# Jina v3 reranker: listwise hidden-state scoring + projector
|
|
self.model, self.processor = self._load_jina_reranker()
|
|
self._is_jina_reranker = True
|
|
self._num_labels = 1
|
|
elif arch in CAUSAL_LM_RERANKER_ARCHITECTURES:
|
|
# CausalLM-based reranker (e.g., Qwen3-Reranker)
|
|
self.model, self.processor = self._load_causal_lm()
|
|
self._is_causal_lm = True
|
|
self._num_labels = 2 # yes/no
|
|
elif arch == "XLMRobertaForSequenceClassification":
|
|
# Use omlx native implementation
|
|
self.model, self.processor = self._load_xlm_roberta()
|
|
self._num_labels = getattr(self.model.config, "num_labels", None)
|
|
else:
|
|
# Use mlx-embeddings for other architectures (ModernBert, etc.)
|
|
patch_qwen3_vl_processor_for_torch_free_image_loading()
|
|
from mlx_embeddings import load
|
|
|
|
self.model, self.processor = load(
|
|
self.model_name,
|
|
tokenizer_config={"trust_remote_code": self.trust_remote_code},
|
|
)
|
|
|
|
# Get num_labels from model config
|
|
if hasattr(self.model, "config"):
|
|
config = self.model.config
|
|
self._num_labels = getattr(config, "num_labels", None)
|
|
|
|
# Try mx.compile for persistent Metal kernel caching
|
|
self._is_compiled = self._try_compile()
|
|
|
|
self._loaded = True
|
|
logger.info(
|
|
f"Reranker model loaded successfully: {self.model_name} "
|
|
f"(arch={arch}, num_labels={self._num_labels}, "
|
|
f"causal_lm={self._is_causal_lm}, vl={self._is_vl_reranker}, "
|
|
f"compiled={self._is_compiled})"
|
|
)
|
|
|
|
except ImportError as e:
|
|
raise ImportError(
|
|
"mlx-lm, mlx-embeddings, or transformers is required for reranking. "
|
|
"Install with: pip install mlx-lm mlx-embeddings transformers"
|
|
) from e
|
|
except FileNotFoundError:
|
|
raise FileNotFoundError(
|
|
f"No safetensors weight files found for '{self.model_name}'. "
|
|
f"Reranker 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 reranker model: {e}")
|
|
raise
|
|
|
|
def _try_compile(self) -> bool:
|
|
"""Compile reranker scoring path to return primitive logits arrays.
|
|
|
|
Root-cause fix:
|
|
- Compiling model.__call__ directly can yield arrays without primitives
|
|
in some MLX output containers.
|
|
- Compile a narrow function that returns logits only.
|
|
"""
|
|
if self._is_causal_lm or self._is_vl_reranker:
|
|
# CausalLM / VL reranker paths use custom scoring (yes/no logits or
|
|
# mlx-embeddings model.process). VL forward needs pixel_values and
|
|
# lacks pooler_output, so the compile wrapper here wouldn't apply.
|
|
logger.info(f"mx.compile skipped for {self.model_name}")
|
|
self._compiled_seq_logits = None
|
|
return False
|
|
|
|
base_model = self.model
|
|
if not callable(base_model):
|
|
return False
|
|
try:
|
|
|
|
def _compiled_seq_logits(inputs):
|
|
outputs = base_model(**inputs)
|
|
if (
|
|
hasattr(outputs, "pooler_output")
|
|
and outputs.pooler_output is not None
|
|
):
|
|
return outputs.pooler_output
|
|
raise ValueError(
|
|
"Model output does not contain pooler_output. "
|
|
"Ensure the model is a SequenceClassification model."
|
|
)
|
|
|
|
# NOTE: use default compile mode. shapeless=True can fail shape
|
|
# inference for some linear ops in embedding/reranker stacks.
|
|
self._compiled_seq_logits = mx.compile(_compiled_seq_logits)
|
|
|
|
# Warmup: verify compilation actually works with a dummy forward pass
|
|
test_inputs = {
|
|
"input_ids": mx.zeros((1, 4), dtype=mx.int32),
|
|
"attention_mask": mx.ones((1, 4), dtype=mx.int32),
|
|
}
|
|
_ = self._compiled_seq_logits(test_inputs)
|
|
|
|
logger.info(
|
|
f"mx.compile enabled for {self.model_name} "
|
|
f"(primitive reranker logits path)"
|
|
)
|
|
return True
|
|
except Exception as e:
|
|
logger.info(f"mx.compile unavailable for {self.model_name}: {e}")
|
|
self._compiled_seq_logits = None
|
|
return False
|
|
|
|
def close(self) -> None:
|
|
"""Release model, processor, projector, and compiled reranker resources."""
|
|
self._compiled_seq_logits = None
|
|
self._is_compiled = False
|
|
|
|
self.model = None
|
|
self.processor = None
|
|
self._loaded = False
|
|
self._num_labels = None
|
|
self._is_causal_lm = False
|
|
self._is_jina_reranker = False
|
|
self._is_vl_reranker = False
|
|
self._token_true_id = None
|
|
self._token_false_id = None
|
|
self._doc_embed_token_id = None
|
|
self._query_embed_token_id = None
|
|
self._jina_projector = None
|
|
self._prefix_tokens = None
|
|
self._suffix_tokens = None
|
|
|
|
gc.collect()
|
|
mx.synchronize()
|
|
mx.clear_cache()
|
|
clear_thread_compile_cache()
|
|
gc.collect()
|
|
|
|
# Default max_length per model type
|
|
_DEFAULT_MAX_LENGTH_SEQ_CLASSIFICATION = 512
|
|
_DEFAULT_MAX_LENGTH_CAUSAL_LM = 8192
|
|
|
|
def rerank(
|
|
self,
|
|
query: "str | dict",
|
|
documents: "list[str] | list[dict]",
|
|
max_length: int | None = None,
|
|
) -> RerankOutput:
|
|
"""
|
|
Rerank documents by relevance to the query.
|
|
|
|
Args:
|
|
query: The search query. String for text-only rerankers. Dict with
|
|
'text' and/or 'image' for multimodal rerankers.
|
|
documents: List of documents to rerank. Each item can be a string
|
|
or a dict with 'text' and/or 'image' keys.
|
|
max_length: Maximum token length for each query-document pair.
|
|
If None, uses model-appropriate default (512 for encoder,
|
|
8192 for CausalLM).
|
|
|
|
Returns:
|
|
RerankOutput with scores, sorted indices, and token count
|
|
"""
|
|
if not self._loaded:
|
|
self.load()
|
|
|
|
if not documents:
|
|
return RerankOutput(scores=[], indices=[], total_tokens=0)
|
|
|
|
if self._is_vl_reranker:
|
|
effective_max_length = (
|
|
max_length
|
|
if max_length is not None
|
|
else self._DEFAULT_MAX_LENGTH_CAUSAL_LM
|
|
)
|
|
return self._rerank_vl(query, documents, effective_max_length)
|
|
|
|
# Text-only paths: coerce dict inputs down to text so existing
|
|
# _rerank_* methods keep their str-only contract.
|
|
query_str = _coerce_item_to_text(query)
|
|
docs_str = [_coerce_item_to_text(d) for d in documents]
|
|
|
|
if self._is_jina_reranker:
|
|
effective_max_length = (
|
|
max_length
|
|
if max_length is not None
|
|
else self._DEFAULT_MAX_LENGTH_CAUSAL_LM
|
|
)
|
|
return self._rerank_jina(query_str, docs_str, effective_max_length)
|
|
elif self._is_causal_lm:
|
|
effective_max_length = (
|
|
max_length
|
|
if max_length is not None
|
|
else self._DEFAULT_MAX_LENGTH_CAUSAL_LM
|
|
)
|
|
return self._rerank_causal_lm(query_str, docs_str, effective_max_length)
|
|
else:
|
|
effective_max_length = (
|
|
max_length
|
|
if max_length is not None
|
|
else self._DEFAULT_MAX_LENGTH_SEQ_CLASSIFICATION
|
|
)
|
|
return self._rerank_seq_classification(
|
|
query_str, docs_str, effective_max_length
|
|
)
|
|
|
|
def _rerank_causal_lm(
|
|
self,
|
|
query: str,
|
|
documents: list[str],
|
|
max_length: int = 8192,
|
|
) -> RerankOutput:
|
|
"""
|
|
Rerank using CausalLM yes/no logit scoring (e.g., Qwen3-Reranker).
|
|
|
|
Constructs instruction prompts, runs per-document forward passes, and
|
|
extracts relevance scores from the logits of yes/no tokens at the last
|
|
position. Each document is processed individually since mlx-lm models
|
|
generate their own causal mask internally and don't accept an external
|
|
padding mask.
|
|
"""
|
|
import mlx.core as mx
|
|
|
|
tokenizer = self.processor
|
|
prefix_tokens = self._prefix_tokens
|
|
suffix_tokens = self._suffix_tokens
|
|
if not callable(tokenizer):
|
|
raise ValueError("CausalLM reranker tokenizer is not initialized.")
|
|
if prefix_tokens is None or suffix_tokens is None:
|
|
raise ValueError("CausalLM reranker prompt tokens are not initialized.")
|
|
if not callable(self.model):
|
|
raise ValueError("CausalLM reranker model is not initialized.")
|
|
|
|
# Compute max tokens available for the instruction content
|
|
max_content_tokens = max_length - len(prefix_tokens) - len(suffix_tokens)
|
|
|
|
# Format and tokenize each query-document pair
|
|
pairs_text = []
|
|
for doc in documents:
|
|
content = (
|
|
f"<Instruct>: {self._CAUSAL_LM_DEFAULT_INSTRUCTION}\n"
|
|
f"<Query>: {query}\n"
|
|
f"<Document>: {doc}"
|
|
)
|
|
pairs_text.append(content)
|
|
|
|
# Tokenize content parts (without prefix/suffix)
|
|
content_encodings = tokenizer(
|
|
pairs_text,
|
|
padding=False,
|
|
truncation=True,
|
|
return_attention_mask=False,
|
|
max_length=max_content_tokens,
|
|
add_special_tokens=False,
|
|
)
|
|
|
|
# Assemble full token sequences: prefix + content + suffix
|
|
all_input_ids = []
|
|
for content_ids in content_encodings["input_ids"]:
|
|
full_ids = prefix_tokens + content_ids + suffix_tokens
|
|
all_input_ids.append(full_ids)
|
|
|
|
# Per-document forward pass and score extraction.
|
|
# mlx-lm models generate their own causal attention mask internally
|
|
# and don't support external padding masks, so we process each
|
|
# document individually to ensure correct attention computation.
|
|
scores = []
|
|
total_tokens = 0
|
|
for ids in all_input_ids:
|
|
input_ids = mx.array([ids]) # (1, seq_len)
|
|
logits = self.model(input_ids)
|
|
# Extract yes/no logits at the last position
|
|
last_logits = logits[0, -1, :]
|
|
true_logit = last_logits[self._token_true_id]
|
|
false_logit = last_logits[self._token_false_id]
|
|
paired = mx.array([false_logit, true_logit])
|
|
probs = mx.softmax(paired)
|
|
mx.eval(probs)
|
|
scores.append(probs[1].item())
|
|
total_tokens += len(ids)
|
|
|
|
# Sort indices by score (descending)
|
|
indexed_scores = list(enumerate(scores))
|
|
indexed_scores.sort(key=lambda x: x[1], reverse=True)
|
|
sorted_indices = [idx for idx, _ in indexed_scores]
|
|
|
|
return RerankOutput(
|
|
scores=scores,
|
|
indices=sorted_indices,
|
|
total_tokens=total_tokens,
|
|
)
|
|
|
|
def _rerank_jina(
|
|
self,
|
|
query: str,
|
|
documents: list[str],
|
|
max_length: int = 8192,
|
|
) -> RerankOutput:
|
|
"""
|
|
Rerank using Jina v3 listwise embedding-based scoring.
|
|
|
|
Builds multi-document prompts, extracts hidden states at special token
|
|
positions, applies the projector, and computes query-document cosine
|
|
similarities. Uses deterministic greedy chunking under max_length.
|
|
"""
|
|
tokenizer = self.processor
|
|
doc_embed_token_id = self._doc_embed_token_id
|
|
query_embed_token_id = self._query_embed_token_id
|
|
projector = self._jina_projector
|
|
if tokenizer is None:
|
|
raise ValueError("Jina reranker tokenizer is not initialized.")
|
|
|
|
encode = getattr(tokenizer, "encode", None)
|
|
if not callable(encode):
|
|
raise ValueError("Jina reranker tokenizer does not provide encode().")
|
|
|
|
if (
|
|
doc_embed_token_id is None
|
|
or query_embed_token_id is None
|
|
or projector is None
|
|
):
|
|
raise ValueError(
|
|
"Jina reranker is not fully initialized. "
|
|
"Missing special-token IDs or projector."
|
|
)
|
|
|
|
def _to_token_ids(text: str) -> list[int]:
|
|
encoded = encode(text, add_special_tokens=False)
|
|
if hasattr(encoded, "ids"):
|
|
return list(encoded.ids)
|
|
return list(encoded)
|
|
|
|
decode = getattr(tokenizer, "decode", None)
|
|
|
|
def _truncate_doc_to_fit(
|
|
query_text: str, doc_text: str
|
|
) -> Tuple[str, list[int]]:
|
|
doc_token_ids = _to_token_ids(doc_text)
|
|
if not doc_token_ids:
|
|
prompt = self._format_jina_prompt(query_text, [""])
|
|
prompt_ids = _to_token_ids(prompt)[:max_length]
|
|
return "", prompt_ids
|
|
|
|
best_doc = ""
|
|
best_ids: list[int] = []
|
|
lo = 0
|
|
hi = len(doc_token_ids)
|
|
while lo <= hi:
|
|
mid = (lo + hi) // 2
|
|
if callable(decode):
|
|
candidate_doc = decode(
|
|
doc_token_ids[:mid], skip_special_tokens=False
|
|
)
|
|
else:
|
|
candidate_doc = doc_text[:mid]
|
|
|
|
prompt = self._format_jina_prompt(query_text, [candidate_doc])
|
|
prompt_ids = _to_token_ids(prompt)
|
|
if len(prompt_ids) <= max_length:
|
|
best_doc = candidate_doc
|
|
best_ids = prompt_ids
|
|
lo = mid + 1
|
|
else:
|
|
hi = mid - 1
|
|
|
|
if not best_ids:
|
|
raise ValueError(
|
|
"Could not fit even a minimally truncated document into max_length. "
|
|
f"max_length={max_length}"
|
|
)
|
|
|
|
return best_doc, best_ids
|
|
|
|
sanitized_query = self._sanitize_jina_text(query)
|
|
sanitized_docs = [self._sanitize_jina_text(doc) for doc in documents]
|
|
|
|
scores = [0.0] * len(documents)
|
|
total_tokens = 0
|
|
start = 0
|
|
while start < len(sanitized_docs):
|
|
chunk_doc_indices: list[int] = []
|
|
chunk_docs: list[str] = []
|
|
chunk_input_ids: list[int] | None = None
|
|
cursor = start
|
|
|
|
while cursor < len(sanitized_docs):
|
|
candidate_docs = chunk_docs + [sanitized_docs[cursor]]
|
|
candidate_prompt = self._format_jina_prompt(
|
|
sanitized_query, candidate_docs
|
|
)
|
|
candidate_ids = _to_token_ids(candidate_prompt)
|
|
|
|
if len(candidate_ids) <= max_length:
|
|
chunk_docs = candidate_docs
|
|
chunk_doc_indices.append(cursor)
|
|
chunk_input_ids = candidate_ids
|
|
cursor += 1
|
|
continue
|
|
|
|
if chunk_docs:
|
|
break
|
|
|
|
truncated_doc, truncated_ids = _truncate_doc_to_fit(
|
|
sanitized_query,
|
|
sanitized_docs[cursor],
|
|
)
|
|
chunk_docs = [truncated_doc]
|
|
chunk_doc_indices = [cursor]
|
|
chunk_input_ids = truncated_ids
|
|
cursor += 1
|
|
break
|
|
|
|
if chunk_input_ids is None or not chunk_doc_indices:
|
|
raise ValueError("Failed to create a valid Jina reranker chunk.")
|
|
|
|
input_array = mx.array([chunk_input_ids])
|
|
hidden_states = self._get_jina_hidden_states(input_array)
|
|
|
|
query_positions = [
|
|
pos
|
|
for pos, token_id in enumerate(chunk_input_ids)
|
|
if token_id == query_embed_token_id
|
|
]
|
|
if not query_positions:
|
|
raise ValueError(
|
|
"Jina prompt does not contain '<|rerank_token|>' in tokenized input."
|
|
)
|
|
|
|
doc_positions = [
|
|
pos
|
|
for pos, token_id in enumerate(chunk_input_ids)
|
|
if token_id == doc_embed_token_id
|
|
]
|
|
if len(doc_positions) < len(chunk_docs):
|
|
raise ValueError(
|
|
"Jina prompt/doc mismatch: detected fewer '<|embed_token|>' "
|
|
"positions than documents in chunk."
|
|
)
|
|
|
|
selected_doc_positions = doc_positions[: len(chunk_docs)]
|
|
query_hidden = hidden_states[0, query_positions[0], :]
|
|
doc_hidden = hidden_states[0, selected_doc_positions, :]
|
|
|
|
query_vec = projector(query_hidden)
|
|
doc_vecs = projector(doc_hidden)
|
|
similarities = self._cosine_similarity(query_vec, doc_vecs)
|
|
mx.eval(similarities)
|
|
|
|
chunk_scores = similarities.tolist()
|
|
for original_idx, score in zip(chunk_doc_indices, chunk_scores):
|
|
scores[original_idx] = float(score)
|
|
|
|
total_tokens += len(chunk_input_ids)
|
|
start = cursor
|
|
|
|
# Sort by score descending
|
|
indexed_scores = list(enumerate(scores))
|
|
indexed_scores.sort(key=lambda x: x[1], reverse=True)
|
|
sorted_indices = [idx for idx, _ in indexed_scores]
|
|
|
|
return RerankOutput(
|
|
scores=scores,
|
|
indices=sorted_indices,
|
|
total_tokens=total_tokens,
|
|
)
|
|
|
|
def _rerank_seq_classification(
|
|
self,
|
|
query: str,
|
|
documents: list[str],
|
|
max_length: int = 512,
|
|
) -> RerankOutput:
|
|
"""Rerank using SequenceClassification models (encoder-based)."""
|
|
import mlx.core as mx
|
|
|
|
# Get the underlying tokenizer from TokenizerWrapper (mlx-embeddings only)
|
|
# Don't unwrap transformers tokenizers which also have _tokenizer attribute
|
|
processor = self.processor
|
|
processor_class = type(processor).__name__
|
|
if processor_class == "TokenizerWrapper" and hasattr(processor, "_tokenizer"):
|
|
processor = processor._tokenizer
|
|
if not callable(processor):
|
|
raise ValueError("SequenceClassification processor is not initialized.")
|
|
|
|
# Tokenize query-document pairs
|
|
# SequenceClassification models expect pairs as (query, document)
|
|
pairs = [(query, doc) for doc in documents]
|
|
|
|
# Batch encode all pairs
|
|
inputs = processor(
|
|
[p[0] for p in pairs],
|
|
[p[1] for p in pairs],
|
|
max_length=max_length,
|
|
padding=True,
|
|
truncation=True,
|
|
return_tensors="np",
|
|
)
|
|
|
|
# Convert to MLX arrays
|
|
input_ids = mx.array(inputs["input_ids"])
|
|
attention_mask = mx.array(inputs["attention_mask"])
|
|
|
|
# Forward pass (compiled primitive logits path when available)
|
|
logits = None
|
|
if self._is_compiled and self._compiled_seq_logits is not None:
|
|
try:
|
|
model_inputs = {
|
|
"input_ids": input_ids,
|
|
"attention_mask": attention_mask,
|
|
}
|
|
logits = self._compiled_seq_logits(model_inputs)
|
|
except Exception as e:
|
|
logger.warning(
|
|
f"compiled reranker path failed for {self.model_name}: {e}; "
|
|
f"disabling compile and falling back to eager forward()"
|
|
)
|
|
self._is_compiled = False
|
|
self._compiled_seq_logits = None
|
|
|
|
if logits is None:
|
|
if not callable(self.model):
|
|
raise ValueError("SequenceClassification model is not initialized.")
|
|
outputs = self.model(
|
|
input_ids=input_ids,
|
|
attention_mask=attention_mask,
|
|
)
|
|
|
|
# Extract scores from pooler_output
|
|
# pooler_output shape: (batch_size, num_labels)
|
|
if hasattr(outputs, "pooler_output") and outputs.pooler_output is not None:
|
|
logits = outputs.pooler_output
|
|
else:
|
|
raise ValueError(
|
|
"Model output does not contain pooler_output. "
|
|
"Ensure the model is a SequenceClassification model."
|
|
)
|
|
|
|
# Ensure computation is done
|
|
mx.eval(logits)
|
|
|
|
# Extract relevance scores
|
|
# For binary classification (num_labels=1), score is already sigmoid applied
|
|
# For multi-class, take the positive class probability
|
|
if logits.shape[-1] == 1:
|
|
# Binary classification: sigmoid already applied by model
|
|
scores = logits.squeeze(-1).tolist()
|
|
else:
|
|
# Multi-class: take last column (typically "relevant" class)
|
|
scores = logits[:, -1].tolist()
|
|
|
|
# Sort indices by score (descending)
|
|
indexed_scores = list(enumerate(scores))
|
|
indexed_scores.sort(key=lambda x: x[1], reverse=True)
|
|
sorted_indices = [idx for idx, _ in indexed_scores]
|
|
|
|
# Count tokens
|
|
total_tokens = self._count_tokens(query, documents)
|
|
|
|
return RerankOutput(
|
|
scores=scores,
|
|
indices=sorted_indices,
|
|
total_tokens=total_tokens,
|
|
)
|
|
|
|
def _count_tokens(self, query: str, documents: list[str]) -> int:
|
|
"""Count total tokens in query-document pairs."""
|
|
total = 0
|
|
|
|
processor = self.processor
|
|
processor_class = type(processor).__name__
|
|
if processor_class == "TokenizerWrapper" and hasattr(processor, "_tokenizer"):
|
|
processor = processor._tokenizer
|
|
|
|
def get_token_count(text: str, add_special: bool = True) -> int:
|
|
"""Get token count for text, handling different tokenizer types."""
|
|
if hasattr(processor, "encode"):
|
|
tokens = processor.encode(text, add_special_tokens=add_special)
|
|
# Handle different return types
|
|
if isinstance(tokens, list):
|
|
return len(tokens)
|
|
elif hasattr(tokens, "ids"):
|
|
# tokenizers.Encoding object
|
|
return len(tokens.ids)
|
|
else:
|
|
return len(tokens)
|
|
else:
|
|
# Fallback to word count estimate
|
|
return len(text.split()) + (2 if add_special else 0)
|
|
|
|
# Count query tokens once
|
|
query_len = get_token_count(query, add_special=True)
|
|
|
|
# Count document tokens
|
|
for doc in documents:
|
|
doc_len = get_token_count(doc, add_special=False)
|
|
# Each pair includes query + doc + special tokens
|
|
total += query_len + doc_len + 3 # [CLS], [SEP], [SEP]
|
|
|
|
return total
|
|
|
|
@property
|
|
def num_labels(self) -> int | None:
|
|
"""Get the number of classification labels."""
|
|
return self._num_labels
|
|
|
|
def _validate_architecture(self) -> None:
|
|
"""
|
|
Validate that the model architecture is supported.
|
|
|
|
Raises:
|
|
ValueError: If the architecture is not supported
|
|
"""
|
|
config_path = Path(self.model_name) / "config.json"
|
|
if not config_path.exists():
|
|
# If no config.json, let mlx-embeddings handle validation
|
|
return
|
|
|
|
try:
|
|
with open(config_path) as f:
|
|
config = json.load(f)
|
|
except (json.JSONDecodeError, IOError) as e:
|
|
logger.warning(f"Failed to read config.json: {e}")
|
|
return
|
|
|
|
architectures = config.get("architectures", [])
|
|
if not architectures:
|
|
return
|
|
|
|
arch = architectures[0]
|
|
|
|
# CausalLM reranker architectures require the directory name heuristic
|
|
# to distinguish from regular LLMs with the same architecture.
|
|
if arch in CAUSAL_LM_RERANKER_ARCHITECTURES:
|
|
if not _is_causal_lm_reranker(Path(self.model_name)):
|
|
raise ValueError(
|
|
f"Architecture {arch} is a CausalLM that can be used as a "
|
|
f"reranker, but the model directory name "
|
|
f"'{Path(self.model_name).name}' does not contain "
|
|
f"'reranker' or 'rerank'. Please rename the directory or "
|
|
f"use the correct model."
|
|
)
|
|
return
|
|
|
|
# Multimodal reranker architectures share the arch string with VLM chat
|
|
# models; use the same dir-name heuristic to disambiguate.
|
|
if arch in MULTIMODAL_RERANKER_ARCHITECTURES:
|
|
if not _is_causal_lm_reranker(Path(self.model_name)):
|
|
raise ValueError(
|
|
f"Architecture {arch} is a VLM that can be used as a "
|
|
f"reranker, but the model directory name "
|
|
f"'{Path(self.model_name).name}' does not contain "
|
|
f"'reranker' or 'rerank'. Please rename the directory or "
|
|
f"use the correct model."
|
|
)
|
|
return
|
|
|
|
if arch not in SUPPORTED_RERANKER_ARCHITECTURES:
|
|
supported_list = ", ".join(
|
|
sorted(
|
|
SUPPORTED_RERANKER_ARCHITECTURES
|
|
| CAUSAL_LM_RERANKER_ARCHITECTURES
|
|
| MULTIMODAL_RERANKER_ARCHITECTURES
|
|
)
|
|
)
|
|
raise ValueError(
|
|
f"Unsupported reranker architecture: {arch}. "
|
|
f"Currently supported architectures: {supported_list}."
|
|
)
|
|
|
|
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,
|
|
"num_labels": self._num_labels,
|
|
}
|
|
|
|
# Try to get model config
|
|
if hasattr(self.model, "config"):
|
|
config = self.model.config
|
|
info.update(
|
|
{
|
|
"model_type": getattr(config, "model_type", None),
|
|
"hidden_size": getattr(config, "hidden_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"
|
|
return f"<MLXRerankerModel model={self.model_name} status={status}>"
|