| from typing import Any |
|
|
| import mlx.core as mx |
| from mlx_lm.models.qwen3 import ModelArgs, Qwen3Model |
|
|
| SUPPORTED_MODEL_TYPE = "bidirectional_pplx_qwen3" |
|
|
|
|
| def make_bidirectional_padding_mask(attention_mask: mx.array | None) -> mx.array | None: |
| """Create a full-attention key padding mask from a tokenizer attention mask.""" |
| if attention_mask is None: |
| return None |
|
|
| return attention_mask.astype(mx.bool_)[:, None, None, :] |
|
|
|
|
| class PPLXQwen3Model(Qwen3Model): |
| """Qwen3 encoder variant used by Perplexity contextual embedding models.""" |
|
|
| def __call__( |
| self, |
| inputs: mx.array, |
| attention_mask: mx.array | None = None, |
| input_embeddings: mx.array | None = None, |
| ) -> mx.array: |
| if input_embeddings is not None: |
| hidden_states = input_embeddings |
| else: |
| hidden_states = self.embed_tokens(inputs) |
|
|
| mask = make_bidirectional_padding_mask(attention_mask) |
|
|
| for layer in self.layers: |
| hidden_states = layer(hidden_states, mask, None) |
|
|
| return self.norm(hidden_states) |
|
|
|
|
| def get_pplx_model_classes(config: dict[str, Any]) -> tuple[type[PPLXQwen3Model], type[ModelArgs]]: |
| model_type = config.get("model_type") |
| if model_type != SUPPORTED_MODEL_TYPE: |
| msg = ( |
| f"Unsupported Perplexity model type {model_type!r}; expected {SUPPORTED_MODEL_TYPE!r}." |
| ) |
| raise ValueError(msg) |
|
|
| return PPLXQwen3Model, ModelArgs |
|
|
