Add files using upload-large-folder tool

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,7 @@
+---
+language: en
+pipeline_tag: text-generation
+tags:
+- mlx
+library_name: mlx
+---

config.json

@@ -0,0 +1,79 @@
+{
+    "architectures": [
+        "PPLXQwen3Model"
+    ],
+    "attention_bias": false,
+    "attention_dropout": 0.0,
+    "attn_implementation": "sdpa",
+    "auto_map": {
+        "AutoConfig": "configuration.PPLXQwen3Config",
+        "AutoModel": "modeling.PPLXQwen3Model"
+    },
+    "bos_token_id": 151643,
+    "dtype": "float32",
+    "eos_token_id": 151643,
+    "head_dim": 128,
+    "hidden_act": "silu",
+    "hidden_size": 2560,
+    "initializer_range": 0.02,
+    "intermediate_size": 9728,
+    "layer_types": [
+        "full_attention",
+        "full_attention",
+        "full_attention",
+        "full_attention",
+        "full_attention",
+        "full_attention",
+        "full_attention",
+        "full_attention",
+        "full_attention",
+        "full_attention",
+        "full_attention",
+        "full_attention",
+        "full_attention",
+        "full_attention",
+        "full_attention",
+        "full_attention",
+        "full_attention",
+        "full_attention",
+        "full_attention",
+        "full_attention",
+        "full_attention",
+        "full_attention",
+        "full_attention",
+        "full_attention",
+        "full_attention",
+        "full_attention",
+        "full_attention",
+        "full_attention",
+        "full_attention",
+        "full_attention",
+        "full_attention",
+        "full_attention",
+        "full_attention",
+        "full_attention",
+        "full_attention",
+        "full_attention"
+    ],
+    "max_position_embeddings": 32768,
+    "max_window_layers": 36,
+    "model_file": "mlx_pplx_qwen3.py",
+    "model_type": "bidirectional_pplx_qwen3",
+    "num_attention_heads": 32,
+    "num_hidden_layers": 36,
+    "num_key_value_heads": 8,
+    "rms_norm_eps": 1e-06,
+    "rope_parameters": {
+        "rope_theta": 1000000,
+        "rope_type": "default"
+    },
+    "rope_theta": 1000000,
+    "sliding_window": null,
+    "source_model_type": "bidirectional_pplx_qwen3",
+    "tie_word_embeddings": true,
+    "transformers_version": "5.0.0.dev0",
+    "use_bidirectional_attention": true,
+    "use_cache": false,
+    "use_sliding_window": false,
+    "vocab_size": 151936
+}

configuration.py

@@ -0,0 +1,5 @@
+from transformers.models.qwen3.configuration_qwen3 import Qwen3Config
+
+
+class PPLXQwen3Config(Qwen3Config):
+    model_type = "bidirectional_pplx_qwen3"

mlx_pplx_qwen3.py

@@ -0,0 +1,230 @@
+from dataclasses import dataclass
+from typing import Any, Dict, Optional, Union
+
+import mlx.core as mx
+import mlx.nn as nn
+from mlx.nn.layers.distributed import shard_linear
+
+from mlx_lm.models.activations import swiglu
+from mlx_lm.models.base import (
+    BaseModelArgs,
+    scaled_dot_product_attention,
+)
+from mlx_lm.models.rope_utils import initialize_rope
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    hidden_size: int
+    num_hidden_layers: int
+    intermediate_size: int
+    num_attention_heads: int
+    rms_norm_eps: float
+    vocab_size: int
+    num_key_value_heads: int
+    max_position_embeddings: int
+    rope_theta: float
+    head_dim: int
+    tie_word_embeddings: bool
+    rope_scaling: Optional[Dict[str, Union[float, str]]] = None
+
+
+def _make_bidirectional_mask(
+    attention_mask: mx.array,
+    batch_size: int,
+    seq_len: int,
+    offset: int = 0,
+) -> mx.array:
+    if attention_mask.ndim != 2:
+        raise ValueError(
+            f"Expected 2D attention_mask with shape [batch, seq], got {attention_mask.shape}"
+        )
+    if attention_mask.shape[0] != batch_size:
+        raise ValueError(
+            "attention_mask batch size does not match input batch size"
+        )
+    if attention_mask.shape[1] < offset + seq_len:
+        raise ValueError(
+            "attention_mask sequence length is shorter than the required cached length"
+        )
+
+    # Build a full (non-causal) valid-token mask from the 2D attention mask.
+    q = attention_mask[:, offset : offset + seq_len].astype(mx.bool_)
+    k = attention_mask[:, : offset + seq_len].astype(mx.bool_)
+    mask = q[:, :, None] & k[:, None, :]
+    return mask[:, None, :, :]
+
+
+class Attention(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+
+        dim = args.hidden_size
+        self.n_heads = n_heads = args.num_attention_heads
+        self.n_kv_heads = n_kv_heads = args.num_key_value_heads
+        head_dim = args.head_dim
+        self.scale = head_dim**-0.5
+
+        self.q_proj = nn.Linear(dim, n_heads * head_dim, bias=False)
+        self.k_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=False)
+        self.v_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=False)
+        self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=False)
+
+        self.q_norm = nn.RMSNorm(head_dim, eps=args.rms_norm_eps)
+        self.k_norm = nn.RMSNorm(head_dim, eps=args.rms_norm_eps)
+        self.rope = initialize_rope(
+            head_dim,
+            base=args.rope_theta,
+            traditional=False,
+            scaling_config=args.rope_scaling,
+            max_position_embeddings=args.max_position_embeddings,
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        bsz, seq_len, _ = x.shape
+
+        queries = self.q_proj(x)
+        keys = self.k_proj(x)
+        values = self.v_proj(x)
+
+        queries = self.q_norm(queries.reshape(bsz, seq_len, self.n_heads, -1)).transpose(
+            0, 2, 1, 3
+        )
+        keys = self.k_norm(keys.reshape(bsz, seq_len, self.n_kv_heads, -1)).transpose(
+            0, 2, 1, 3
+        )
+        values = values.reshape(bsz, seq_len, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
+
+        if cache is not None:
+            queries = self.rope(queries, offset=cache.offset)
+            keys = self.rope(keys, offset=cache.offset)
+            keys, values = cache.update_and_fetch(keys, values)
+        else:
+            queries = self.rope(queries)
+            keys = self.rope(keys)
+
+        output = scaled_dot_product_attention(
+            queries,
+            keys,
+            values,
+            cache=cache,
+            scale=self.scale,
+            mask=mask,
+        )
+        output = output.transpose(0, 2, 1, 3).reshape(bsz, seq_len, -1)
+        return self.o_proj(output)
+
+
+class MLP(nn.Module):
+    def __init__(self, dim: int, hidden_dim: int):
+        super().__init__()
+        self.gate_proj = nn.Linear(dim, hidden_dim, bias=False)
+        self.down_proj = nn.Linear(hidden_dim, dim, bias=False)
+        self.up_proj = nn.Linear(dim, hidden_dim, bias=False)
+
+    def __call__(self, x: mx.array) -> mx.array:
+        return self.down_proj(swiglu(self.gate_proj(x), self.up_proj(x)))
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.self_attn = Attention(args)
+        self.mlp = MLP(args.hidden_size, args.intermediate_size)
+        self.input_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+        self.post_attention_layernorm = nn.RMSNorm(
+            args.hidden_size, eps=args.rms_norm_eps
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        residual = x + self.self_attn(self.input_layernorm(x), mask, cache)
+        return residual + self.mlp(self.post_attention_layernorm(residual))
+
+
+class Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.model_type = args.model_type
+        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
+        self.layers = [TransformerBlock(args) for _ in range(args.num_hidden_layers)]
+        self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache=None,
+        input_embeddings: Optional[mx.array] = None,
+        attention_mask: Optional[mx.array] = None,
+    ) -> mx.array:
+        if input_embeddings is not None:
+            h = input_embeddings
+        else:
+            h = self.embed_tokens(inputs)
+
+        if cache is None:
+            cache = [None] * len(self.layers)
+        elif len(cache) != len(self.layers):
+            raise ValueError(
+                f"Expected cache with {len(self.layers)} layers, got {len(cache)}"
+            )
+
+        if any(layer_cache is not None for layer_cache in cache):
+            raise ValueError(
+                "KV cache is not supported for this bidirectional embedding model."
+            )
+
+        if attention_mask is not None:
+            mask = _make_bidirectional_mask(
+                attention_mask,
+                batch_size=h.shape[0],
+                seq_len=h.shape[1],
+                offset=0,
+            )
+        else:
+            mask = None
+
+        for layer, layer_cache in zip(self.layers, cache):
+            h = layer(h, mask, layer_cache)
+
+        return self.norm(h)
+
+    def shard(self, group: Optional[mx.distributed.Group] = None):
+        group = group or mx.distributed.init()
+        n = group.size()
+        for layer in self.layers:
+            layer.self_attn.q_proj = shard_linear(
+                layer.self_attn.q_proj, "all-to-sharded", group=group
+            )
+            layer.self_attn.k_proj = shard_linear(
+                layer.self_attn.k_proj, "all-to-sharded", group=group
+            )
+            layer.self_attn.v_proj = shard_linear(
+                layer.self_attn.v_proj, "all-to-sharded", group=group
+            )
+            layer.self_attn.o_proj = shard_linear(
+                layer.self_attn.o_proj, "sharded-to-all", group=group
+            )
+            layer.self_attn.n_heads //= n
+            layer.self_attn.n_kv_heads //= n
+
+            layer.mlp.gate_proj = shard_linear(
+                layer.mlp.gate_proj, "all-to-sharded", group=group
+            )
+            layer.mlp.down_proj = shard_linear(
+                layer.mlp.down_proj, "sharded-to-all", group=group
+            )
+            layer.mlp.up_proj = shard_linear(
+                layer.mlp.up_proj, "all-to-sharded", group=group
+            )

model-00001-of-00002.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8f384893878795ee18bd9a351ef2f1173b1949a013d6c64afab71b353c564336
+size 5321131845

model-00002-of-00002.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1e32d5e4c3d725972a8f4fc6938668d44c4dd096040483783b8e0efcb2cf5c4e
+size 2723847252

model.safetensors.index.json

@@ -0,0 +1,406 @@
+{
+    "metadata": {
+        "total_size": 8044936192,
+        "total_parameters": 4022468096
+    },
+    "weight_map": {
+        "embed_tokens.weight": "model-00001-of-00002.safetensors",
+        "layers.0.input_layernorm.weight": "model-00001-of-00002.safetensors",
+        "layers.0.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.0.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.0.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.0.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+        "layers.0.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
+        "layers.0.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.0.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.0.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
+        "layers.0.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.0.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.1.input_layernorm.weight": "model-00001-of-00002.safetensors",
+        "layers.1.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.1.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.1.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.1.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+        "layers.1.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
+        "layers.1.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.1.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.1.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
+        "layers.1.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.1.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.10.input_layernorm.weight": "model-00001-of-00002.safetensors",
+        "layers.10.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.10.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.10.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.10.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+        "layers.10.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
+        "layers.10.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.10.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.10.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
+        "layers.10.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.10.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.11.input_layernorm.weight": "model-00001-of-00002.safetensors",
+        "layers.11.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.11.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.11.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.11.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+        "layers.11.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
+        "layers.11.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.11.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.11.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
+        "layers.11.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.11.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.12.input_layernorm.weight": "model-00001-of-00002.safetensors",
+        "layers.12.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.12.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.12.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.12.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+        "layers.12.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
+        "layers.12.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.12.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.12.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
+        "layers.12.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.12.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.13.input_layernorm.weight": "model-00001-of-00002.safetensors",
+        "layers.13.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.13.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.13.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.13.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+        "layers.13.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
+        "layers.13.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.13.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.13.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
+        "layers.13.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.13.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.14.input_layernorm.weight": "model-00001-of-00002.safetensors",
+        "layers.14.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.14.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.14.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.14.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+        "layers.14.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
+        "layers.14.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.14.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.14.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
+        "layers.14.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.14.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.15.input_layernorm.weight": "model-00001-of-00002.safetensors",
+        "layers.15.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.15.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.15.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.15.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+        "layers.15.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
+        "layers.15.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.15.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.15.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
+        "layers.15.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.15.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.16.input_layernorm.weight": "model-00001-of-00002.safetensors",
+        "layers.16.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.16.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.16.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.16.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+        "layers.16.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
+        "layers.16.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.16.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.16.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
+        "layers.16.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.16.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.17.input_layernorm.weight": "model-00001-of-00002.safetensors",
+        "layers.17.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.17.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.17.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.17.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+        "layers.17.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
+        "layers.17.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.17.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.17.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
+        "layers.17.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.17.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.18.input_layernorm.weight": "model-00001-of-00002.safetensors",
+        "layers.18.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.18.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.18.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.18.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+        "layers.18.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
+        "layers.18.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.18.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.18.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
+        "layers.18.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.18.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.19.input_layernorm.weight": "model-00001-of-00002.safetensors",
+        "layers.19.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.19.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.19.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.19.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+        "layers.19.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
+        "layers.19.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.19.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.19.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
+        "layers.19.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.19.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.2.input_layernorm.weight": "model-00001-of-00002.safetensors",
+        "layers.2.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.2.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.2.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.2.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+        "layers.2.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
+        "layers.2.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.2.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.2.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
+        "layers.2.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.2.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.20.input_layernorm.weight": "model-00001-of-00002.safetensors",
+        "layers.20.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.20.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.20.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.20.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+        "layers.20.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
+        "layers.20.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.20.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.20.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
+        "layers.20.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.20.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.21.input_layernorm.weight": "model-00001-of-00002.safetensors",
+        "layers.21.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.21.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.21.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.21.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+        "layers.21.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
+        "layers.21.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.21.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.21.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
+        "layers.21.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.21.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.22.input_layernorm.weight": "model-00002-of-00002.safetensors",
+        "layers.22.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.22.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.22.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.22.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
+        "layers.22.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
+        "layers.22.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.22.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.22.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
+        "layers.22.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.22.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.23.input_layernorm.weight": "model-00002-of-00002.safetensors",
+        "layers.23.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.23.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.23.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.23.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
+        "layers.23.self_attn.k_norm.weight": "model-00002-of-00002.safetensors",
+        "layers.23.self_attn.k_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.23.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.23.self_attn.q_norm.weight": "model-00002-of-00002.safetensors",
+        "layers.23.self_attn.q_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.23.self_attn.v_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.24.input_layernorm.weight": "model-00002-of-00002.safetensors",
+        "layers.24.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.24.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.24.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.24.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
+        "layers.24.self_attn.k_norm.weight": "model-00002-of-00002.safetensors",
+        "layers.24.self_attn.k_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.24.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.24.self_attn.q_norm.weight": "model-00002-of-00002.safetensors",
+        "layers.24.self_attn.q_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.24.self_attn.v_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.25.input_layernorm.weight": "model-00002-of-00002.safetensors",
+        "layers.25.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.25.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.25.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.25.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
+        "layers.25.self_attn.k_norm.weight": "model-00002-of-00002.safetensors",
+        "layers.25.self_attn.k_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.25.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.25.self_attn.q_norm.weight": "model-00002-of-00002.safetensors",
+        "layers.25.self_attn.q_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.25.self_attn.v_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.26.input_layernorm.weight": "model-00002-of-00002.safetensors",
+        "layers.26.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.26.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.26.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.26.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
+        "layers.26.self_attn.k_norm.weight": "model-00002-of-00002.safetensors",
+        "layers.26.self_attn.k_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.26.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.26.self_attn.q_norm.weight": "model-00002-of-00002.safetensors",
+        "layers.26.self_attn.q_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.26.self_attn.v_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.27.input_layernorm.weight": "model-00002-of-00002.safetensors",
+        "layers.27.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.27.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.27.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.27.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
+        "layers.27.self_attn.k_norm.weight": "model-00002-of-00002.safetensors",
+        "layers.27.self_attn.k_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.27.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.27.self_attn.q_norm.weight": "model-00002-of-00002.safetensors",
+        "layers.27.self_attn.q_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.27.self_attn.v_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.28.input_layernorm.weight": "model-00002-of-00002.safetensors",
+        "layers.28.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.28.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.28.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.28.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
+        "layers.28.self_attn.k_norm.weight": "model-00002-of-00002.safetensors",
+        "layers.28.self_attn.k_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.28.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.28.self_attn.q_norm.weight": "model-00002-of-00002.safetensors",
+        "layers.28.self_attn.q_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.28.self_attn.v_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.29.input_layernorm.weight": "model-00002-of-00002.safetensors",
+        "layers.29.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.29.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.29.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.29.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
+        "layers.29.self_attn.k_norm.weight": "model-00002-of-00002.safetensors",
+        "layers.29.self_attn.k_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.29.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.29.self_attn.q_norm.weight": "model-00002-of-00002.safetensors",
+        "layers.29.self_attn.q_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.29.self_attn.v_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.3.input_layernorm.weight": "model-00001-of-00002.safetensors",
+        "layers.3.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.3.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.3.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.3.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+        "layers.3.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
+        "layers.3.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.3.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.3.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
+        "layers.3.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.3.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.30.input_layernorm.weight": "model-00002-of-00002.safetensors",
+        "layers.30.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.30.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.30.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.30.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
+        "layers.30.self_attn.k_norm.weight": "model-00002-of-00002.safetensors",
+        "layers.30.self_attn.k_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.30.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.30.self_attn.q_norm.weight": "model-00002-of-00002.safetensors",
+        "layers.30.self_attn.q_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.30.self_attn.v_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.31.input_layernorm.weight": "model-00002-of-00002.safetensors",
+        "layers.31.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.31.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.31.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.31.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
+        "layers.31.self_attn.k_norm.weight": "model-00002-of-00002.safetensors",
+        "layers.31.self_attn.k_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.31.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.31.self_attn.q_norm.weight": "model-00002-of-00002.safetensors",
+        "layers.31.self_attn.q_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.31.self_attn.v_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.32.input_layernorm.weight": "model-00002-of-00002.safetensors",
+        "layers.32.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.32.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.32.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.32.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
+        "layers.32.self_attn.k_norm.weight": "model-00002-of-00002.safetensors",
+        "layers.32.self_attn.k_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.32.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.32.self_attn.q_norm.weight": "model-00002-of-00002.safetensors",
+        "layers.32.self_attn.q_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.32.self_attn.v_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.33.input_layernorm.weight": "model-00002-of-00002.safetensors",
+        "layers.33.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.33.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.33.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.33.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
+        "layers.33.self_attn.k_norm.weight": "model-00002-of-00002.safetensors",
+        "layers.33.self_attn.k_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.33.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.33.self_attn.q_norm.weight": "model-00002-of-00002.safetensors",
+        "layers.33.self_attn.q_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.33.self_attn.v_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.34.input_layernorm.weight": "model-00002-of-00002.safetensors",
+        "layers.34.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.34.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.34.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.34.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
+        "layers.34.self_attn.k_norm.weight": "model-00002-of-00002.safetensors",
+        "layers.34.self_attn.k_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.34.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.34.self_attn.q_norm.weight": "model-00002-of-00002.safetensors",
+        "layers.34.self_attn.q_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.34.self_attn.v_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.35.input_layernorm.weight": "model-00002-of-00002.safetensors",
+        "layers.35.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.35.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.35.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.35.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
+        "layers.35.self_attn.k_norm.weight": "model-00002-of-00002.safetensors",
+        "layers.35.self_attn.k_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.35.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.35.self_attn.q_norm.weight": "model-00002-of-00002.safetensors",
+        "layers.35.self_attn.q_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.35.self_attn.v_proj.weight": "model-00002-of-00002.safetensors",
+        "layers.4.input_layernorm.weight": "model-00001-of-00002.safetensors",
+        "layers.4.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.4.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.4.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.4.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+        "layers.4.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
+        "layers.4.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.4.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.4.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
+        "layers.4.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.4.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.5.input_layernorm.weight": "model-00001-of-00002.safetensors",
+        "layers.5.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.5.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.5.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.5.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+        "layers.5.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
+        "layers.5.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.5.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.5.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
+        "layers.5.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.5.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.6.input_layernorm.weight": "model-00001-of-00002.safetensors",
+        "layers.6.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.6.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.6.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.6.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+        "layers.6.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
+        "layers.6.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.6.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.6.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
+        "layers.6.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.6.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.7.input_layernorm.weight": "model-00001-of-00002.safetensors",
+        "layers.7.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.7.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.7.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.7.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+        "layers.7.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
+        "layers.7.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.7.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.7.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
+        "layers.7.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.7.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.8.input_layernorm.weight": "model-00001-of-00002.safetensors",
+        "layers.8.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.8.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.8.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.8.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+        "layers.8.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
+        "layers.8.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.8.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.8.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
+        "layers.8.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.8.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.9.input_layernorm.weight": "model-00001-of-00002.safetensors",
+        "layers.9.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.9.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.9.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.9.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+        "layers.9.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
+        "layers.9.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.9.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.9.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
+        "layers.9.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+        "layers.9.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+        "norm.weight": "model-00002-of-00002.safetensors"
+    }
+}

modeling.py

@@ -0,0 +1,83 @@
+from typing import Callable
+import torch
+from transformers import Qwen3Model
+from transformers.cache_utils import Cache
+from transformers.masking_utils import create_causal_mask
+from transformers.modeling_outputs import BaseModelOutputWithPooling
+from transformers.processing_utils import Unpack
+from transformers.utils import TransformersKwargs
+from .configuration import PPLXQwen3Config
+
+
+# From modeling_t5gemma.py
+def bidirectional_mask_function(attention_mask: torch.Tensor | None) -> Callable:
+    """
+    This creates bidirectional attention mask.
+    """
+
+    def inner_mask(batch_idx: int, head_idx: int, q_idx: int, kv_idx: int) -> bool:
+        if attention_mask is None:
+            return torch.ones((), dtype=torch.bool)
+        return attention_mask[batch_idx, kv_idx].to(torch.bool)
+
+    return inner_mask
+
+
+class PPLXQwen3Model(Qwen3Model):
+    _supports_flash_attn = True
+    _supports_sdpa = True
+
+    config_class = PPLXQwen3Config
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.post_init()
+
+    def post_init(self):
+        super().post_init()
+        # Override to set all layers to non-causal attention. This'll work with attn_implementation="flash_attention_2" or "sdpa"
+        for layer in self.layers:
+            layer.self_attn.is_causal = False
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        use_cache: bool | None = None,
+        cache_position: torch.LongTensor | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> BaseModelOutputWithPooling:
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+            input_ids = None
+
+        # We construct a dummy tensor imitating initial positions
+        dummy_cache_position = torch.arange(
+            inputs_embeds.shape[1], device=inputs_embeds.device, dtype=torch.long
+        )
+        attention_mask = {
+            "full_attention": create_causal_mask(
+                config=self.config,
+                input_embeds=inputs_embeds,
+                attention_mask=attention_mask,
+                cache_position=dummy_cache_position,
+                past_key_values=None,
+                position_ids=position_ids,
+                or_mask_function=bidirectional_mask_function(attention_mask),
+            )
+        }
+
+        outputs = super().forward(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            **kwargs,
+        )
+        return outputs

st_quantize.py

@@ -0,0 +1,122 @@
+import torch
+import numpy as np
+from typing import Literal
+from sentence_transformers.models import Module
+
+
+class Quantizer(torch.nn.Module):
+    def __init__(self, hard: bool = True):
+        """
+        Args:
+            hard: Whether to use hard or soft quantization. Defaults to True.
+        """
+        super().__init__()
+        self._hard = hard
+
+    def _hard_quantize(self, x, *args, **kwargs) -> torch.Tensor:
+        raise NotImplementedError
+
+    def _soft_quantize(self, x, *args, **kwargs) -> torch.Tensor:
+        raise NotImplementedError
+
+    def forward(self, x, *args, **kwargs) -> torch.Tensor:
+        soft = self._soft_quantize(x, *args, **kwargs)
+
+        if not self._hard:
+            result = soft
+        else:
+            result = (
+                self._hard_quantize(x, *args, **kwargs).detach() + soft - soft.detach()
+            )
+
+        return result
+
+
+class Int8TanhQuantizer(Quantizer):
+    def __init__(
+        self,
+        hard: bool = True,
+    ):
+        super().__init__(hard=hard)
+        self.qmin = -128
+        self.qmax = 127
+
+    def _soft_quantize(self, x, *args, **kwargs):
+        return torch.tanh(x)
+
+    def _hard_quantize(self, x, *args, **kwargs):
+        soft = self._soft_quantize(x)
+        int_x = torch.round(soft * self.qmax)
+        int_x = torch.clamp(int_x, self.qmin, self.qmax)
+        return int_x
+
+
+class BinaryTanhQuantizer(Quantizer):
+    def __init__(
+        self,
+        hard: bool = True,
+        scale: float = 1.0,
+    ):
+        super().__init__(hard)
+        self._scale = scale
+
+    def _soft_quantize(self, x, *args, **kwargs):
+        return torch.tanh(self._scale * x)
+
+    def _hard_quantize(self, x, *args, **kwargs):
+        return torch.where(x >= 0, 1.0, -1.0)
+    
+
+class PackedBinaryQuantizer:
+    def __call__(self, x: torch.Tensor) -> torch.Tensor:
+        bits = np.where(x.cpu().numpy() >=  0, True, False)
+        packed = np.packbits(bits, axis=-1)
+        return torch.from_numpy(packed).to(x.device)
+
+
+class FlexibleQuantizer(Module):
+    def __init__(self):
+        super().__init__()
+        self._int8_quantizer = Int8TanhQuantizer()
+        self._binary_quantizer = BinaryTanhQuantizer()
+        self._packed_binary_quantizer = PackedBinaryQuantizer()
+
+    def forward(
+        self,
+        features: dict[str, torch.Tensor],
+        quantization: Literal["int8", "binary", "ubinary"] = "int8",
+        **kwargs
+    ) -> dict[str, torch.Tensor]:
+        if quantization == "int8":
+            features["sentence_embedding"] = self._int8_quantizer(
+                features["sentence_embedding"]
+            )
+        elif quantization == "binary":
+            features["sentence_embedding"] = self._binary_quantizer(
+                features["sentence_embedding"]
+            )
+        elif quantization == "ubinary":
+            features["sentence_embedding"] = self._packed_binary_quantizer(
+                features["sentence_embedding"]
+            )
+        else:
+            raise ValueError(
+                f"Invalid quantization type: {quantization}. Must be 'binary', 'ubinary', or 'int8'."
+            )
+        return features
+
+    @classmethod
+    def load(
+        cls,
+        model_name_or_path: str,
+        subfolder: str = "",
+        token: bool | str | None = None,
+        cache_folder: str | None = None,
+        revision: str | None = None,
+        local_files_only: bool = False,
+        **kwargs,
+    ):
+        return cls()
+        
+    def save(self, output_path: str, *args, **kwargs) -> None: 
+        return

tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:db3914d7cce5125c42bbbf875116cef2697023ba144bda8264ad1368595dde2b
+size 11423107

tokenizer_config.json

@@ -0,0 +1,16 @@
+{
+  "add_prefix_space": false,
+  "backend": "tokenizers",
+  "bos_token": null,
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|endoftext|>",
+  "errors": "replace",
+  "is_local": true,
+  "mask_token": "â½Ĺ",
+  "model_max_length": 131072,
+  "pad_token": "<|endoftext|>",
+  "sep_token": "<|endoftext|>",
+  "split_special_tokens": false,
+  "tokenizer_class": "Qwen2Tokenizer",
+  "unk_token": null
+}