Needle MLX

MLX port of Cactus-Compute/needle, a 26M-parameter encoder-decoder tool-calling model.

Run

uv run python - <<'PY'
import mlx.core as mx
from model import SimpleAttentionNetwork
from tokenizer import NeedleTokenizer, build_encoder_input

path = "mlx-model"  # or a downloaded Hugging Face snapshot
model = SimpleAttentionNetwork.from_pretrained(path)
tokenizer = NeedleTokenizer.from_pretrained(path)
tools = '[{"name":"get_weather","parameters":{"location":"string"}}]'
src = mx.array([build_encoder_input(tokenizer, "Weather in Delhi?", tools)])
print(tokenizer.decode(model.generate(src, max_new_tokens=128)).removeprefix("<tool_call>"))
PY

Create the upload directory from the original checkpoint:

uv run convert.py weights mlx-model
uvx --from huggingface_hub hf upload mlx-community/needle-mlx mlx-model .

Compatibility

This repository uses a custom Needle architecture. It runs through the included MLX implementation; it is not currently supported by mlx-lm, Ollama, llama.cpp, or LM Studio. GGUF or an Ollama Modelfile alone cannot add the missing architecture implementation.

License

MIT. Original model and weights: Cactus Compute.