vllm_mlx.rerank

Reranker engine for cross-encoder models.

View the complete module source at #L1-L398.

API details

Each callable below includes its exact signature, type annotations, inputs, defaults, return contract, documented exceptions, implementation source, and parsed docstring sections when the source provides them.

vllm_mlx.rerank

Reranker engine for cross-encoder models.

Provides a dedicated RerankEngine with adapter-based scoring for the OpenAI/Jina-compatible /v1/rerank endpoint. Cross-encoder models use AutoModelForSequenceClassification-style loading, not mlx_lm.load.

vllm_mlx.rerank.logger module-attribute

logger = logging.getLogger(__name__)

vllm_mlx.rerank._ADAPTER_REGISTRY module-attribute

_ADAPTER_REGISTRY: dict[str, type[RerankAdapter]] = {'default': SigmoidAdapter}

vllm_mlx.rerank.RerankAdapter

Bases: ABC

Per-family adapter for reranker models.

Different cross-encoder families use different tokenization patterns, score extraction logic, and normalization functions. This contract isolates those differences so RerankEngine stays family-agnostic.

vllm_mlx.rerank.RerankAdapter.tokenize_pair abstractmethod

tokenize_pair(tokenizer, query: str, document: str, max_length: int) -> dict

Tokenize a (query, document) pair for the cross-encoder.

Parameters:

• tokenizer –

  The HuggingFace tokenizer instance.

• query (str) –

  The query string.

• document (str) –

  The document string.

• max_length (int) –

  Truncation length (from the model's context window).

Returns:

• dict –

  Dict with 'input_ids' and 'attention_mask' as numpy arrays.

Source code in vllm_mlx/rerank.py

@abstractmethod
def tokenize_pair(
    self, tokenizer, query: str, document: str, max_length: int
) -> dict:
    """
    Tokenize a (query, document) pair for the cross-encoder.

    Args:
        tokenizer: The HuggingFace tokenizer instance.
        query: The query string.
        document: The document string.
        max_length: Truncation length (from the model's context window).

    Returns:
        Dict with 'input_ids' and 'attention_mask' as numpy arrays.
    """
    ...

vllm_mlx.rerank.RerankAdapter.extract_score abstractmethod

extract_score(logits) -> float

Extract a raw relevance score from model output logits.

Parameters:

• logits –

  Model output logits (list or array), shape varies by model.

Returns:

• float –

  A single float raw score.

Source code in vllm_mlx/rerank.py

@abstractmethod
def extract_score(self, logits) -> float:
    """
    Extract a raw relevance score from model output logits.

    Args:
        logits: Model output logits (list or array), shape varies by model.

    Returns:
        A single float raw score.
    """
    ...

vllm_mlx.rerank.RerankAdapter.normalize abstractmethod

normalize(raw_score: float) -> float

Normalize a raw score to [0, 1] range.

Parameters:

• raw_score (float) –

  The raw score from extract_score().

Returns:

• float –

  Normalized relevance score in [0, 1].

Source code in vllm_mlx/rerank.py

@abstractmethod
def normalize(self, raw_score: float) -> float:
    """
    Normalize a raw score to [0, 1] range.

    Args:
        raw_score: The raw score from extract_score().

    Returns:
        Normalized relevance score in [0, 1].
    """
    ...

vllm_mlx.rerank.SigmoidAdapter

Bases: RerankAdapter

Default adapter for single-logit sigmoid rerankers.

Works with Jina Reranker v2, BGE Reranker v2, and MS-MARCO MiniLM families. These models output a single relevance logit at position 0, normalized via sigmoid.

vllm_mlx.rerank.SigmoidAdapter.tokenize_pair

tokenize_pair(tokenizer, query: str, document: str, max_length: int) -> dict

Tokenize as a sentence pair (query, document).

Source code in vllm_mlx/rerank.py

def tokenize_pair(
    self, tokenizer, query: str, document: str, max_length: int
) -> dict:
    """Tokenize as a sentence pair (query, document)."""
    return tokenizer(
        query,
        document,
        padding=True,
        truncation=True,
        max_length=max_length,
        return_tensors="np",
    )

vllm_mlx.rerank.SigmoidAdapter.extract_score

extract_score(logits) -> float

Extract the first logit as the relevance score.

Source code in vllm_mlx/rerank.py

def extract_score(self, logits) -> float:
    """Extract the first logit as the relevance score."""
    return float(logits[0])

vllm_mlx.rerank.SigmoidAdapter.normalize

normalize(raw_score: float) -> float

Apply sigmoid normalization.

Source code in vllm_mlx/rerank.py

def normalize(self, raw_score: float) -> float:
    """Apply sigmoid normalization."""
    return 1.0 / (1.0 + math.exp(-raw_score))

vllm_mlx.rerank.RerankEngine

RerankEngine(model_name: str, token_budget: int = 4096, max_concurrency: int = 1)

Reranker engine for cross-encoder sequence classification models.

Loads cross-encoder models via transformers + MLX (safetensors weights). Scores (query, document) pairs using the adapter contract for family-specific tokenization, score extraction, and normalization.

Supports token-budget batching to avoid OOM on large document lists.

Source code in vllm_mlx/rerank.py

def __init__(
    self,
    model_name: str,
    token_budget: int = 4096,
    max_concurrency: int = 1,
):
    self.model_name = model_name
    self.token_budget = token_budget
    self.max_concurrency = max_concurrency
    self._semaphore = asyncio.Semaphore(max_concurrency)
    self._model = None
    self._tokenizer = None
    self._adapter: RerankAdapter | None = None

vllm_mlx.rerank.RerankEngine.model_name instance-attribute

model_name = model_name

vllm_mlx.rerank.RerankEngine.token_budget instance-attribute

token_budget = token_budget

vllm_mlx.rerank.RerankEngine.max_concurrency instance-attribute

max_concurrency = max_concurrency

vllm_mlx.rerank.RerankEngine._semaphore instance-attribute

_semaphore = asyncio.Semaphore(max_concurrency)

vllm_mlx.rerank.RerankEngine._model instance-attribute

_model = None

vllm_mlx.rerank.RerankEngine._tokenizer instance-attribute

_tokenizer = None

vllm_mlx.rerank.RerankEngine._adapter instance-attribute

_adapter: RerankAdapter | None = None

vllm_mlx.rerank.RerankEngine.is_loaded property

is_loaded: bool

Return whether the reranking model has been loaded.

vllm_mlx.rerank.RerankEngine.load

load() -> None

Load the cross-encoder model and tokenizer.

Uses transformers AutoTokenizer and loads MLX weights from safetensors via the model's from_pretrained or equivalent MLX loading path.

Source code in vllm_mlx/rerank.py

def load(self) -> None:
    """
    Load the cross-encoder model and tokenizer.

    Uses transformers AutoTokenizer and loads MLX weights from safetensors
    via the model's from_pretrained or equivalent MLX loading path.
    """
    from transformers import AutoTokenizer

    logger.info(f"Loading reranker model: {self.model_name}")
    start = time.perf_counter()

    self._tokenizer = AutoTokenizer.from_pretrained(self.model_name)
    self._model = self._load_mlx_model(self.model_name)
    self._adapter = get_adapter(self.model_name)

    elapsed = time.perf_counter() - start
    logger.info(f"Reranker model loaded in {elapsed:.2f}s: {self.model_name}")

vllm_mlx.rerank.RerankEngine._load_mlx_model staticmethod

_load_mlx_model(model_name: str)

Load an MLX cross-encoder model from HuggingFace Hub.

Attempts mlx-community weights first (safetensors), then falls back to transformers AutoModelForSequenceClassification with MLX conversion.

Source code in vllm_mlx/rerank.py

@staticmethod
def _load_mlx_model(model_name: str):
    """
    Load an MLX cross-encoder model from HuggingFace Hub.

    Attempts mlx-community weights first (safetensors), then falls back
    to transformers AutoModelForSequenceClassification with MLX conversion.
    """
    try:
        from huggingface_hub import snapshot_download
        from safetensors import safe_open

        model_path = snapshot_download(model_name)

        import glob
        import json
        import os

        # Load model config
        config_path = os.path.join(model_path, "config.json")
        with open(config_path) as f:
            config = json.load(f)

        # Load weights from safetensors
        weight_files = glob.glob(os.path.join(model_path, "*.safetensors"))
        if not weight_files:
            raise FileNotFoundError(f"No safetensors files found in {model_path}")

        weights = {}
        for wf in weight_files:
            with safe_open(wf, framework="numpy") as f:
                for key in f.keys():
                    weights[key] = mx.array(f.get_tensor(key))

        # Build model based on architecture
        model_type = config.get("model_type", "")
        num_labels = config.get("num_labels", 1)

        model = _build_classifier_model(model_type, config, weights, num_labels)
        mx.eval(model.parameters())
        return model

    except Exception as e:
        logger.error(f"Failed to load reranker model: {e}")
        raise

vllm_mlx.rerank.RerankEngine._ensure_loaded

_ensure_loaded() -> None

Source code in vllm_mlx/rerank.py

def _ensure_loaded(self) -> None:
    if not self.is_loaded:
        self.load()

vllm_mlx.rerank.RerankEngine.score_pairs

score_pairs(query: str, documents: list[str]) -> tuple[list[float], int]

Score each (query, document) pair and return normalized relevance scores.

Pairs are batched by token budget to control memory usage. Each batch is tokenized together and scored in a single forward pass. Returns (scores, total_tokens) where total_tokens reflects the actual tokenization used for scoring (consistent with adapter).

Parameters:

• query (str) –

  The query string.

• documents (list[str]) –

  List of document strings.

Returns:

• list[float] –

  List of normalized relevance scores, one per document,

• int –

  in the same order as the input documents.

Source code in vllm_mlx/rerank.py

def score_pairs(self, query: str, documents: list[str]) -> tuple[list[float], int]:
    """
    Score each (query, document) pair and return normalized relevance scores.

    Pairs are batched by token budget to control memory usage. Each batch
    is tokenized together and scored in a single forward pass.
    Returns (scores, total_tokens) where total_tokens reflects the
    actual tokenization used for scoring (consistent with adapter).

    Args:
        query: The query string.
        documents: List of document strings.

    Returns:
        List of normalized relevance scores, one per document,
        in the same order as the input documents.
    """
    self._ensure_loaded()

    max_length = resolve_max_length(
        getattr(self._model, "config", None),
        self._tokenizer,
    )

    # Tokenize each pair individually to measure token counts
    pair_encodings = []
    pair_token_counts = []
    for doc in documents:
        enc = self._adapter.tokenize_pair(self._tokenizer, query, doc, max_length)
        pair_encodings.append(enc)
        seq_len = (
            len(enc["input_ids"][0])
            if hasattr(enc["input_ids"][0], "__len__")
            else enc["input_ids"].shape[1]
        )
        pair_token_counts.append(seq_len)

    # Build batches by token budget
    batches = []
    current_batch = []
    current_tokens = 0
    for i, (enc, tok_count) in enumerate(zip(pair_encodings, pair_token_counts)):
        if current_batch and current_tokens + tok_count > self.token_budget:
            batches.append(current_batch)
            current_batch = []
            current_tokens = 0
        current_batch.append((i, enc))
        current_tokens += tok_count
    if current_batch:
        batches.append(current_batch)

    # Score each batch
    all_scores: list[tuple[int, float]] = []
    for batch in batches:
        if len(batch) == 1:
            # Single pair — use encoding directly
            idx, enc = batch[0]
            input_ids = mx.array(enc["input_ids"])
            attention_mask = mx.array(enc["attention_mask"])
        else:
            # Pad and stack multiple pairs
            max_len = max(
                (
                    len(enc["input_ids"][0])
                    if hasattr(enc["input_ids"][0], "__len__")
                    else enc["input_ids"].shape[1]
                )
                for _, enc in batch
            )
            padded_ids = []
            padded_mask = []
            for _, enc in batch:
                raw_ids = enc["input_ids"][0]
                ids = (
                    raw_ids.tolist()
                    if hasattr(raw_ids, "tolist")
                    else list(raw_ids)
                )
                raw_mask = enc["attention_mask"][0]
                mask = (
                    raw_mask.tolist()
                    if hasattr(raw_mask, "tolist")
                    else list(raw_mask)
                )
                pad_len = max_len - len(ids)
                padded_ids.append(ids + [0] * pad_len)
                padded_mask.append(mask + [0] * pad_len)
            input_ids = mx.array(padded_ids)
            attention_mask = mx.array(padded_mask)

        output = self._model(input_ids, attention_mask=attention_mask)
        logits_list = output.logits.tolist()

        for j, (idx, _enc) in enumerate(batch):
            logits_row = logits_list[j] if len(batch) > 1 else logits_list[0]
            raw_score = self._adapter.extract_score(logits_row)
            normalized = self._adapter.normalize(raw_score)
            all_scores.append((idx, normalized))

    # Sort by original index to restore input order
    all_scores.sort(key=lambda x: x[0])
    total_tokens = sum(pair_token_counts)
    return [score for _, score in all_scores], total_tokens

vllm_mlx.rerank._MLXClassifierWrapper

_MLXClassifierWrapper(config: dict, weights: dict, num_labels: int)

Minimal MLX wrapper for sequence classification models.

Wraps loaded safetensors weights into a callable that returns logits for (input_ids, attention_mask) pairs. Supports BERT-family and XLM-RoBERTa-family architectures commonly used as cross-encoders.

Source code in vllm_mlx/rerank.py

def __init__(self, config: dict, weights: dict, num_labels: int):
    self.config = config
    self.weights = weights
    self.num_labels = num_labels
    self._params = list(weights.values())

vllm_mlx.rerank._MLXClassifierWrapper.config instance-attribute

config = config

vllm_mlx.rerank._MLXClassifierWrapper.weights instance-attribute

weights = weights

vllm_mlx.rerank._MLXClassifierWrapper.num_labels instance-attribute

num_labels = num_labels

vllm_mlx.rerank._MLXClassifierWrapper._params instance-attribute

_params = list(weights.values())

vllm_mlx.rerank._MLXClassifierWrapper.parameters

parameters()

Return model parameters for mx.eval.

Source code in vllm_mlx/rerank.py

def parameters(self):
    """Return model parameters for mx.eval."""
    return self._params

vllm_mlx.rerank._MLXClassifierWrapper.__call__

__call__(input_ids: array, attention_mask: array = None)

Forward pass through the classifier.

For encoder-only cross-encoders, this runs the full transformer encoder and classification head. The exact layer wiring depends on the model architecture.

This initial implementation uses a weight-lookup forward pass that works for standard BERT/XLM-RoBERTa classifiers. For models with non-standard architectures, register a custom adapter via _ADAPTER_REGISTRY.

Source code in vllm_mlx/rerank.py

def __call__(self, input_ids: mx.array, attention_mask: mx.array = None):
    """
    Forward pass through the classifier.

    For encoder-only cross-encoders, this runs the full transformer
    encoder and classification head. The exact layer wiring depends
    on the model architecture.

    This initial implementation uses a weight-lookup forward pass
    that works for standard BERT/XLM-RoBERTa classifiers. For
    models with non-standard architectures, register a custom
    adapter via _ADAPTER_REGISTRY.
    """
    # Use the transformers-style weight naming convention
    # to walk through embeddings -> encoder layers -> classifier
    from vllm_mlx.rerank_forward import classifier_forward

    logits = classifier_forward(
        input_ids, attention_mask, self.weights, self.config
    )
    return _ClassifierOutput(logits=logits)

vllm_mlx.rerank._ClassifierOutput

_ClassifierOutput(logits: array)

Simple container for classifier output logits.

Source code in vllm_mlx/rerank.py

def __init__(self, logits: mx.array):
    self.logits = logits

vllm_mlx.rerank._ClassifierOutput.logits instance-attribute

logits = logits

vllm_mlx.rerank.get_adapter

get_adapter(model_name: str) -> RerankAdapter

Return the appropriate adapter for a model.

Falls back to SigmoidAdapter (works for Jina, BGE, MS-MARCO families). Extend _ADAPTER_REGISTRY for families that need different scoring.

Source code in vllm_mlx/rerank.py

def get_adapter(model_name: str) -> RerankAdapter:
    """
    Return the appropriate adapter for a model.

    Falls back to SigmoidAdapter (works for Jina, BGE, MS-MARCO families).
    Extend _ADAPTER_REGISTRY for families that need different scoring.
    """
    # Future: inspect model config to select adapter automatically.
    # For now, all known MLX reranker models use the sigmoid pattern.
    return _ADAPTER_REGISTRY["default"]()

vllm_mlx.rerank._build_classifier_model

_build_classifier_model(model_type, config, weights, num_labels)

Build an MLX sequence classification model from config and weights.

This is a thin wrapper that constructs the appropriate encoder architecture with a classification head on top.

Source code in vllm_mlx/rerank.py

def _build_classifier_model(model_type, config, weights, num_labels):
    """
    Build an MLX sequence classification model from config and weights.

    This is a thin wrapper that constructs the appropriate encoder
    architecture with a classification head on top.
    """
    # Import here to avoid top-level dependency on specific model implementations
    return _MLXClassifierWrapper(config, weights, num_labels)

Complete contract reference

Expand any definition for its exact inputs, annotations, defaults, return contract, directly raised exceptions, source-grounded behavior, and immutable line link. This section includes private and nested definitions that ordinary API generators omit.

vllm_mlx.rerank.RerankAdapter · class

vllm_mlx.rerank.RerankAdapter()

Per-family adapter for reranker models.

Parameters

This callable has no explicit inputs.

Returns

• Constructs: vllm_mlx.rerank.RerankAdapter

Exceptions and behavior

Class RerankAdapter derives from ABC and declares 3 direct member(s). No direct raise statement appears in this definition.

View source #L29-L80.

vllm_mlx.rerank.RerankAdapter.tokenize_pair · method

vllm_mlx.rerank.RerankAdapter.tokenize_pair(tokenizer, query: str, document: str, max_length: int) -> dict

Tokenize a (query, document) pair for the cross-encoder.

Parameters

Name	Type	Required	Default	Description
tokenizer	not annotated	yes	none	The HuggingFace tokenizer instance.
query	str	yes	none	The query string.
document	str	yes	none	The document string.
max_length	int	yes	none	Truncation length (from the model's context window).

Returns

• Type: dict

Exceptions and behavior

Method RerankAdapter.tokenize_pair contains no state mutation, call, raise, return, await, or yield. No direct raise statement appears in this definition.

View source #L39-L54.

vllm_mlx.rerank.RerankAdapter.extract_score · method

vllm_mlx.rerank.RerankAdapter.extract_score(logits) -> float

Extract a raw relevance score from model output logits.

Parameters

Name	Type	Required	Default	Description
logits	not annotated	yes	none	Model output logits (list or array), shape varies by model.

Returns

• Type: float

Exceptions and behavior

Method RerankAdapter.extract_score contains no state mutation, call, raise, return, await, or yield. No direct raise statement appears in this definition.

View source #L57-L67.

vllm_mlx.rerank.RerankAdapter.normalize · method

vllm_mlx.rerank.RerankAdapter.normalize(raw_score: float) -> float

Normalize a raw score to [0, 1] range.

Parameters

Name	Type	Required	Default	Description
raw_score	float	yes	none	The raw score from extract_score().

Returns

• Type: float

Exceptions and behavior

Method RerankAdapter.normalize contains no state mutation, call, raise, return, await, or yield. No direct raise statement appears in this definition.

View source #L70-L80.

vllm_mlx.rerank.SigmoidAdapter · class

vllm_mlx.rerank.SigmoidAdapter()

Default adapter for single-logit sigmoid rerankers.

Parameters

This callable has no explicit inputs.

Returns

• Constructs: vllm_mlx.rerank.SigmoidAdapter

Exceptions and behavior

Class SigmoidAdapter derives from RerankAdapter and declares 3 direct member(s). No direct raise statement appears in this definition.

View source #L83-L111.

vllm_mlx.rerank.SigmoidAdapter.tokenize_pair · method

vllm_mlx.rerank.SigmoidAdapter.tokenize_pair(tokenizer, query: str, document: str, max_length: int) -> dict

Tokenize as a sentence pair (query, document).

Parameters

Name	Type	Required	Default	Description
tokenizer	not annotated	yes	none	Required positional or keyword input.
query	str	yes	none	Required positional or keyword input.
document	str	yes	none	Required positional or keyword input.
max_length	int	yes	none	Required positional or keyword input.

Returns

• Type: dict
• Direct return expressions: tokenizer(query, document, padding=True, truncation=True, max_length=max_length, return_tensors='np')

Exceptions and behavior

Method SigmoidAdapter.tokenize_pair calls tokenizer; returns tokenizer(query, document, padding=True, truncation=True, max_length=max_length, return_tensors='np'). No direct raise statement appears in this definition.

View source #L92-L103.

vllm_mlx.rerank.SigmoidAdapter.extract_score · method

vllm_mlx.rerank.SigmoidAdapter.extract_score(logits) -> float

Extract the first logit as the relevance score.

Parameters

Name	Type	Required	Default	Description
logits	not annotated	yes	none	Required positional or keyword input.

Returns

• Type: float
• Direct return expressions: float(logits[0])

Exceptions and behavior

Method SigmoidAdapter.extract_score calls float; returns float(logits[0]). No direct raise statement appears in this definition.

View source #L105-L107.

vllm_mlx.rerank.SigmoidAdapter.normalize · method

vllm_mlx.rerank.SigmoidAdapter.normalize(raw_score: float) -> float

Apply sigmoid normalization.

Parameters

Name	Type	Required	Default	Description
raw_score	float	yes	none	Required positional or keyword input.

Returns

• Type: float
• Direct return expressions: 1.0 / (1.0 + math.exp(-raw_score))

Exceptions and behavior

Method SigmoidAdapter.normalize calls math.exp; returns 1.0 / (1.0 + math.exp(-raw_score)). No direct raise statement appears in this definition.

View source #L109-L111.

vllm_mlx.rerank.get_adapter · function

vllm_mlx.rerank.get_adapter(model_name: str) -> RerankAdapter

Return the appropriate adapter for a model.

Parameters

Name	Type	Required	Default	Description
model_name	str	yes	none	Required positional or keyword input.

Returns

• Type: RerankAdapter
• Direct return expressions: _ADAPTER_REGISTRY['default']()

Exceptions and behavior

Function get_adapter calls _ADAPTER_REGISTRY['default']; returns _ADAPTER_REGISTRY['default'](). No direct raise statement appears in this definition.

View source #L124-L133.

vllm_mlx.rerank.RerankEngine · class

vllm_mlx.rerank.RerankEngine(model_name: str, token_budget: int = 4096, max_concurrency: int = 1)

Reranker engine for cross-encoder sequence classification models.

Parameters

Name	Type	Required	Default	Description
model_name	str	yes	none	Required positional or keyword input.
token_budget	int	no	4096	Optional positional or keyword input; defaults to 4096.
max_concurrency	int	no	1	Optional positional or keyword input; defaults to 1.

Returns

• Constructs: vllm_mlx.rerank.RerankEngine

Exceptions and behavior

Class RerankEngine declares 6 direct member(s). No direct raise statement appears in this definition.

View source #L136-L338.

vllm_mlx.rerank.RerankEngine.__init__ · method

vllm_mlx.rerank.RerankEngine.__init__(model_name: str, token_budget: int = 4096, max_concurrency: int = 1) -> not annotated

Method RerankEngine.__init__ updates self.model_name, self.token_budget, self.max_concurrency, self._semaphore; calls asyncio.Semaphore.

Parameters

Name	Type	Required	Default	Description
model_name	str	yes	none	Required positional or keyword input.
token_budget	int	no	4096	Optional positional or keyword input; defaults to 4096.
max_concurrency	int	no	1	Optional positional or keyword input; defaults to 1.

Returns

• Type: not annotated

Exceptions and behavior

Method RerankEngine.__init__ updates self.model_name, self.token_budget, self.max_concurrency, self._semaphore; calls asyncio.Semaphore. No direct raise statement appears in this definition.

View source #L147-L159.

vllm_mlx.rerank.RerankEngine.is_loaded · method

vllm_mlx.rerank.RerankEngine.is_loaded() -> bool

Return whether the reranking model has been loaded.

Parameters

This callable has no explicit inputs.

Returns

• Type: bool
• Direct return expressions: self._model is not None

Exceptions and behavior

Method RerankEngine.is_loaded returns self._model is not None. No direct raise statement appears in this definition.

View source #L162-L165.

vllm_mlx.rerank.RerankEngine.load · method

vllm_mlx.rerank.RerankEngine.load() -> None

Load the cross-encoder model and tokenizer.

Parameters

This callable has no explicit inputs.

Returns

• Type: None

Exceptions and behavior

Method RerankEngine.load updates self._tokenizer, self._model, self._adapter; calls logger.info, time.perf_counter, AutoTokenizer.from_pretrained, self._load_mlx_model. No direct raise statement appears in this definition.

View source #L167-L184.

vllm_mlx.rerank.RerankEngine._load_mlx_model · method

vllm_mlx.rerank.RerankEngine._load_mlx_model(model_name: str) -> not annotated

Load an MLX cross-encoder model from HuggingFace Hub.

Parameters

Name	Type	Required	Default	Description
model_name	str	yes	none	Required positional or keyword input.

Returns

• Type: not annotated
• Direct return expressions: model

Exceptions and behavior

Method RerankEngine._load_mlx_model calls snapshot_download, os.path.join, open, json.load; can raise FileNotFoundError; returns model. Directly raised exceptions: FileNotFoundError.

View source #L187-L230.

vllm_mlx.rerank.RerankEngine._ensure_loaded · method

vllm_mlx.rerank.RerankEngine._ensure_loaded() -> None

Method RerankEngine._ensure_loaded calls self.load.

Parameters

This callable has no explicit inputs.

Returns

• Type: None

Exceptions and behavior

Method RerankEngine._ensure_loaded calls self.load. No direct raise statement appears in this definition.

View source #L232-L234.

vllm_mlx.rerank.RerankEngine.score_pairs · method

vllm_mlx.rerank.RerankEngine.score_pairs(query: str, documents: list[str]) -> tuple[list[float], int]

Score each (query, document) pair and return normalized relevance scores.

Parameters

Name	Type	Required	Default	Description
query	str	yes	none	The query string.
documents	list[str]	yes	none	List of document strings.

Returns

• Type: tuple[list[float], int]
• Direct return expressions: ([score for _, score in all_scores], total_tokens)

Exceptions and behavior

Method RerankEngine.score_pairs calls self._ensure_loaded, resolve_max_length, getattr, self._adapter.tokenize_pair; returns ([score for _, score in all_scores], total_tokens). No direct raise statement appears in this definition.

View source #L236-L338.

vllm_mlx.rerank._build_classifier_model · function

vllm_mlx.rerank._build_classifier_model(model_type, config, weights, num_labels) -> not annotated

Build an MLX sequence classification model from config and weights.

Parameters

Name	Type	Required	Default	Description
model_type	not annotated	yes	none	Required positional or keyword input.
config	not annotated	yes	none	Required positional or keyword input.
weights	not annotated	yes	none	Required positional or keyword input.
num_labels	not annotated	yes	none	Required positional or keyword input.

Returns

• Type: not annotated
• Direct return expressions: _MLXClassifierWrapper(config, weights, num_labels)

Exceptions and behavior

Function _build_classifier_model calls _MLXClassifierWrapper; returns _MLXClassifierWrapper(config, weights, num_labels). No direct raise statement appears in this definition.

View source #L341-L349.

vllm_mlx.rerank._MLXClassifierWrapper · class

vllm_mlx.rerank._MLXClassifierWrapper(config: dict, weights: dict, num_labels: int)

Minimal MLX wrapper for sequence classification models.

Parameters

Name	Type	Required	Default	Description
config	dict	yes	none	Required positional or keyword input.
weights	dict	yes	none	Required positional or keyword input.
num_labels	int	yes	none	Required positional or keyword input.

Returns

• Constructs: vllm_mlx.rerank._MLXClassifierWrapper

Exceptions and behavior

Class _MLXClassifierWrapper declares 3 direct member(s). No direct raise statement appears in this definition.

View source #L352-L391.

vllm_mlx.rerank._MLXClassifierWrapper.__init__ · method

vllm_mlx.rerank._MLXClassifierWrapper.__init__(config: dict, weights: dict, num_labels: int) -> not annotated

Method _MLXClassifierWrapper.__init__ updates self.config, self.weights, self.num_labels, self._params; calls list, weights.values.

Parameters

Name	Type	Required	Default	Description
config	dict	yes	none	Required positional or keyword input.
weights	dict	yes	none	Required positional or keyword input.
num_labels	int	yes	none	Required positional or keyword input.

Returns

• Type: not annotated

Exceptions and behavior

Method _MLXClassifierWrapper.__init__ updates self.config, self.weights, self.num_labels, self._params; calls list, weights.values. No direct raise statement appears in this definition.

View source #L361-L365.

vllm_mlx.rerank._MLXClassifierWrapper.parameters · method

vllm_mlx.rerank._MLXClassifierWrapper.parameters() -> not annotated

Return model parameters for mx.eval.

Parameters

This callable has no explicit inputs.

Returns

• Type: not annotated
• Direct return expressions: self._params

Exceptions and behavior

Method _MLXClassifierWrapper.parameters returns self._params. No direct raise statement appears in this definition.

View source #L367-L369.

vllm_mlx.rerank._MLXClassifierWrapper.__call__ · method

vllm_mlx.rerank._MLXClassifierWrapper.__call__(input_ids: mx.array, attention_mask: mx.array = None) -> not annotated

Forward pass through the classifier.

Parameters

Name	Type	Required	Default	Description
input_ids	mx.array	yes	none	Required positional or keyword input.
attention_mask	mx.array	no	None	Optional positional or keyword input; defaults to None.

Returns

• Type: not annotated
• Direct return expressions: _ClassifierOutput(logits=logits)

Exceptions and behavior

Method _MLXClassifierWrapper.__call__ calls classifier_forward, _ClassifierOutput; returns _ClassifierOutput(logits=logits). No direct raise statement appears in this definition.

View source #L371-L391.

vllm_mlx.rerank._ClassifierOutput · class

vllm_mlx.rerank._ClassifierOutput(logits: mx.array)

Simple container for classifier output logits.

Parameters

Name	Type	Required	Default	Description
logits	mx.array	yes	none	Required positional or keyword input.

Returns

• Constructs: vllm_mlx.rerank._ClassifierOutput

Exceptions and behavior

Class _ClassifierOutput declares 1 direct member(s). No direct raise statement appears in this definition.

View source #L394-L398.

vllm_mlx.rerank._ClassifierOutput.__init__ · method

vllm_mlx.rerank._ClassifierOutput.__init__(logits: mx.array) -> not annotated

Method _ClassifierOutput.__init__ updates self.logits.

Parameters

Name	Type	Required	Default	Description
logits	mx.array	yes	none	Required positional or keyword input.

Returns

• Type: not annotated

Exceptions and behavior

Method _ClassifierOutput.__init__ updates self.logits. No direct raise statement appears in this definition.

View source #L397-L398.

Complete symbol map

This map also includes private definitions and nested helpers. The signature column exposes every explicit input even when an internal helper has no dedicated parameter prose.

Symbol	Kind	Signature and inputs	What it does	Source
RerankAdapter	class	RerankAdapter()	Per-family adapter for reranker models.	#L29-L80
RerankAdapter.tokenize_pair	method	RerankAdapter.tokenize_pair(tokenizer, query: str, document: str, max_length: int) -> dict	Tokenize a (query, document) pair for the cross-encoder.	#L39-L54
RerankAdapter.extract_score	method	RerankAdapter.extract_score(logits) -> float	Extract a raw relevance score from model output logits.	#L57-L67
RerankAdapter.normalize	method	RerankAdapter.normalize(raw_score: float) -> float	Normalize a raw score to [0, 1] range.	#L70-L80
SigmoidAdapter	class	SigmoidAdapter()	Default adapter for single-logit sigmoid rerankers.	#L83-L111
SigmoidAdapter.tokenize_pair	method	SigmoidAdapter.tokenize_pair(tokenizer, query: str, document: str, max_length: int) -> dict	Tokenize as a sentence pair (query, document).	#L92-L103
SigmoidAdapter.extract_score	method	SigmoidAdapter.extract_score(logits) -> float	Extract the first logit as the relevance score.	#L105-L107
SigmoidAdapter.normalize	method	SigmoidAdapter.normalize(raw_score: float) -> float	Apply sigmoid normalization.	#L109-L111
get_adapter	function	get_adapter(model_name: str) -> RerankAdapter	Return the appropriate adapter for a model.	#L124-L133
RerankEngine	class	RerankEngine(model_name: str, token_budget: int = 4096, max_concurrency: int = 1)	Reranker engine for cross-encoder sequence classification models.	#L136-L338
RerankEngine.__init__	method	RerankEngine.__init__(model_name: str, token_budget: int = 4096, max_concurrency: int = 1) -> not annotated	Method RerankEngine.__init__ updates self.model_name, self.token_budget, self.max_concurrency, self._semaphore; calls asyncio.Semaphore.	#L147-L159
RerankEngine.is_loaded	method	RerankEngine.is_loaded() -> bool	Return whether the reranking model has been loaded.	#L162-L165
RerankEngine.load	method	RerankEngine.load() -> None	Load the cross-encoder model and tokenizer.	#L167-L184
RerankEngine._load_mlx_model	method	RerankEngine._load_mlx_model(model_name: str) -> not annotated	Load an MLX cross-encoder model from HuggingFace Hub.	#L187-L230
RerankEngine._ensure_loaded	method	RerankEngine._ensure_loaded() -> None	Method RerankEngine._ensure_loaded calls self.load.	#L232-L234
RerankEngine.score_pairs	method	RerankEngine.score_pairs(query: str, documents: list[str]) -> tuple[list[float], int]	Score each (query, document) pair and return normalized relevance scores.	#L236-L338
_build_classifier_model	function	_build_classifier_model(model_type, config, weights, num_labels) -> not annotated	Build an MLX sequence classification model from config and weights.	#L341-L349
_MLXClassifierWrapper	class	_MLXClassifierWrapper(config: dict, weights: dict, num_labels: int)	Minimal MLX wrapper for sequence classification models.	#L352-L391
_MLXClassifierWrapper.__init__	method	_MLXClassifierWrapper.__init__(config: dict, weights: dict, num_labels: int) -> not annotated	Method _MLXClassifierWrapper.__init__ updates self.config, self.weights, self.num_labels, self._params; calls list, weights.values.	#L361-L365
_MLXClassifierWrapper.parameters	method	_MLXClassifierWrapper.parameters() -> not annotated	Return model parameters for mx.eval.	#L367-L369
_MLXClassifierWrapper.__call__	method	_MLXClassifierWrapper.__call__(input_ids: mx.array, attention_mask: mx.array = None) -> not annotated	Forward pass through the classifier.	#L371-L391
_ClassifierOutput	class	_ClassifierOutput(logits: mx.array)	Simple container for classifier output logits.	#L394-L398
_ClassifierOutput.__init__	method	_ClassifierOutput.__init__(logits: mx.array) -> not annotated	Method _ClassifierOutput.__init__ updates self.logits.	#L397-L398