Upload verified Needle MLX conversion

README.md

@@ -0,0 +1,62 @@
+---
+license: mit
+library_name: mlx
+base_model: Cactus-Compute/needle
+tags:
+  - mlx
+  - function-calling
+  - tool-use
+  - encoder-decoder
+  - apple-silicon
+---
+
+# Needle MLX
+
+An MLX safetensors conversion of
+[Cactus-Compute/needle](https://huggingface.co/Cactus-Compute/needle), a
+26M-parameter encoder-decoder function-calling model.
+
+## What is included
+
+- `model.safetensors`: 31 converted tensors, 26,315,421 parameters.
+- `needle_mlx.py`: a small custom MLX inference implementation for Needle's
+  architecture.
+- `tokenizer/`: the upstream SentencePiece tokenizer files.
+- `config.json` and `manifest.json`: architecture and conversion provenance.
+
+## Use
+
+```python
+from needle_mlx import NeedleModel
+
+model = NeedleModel.from_pretrained(".")
+logits = model.forward([[1, 2, 3]], [[1, 4]])
+```
+
+This is an MLX-native package, not an MLX-LM or oMLX model yet. Needle uses a
+custom JAX/Flax architecture, so load it with the included runtime or add a
+dedicated adapter to the serving runtime.
+
+## Conversion and verification
+
+The published `needle.pkl` checkpoint was read with a restricted NumPy-only
+pickle loader, converted directly to MLX safetensors, and checked
+tensor-for-tensor against the source.
+
+- Source SHA-256:
+  `40a32e91d1d4197bf15ba559b74f6727c342dc8746918742fc7d8e2c1f18df40`
+- Converted `model.safetensors` SHA-256:
+  `7b9d5f0d6ddeb7fbb20f4e45f3f616919357e5d08b5778859fdb762a33d60dae`
+- Verification: 31/31 tensors equal the source; an MLX encoder-decoder smoke
+  pass produced logits with shape `(1, 2, 8192)`.
+
+The converter source is available at
+[seeker-cyber-maker/needle-mlx-depicklinator](https://github.com/seeker-cyber-maker/needle-mlx-depicklinator).
+
+## Credits and license
+
+Needle was created by Cactus Compute. See the
+[upstream model card](https://huggingface.co/Cactus-Compute/needle) and
+[source repository](https://github.com/cactus-compute/needle) for the model,
+training details, and citation. The upstream model card publishes Needle under
+the MIT license; this converted package retains that license.

config.json

@@ -0,0 +1,29 @@
+{
+  "vocab_size": 8192,
+  "d_model": 512,
+  "num_heads": 8,
+  "num_kv_heads": 4,
+  "num_encoder_layers": 12,
+  "num_decoder_layers": 8,
+  "d_ff": 2048,
+  "max_seq_len": 1024,
+  "pad_token_id": 0,
+  "rope_theta": 10000.0,
+  "dtype": "bfloat16",
+  "activation": "swiglu",
+  "num_memory_slots": 64,
+  "n_mels": 80,
+  "dropout_rate": 0.1,
+  "contrastive_dim": 128,
+  "enable_speech": false,
+  "no_feedforward": true,
+  "model_type": "needle_mlx",
+  "architectures": [
+    "NeedleModel"
+  ],
+  "format": "mlx-safetensors",
+  "format_version": 1,
+  "eos_token_id": 1,
+  "source_checkpoint": "needle.pkl",
+  "source_sha256": "40a32e91d1d4197bf15ba559b74f6727c342dc8746918742fc7d8e2c1f18df40"
+}

manifest.json

@@ -0,0 +1,8 @@
+{
+  "format": "needle-mlx",
+  "format_version": 1,
+  "tensor_count": 31,
+  "parameter_count": 26315421,
+  "source_sha256": "40a32e91d1d4197bf15ba559b74f6727c342dc8746918742fc7d8e2c1f18df40",
+  "source_checkpoint": "needle.pkl"
+}

model.safetensors

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

needle_mlx.py

@@ -0,0 +1,127 @@
+"""Small MLX inference implementation for Cactus Needle."""
+
+from __future__ import annotations
+
+import json
+import math
+from pathlib import Path
+from typing import Any
+
+import mlx.core as mx
+
+
+class NeedleModel:
+    def __init__(self, weights: dict[str, mx.array], config: dict[str, Any]):
+        self.weights = weights
+        self.config = config
+        self.num_heads = int(config["num_heads"])
+        self.num_kv_heads = int(config["num_kv_heads"])
+        self.d_model = int(config["d_model"])
+        self.head_dim = self.d_model // self.num_heads
+
+    @classmethod
+    def from_pretrained(cls, path: str | Path) -> "NeedleModel":
+        path = Path(path)
+        config = json.loads((path / "config.json").read_text())
+        weights = mx.load(str(path / "model.safetensors"))
+        return cls(weights, config)
+
+    def _weight(self, name: str, layer: int | None = None) -> mx.array:
+        value = self.weights[name]
+        return value if layer is None else value[layer]
+
+    @staticmethod
+    def _linear(x: mx.array, kernel: mx.array, bias: mx.array | None = None) -> mx.array:
+        output = mx.matmul(x, kernel)
+        return output if bias is None else output + bias
+
+    @staticmethod
+    def _zcrms_norm(x: mx.array, scale: mx.array) -> mx.array:
+        rms = mx.sqrt(mx.mean(mx.square(x.astype(mx.float32)), axis=-1, keepdims=True) + 1e-6)
+        return ((1.0 + scale) * x / rms).astype(x.dtype)
+
+    def _rope(self, x: mx.array) -> mx.array:
+        sequence_length = x.shape[2]
+        half_dim = self.head_dim // 2
+        positions = mx.arange(sequence_length, dtype=mx.float32)
+        inv_freq = 1.0 / (
+            float(self.config["rope_theta"])
+            ** (mx.arange(0, self.head_dim, 2, dtype=mx.float32) / self.head_dim)
+        )
+        angles = mx.outer(positions, inv_freq)
+        cos = mx.reshape(mx.cos(angles), (1, 1, sequence_length, half_dim))
+        sin = mx.reshape(mx.sin(angles), (1, 1, sequence_length, half_dim))
+        first, second = x[..., :half_dim], x[..., half_dim:]
+        return mx.concatenate([first * cos - second * sin, second * cos + first * sin], axis=-1)
+
+    def _attention(
+        self,
+        q_input: mx.array,
+        kv_input: mx.array,
+        prefix: str,
+        layer: int,
+        mask: mx.array | None,
+        rope: bool,
+    ) -> mx.array:
+        q = self._linear(q_input, self._weight(f"{prefix}/q_proj/kernel", layer))
+        k = self._linear(kv_input, self._weight(f"{prefix}/k_proj/kernel", layer))
+        v = self._linear(kv_input, self._weight(f"{prefix}/v_proj/kernel", layer))
+        batch = q.shape[0]
+        q = mx.transpose(mx.reshape(q, (batch, -1, self.num_heads, self.head_dim)), (0, 2, 1, 3))
+        k = mx.transpose(mx.reshape(k, (batch, -1, self.num_kv_heads, self.head_dim)), (0, 2, 1, 3))
+        v = mx.transpose(mx.reshape(v, (batch, -1, self.num_kv_heads, self.head_dim)), (0, 2, 1, 3))
+        q = self._zcrms_norm(q, self._weight(f"{prefix}/q_norm/scale", layer))
+        k = self._zcrms_norm(k, self._weight(f"{prefix}/k_norm/scale", layer))
+        if self.num_heads != self.num_kv_heads:
+            repeats = self.num_heads // self.num_kv_heads
+            k = mx.repeat(k, repeats, axis=1)
+            v = mx.repeat(v, repeats, axis=1)
+        if rope:
+            q, k = self._rope(q), self._rope(k)
+        scores = mx.matmul(q.astype(mx.float32), mx.transpose(k.astype(mx.float32), (0, 1, 3, 2)))
+        scores = scores / math.sqrt(self.head_dim)
+        if mask is not None:
+            scores = mx.where(mask, scores, mx.array(-1e9, dtype=mx.float32))
+        output = mx.matmul(mx.softmax(scores, axis=-1).astype(v.dtype), v)
+        output = mx.reshape(mx.transpose(output, (0, 2, 1, 3)), (batch, -1, self.d_model))
+        return self._linear(output, self._weight(f"{prefix}/out_proj/kernel", layer))
+
+    def encode(self, source_tokens: mx.array) -> tuple[mx.array, mx.array]:
+        source_tokens = mx.array(source_tokens, dtype=mx.int32)
+        padding_mask = source_tokens != int(self.config["pad_token_id"])
+        mask = mx.reshape(padding_mask, (source_tokens.shape[0], 1, 1, source_tokens.shape[1]))
+        x = self.weights["embedding/embedding"][source_tokens] * math.sqrt(self.d_model)
+        for layer in range(int(self.config["num_encoder_layers"])):
+            prefix = "encoder/layers/EncoderBlock_0"
+            residual = x
+            x = self._zcrms_norm(x, self._weight(f"{prefix}/ZCRMSNorm_0/scale", layer))
+            x = self._attention(x, x, f"{prefix}/self_attn", layer, mask, rope=True)
+            gate = mx.sigmoid(self._weight(f"{prefix}/attn_gate", layer))
+            x = residual + gate * x
+        return self._zcrms_norm(x, self.weights["encoder/final_norm/scale"]), mask
+
+    def decode(self, target_tokens: mx.array, encoder_out: mx.array, cross_mask: mx.array) -> mx.array:
+        target_tokens = mx.array(target_tokens, dtype=mx.int32)
+        length = target_tokens.shape[1]
+        causal = mx.tril(mx.ones((length, length), dtype=mx.bool_))
+        target_valid = target_tokens != int(self.config["pad_token_id"])
+        self_mask = mx.reshape(causal, (1, 1, length, length)) & mx.reshape(
+            target_valid, (target_tokens.shape[0], 1, 1, length)
+        )
+        x = self.weights["embedding/embedding"][target_tokens] * math.sqrt(self.d_model)
+        prefix = "decoder/layers/DecoderBlock_0"
+        for layer in range(int(self.config["num_decoder_layers"])):
+            residual = x
+            x = self._zcrms_norm(x, self._weight(f"{prefix}/ZCRMSNorm_0/scale", layer))
+            x = self._attention(x, x, f"{prefix}/self_attn", layer, self_mask, rope=True)
+            x = residual + mx.sigmoid(self._weight(f"{prefix}/self_attn_gate", layer)) * x
+            residual = x
+            x = self._zcrms_norm(x, self._weight(f"{prefix}/ZCRMSNorm_1/scale", layer))
+            x = self._attention(x, encoder_out, f"{prefix}/cross_attn", layer, cross_mask, rope=False)
+            x = residual + mx.sigmoid(self._weight(f"{prefix}/cross_attn_gate", layer)) * x
+        x = self._zcrms_norm(x, self.weights["decoder/ZCRMSNorm_0/scale"])
+        return mx.matmul(x.astype(mx.float32), mx.transpose(self.weights["embedding/embedding"].astype(mx.float32)))
+
+    def forward(self, source_tokens: mx.array, target_tokens: mx.array) -> mx.array:
+        encoder_out, cross_mask = self.encode(source_tokens)
+        return self.decode(target_tokens, encoder_out, cross_mask)

tokenizer/needle.model

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