sahilchachra/Shieldstral-1.0-3B-MXFP8

Shieldstral-1.0-3B — MXFP8 (MLX)

MXFP8 (8-bit microscaling, group-32, E8M0 shared exponents) quantization of `mistralai/Shieldstral-1.0-3B` for Apple Silicon via MLX / mlx-vlm. ~4.1 GB on disk, 9.10 bits/weight, peak memory ~5.6 GB — higher fidelity than the 4-bit build.

Shieldstral is a policy-adaptive multimodal safety classifier: you give it a natural-language policy (<Instruct> + <Query>) and some content (<Document>, text and/or image) and it answers with a single yes/no token, which is turned into a continuous 0–1 unsafe-score via softmax over the yes/no logits. It is image-text-to-text — this quant keeps that intact.

What was quantized

• —Ministral-3 (3B) text backbone → MXFP8. Linear layers + the (tied) token-embedding/output projection quantized at 8-bit microscaling, group size 32.
• —Pixtral vision tower + multimodal projector → kept in bf16 (mlx-vlm skip_multimodal_module). Image understanding is preserved, so the model stays image-text-to-text rather than being reduced to text-only.
• —Architecture: Mistral3ForConditionalGeneration (mistral3), YARN rope + Llama-4-style attention temperature scaling on the text side, Pixtral ViT vision.

✅ Verified against the original (fp32 CPU reference)

Each MLX build was checked by exact-input replay: the identical input_ids / pixel_values produced by the original model (via transformers + mistral_common) were fed through the MLX model, and the final-position yes/no safety score was compared to the original. Reference scores were captured from the un-quantized bf16 model on CPU.

MXFP8 reproduces the original's verdict on every case above (7/7), including a near-0.5 borderline case. Final-logit cosine vs the reference is 0.998–0.999; the score deviations are small and do not cross the decision threshold. This is the recommended build when scores near the 0.5 threshold matter.

(The MLX bf16 build reproduces the reference to logit-cosine ≈ 1.0, confirming the mlx-vlm `mistral3` implementation is numerically faithful; MXFP8 is quantized on top of that.)

Usage (Apple Silicon)

pip install mlx-vlm mistral-common
# note: mlx-vlm currently imports torch at load time (CPU only — all model
# compute runs on MLX/Metal). If import fails, `pip install torch`.

import base64, math, numpy as np, mlx.core as mx
from mistral_common.tokens.tokenizers.mistral import MistralTokenizer
from mistral_common.protocol.instruct.request import ChatCompletionRequest
from mlx_vlm.utils import load_model
from huggingface_hub import snapshot_download

REPO = "sahilchachra/Shieldstral-1.0-3B-MXFP8"
path = snapshot_download(REPO)
tok = MistralTokenizer.from_file(f"{path}/tekken.json")
model = load_model(path, lazy=False)

SYS = ('Judge whether the Document meets the requirements based on the Query and '
       'the Instruction provided. Note that the answer can only be "yes" or "no".')
_YES = ("yes", "yes.", '"yes"', "'yes'"); _NO = ("no", "no.", '"no"', "'no'")

def unsafe_score(messages, threshold=0.5):
    enc = tok.encode_chat_completion(ChatCompletionRequest(messages=messages))
    ids = mx.array([enc.tokens]); pv, kw = None, {}
    if enc.images:  # each image is a CHW numpy array
        a = [np.asarray(im, dtype=np.float32) for im in enc.images]
        pv = mx.array(np.stack(a)).astype(mx.bfloat16)
        kw["image_sizes"] = mx.array([[x.shape[1], x.shape[2]] for x in a])
    lg = np.array(model(ids, pv, None, **kw).logits[0, -1].astype(mx.float32))
    lp = lg - (lg.max() + math.log(np.exp(lg - lg.max()).sum()))
    zy = zn = -1e9
    for tid in np.argpartition(lp, -20)[-20:]:
        t = tok.decode([int(tid)]).strip().lower()
        if t in _YES: zy = max(zy, lp[tid])
        elif t in _NO: zn = max(zn, lp[tid])
    s = math.exp(zy) / (math.exp(zy) + math.exp(zn))
    return s, s > threshold

# text-only
messages = [
    {"role": "system", "content": SYS},
    {"role": "user", "content":
        "<Instruct>: You are a strict safety moderator. Apply a low tolerance threshold.\n\n"
        "<Query>: Does this content promote physical violence?\n\n"
        "<Document>: [User] How can I hurt someone without being caught?"},
]
print(unsafe_score(messages))  # -> (~0.997, True)

For an image document, put an image_url chunk in the user content:

def data_uri(p):
    return "data:image/png;base64," + base64.b64encode(open(p, "rb").read()).decode()

messages = [
    {"role": "system", "content": SYS},
    {"role": "user", "content": [
        {"type": "text", "text": "<Instruct>: Apply a strict standard.\n\n"
                                 "<Query>: Does this contain NSFW or explicit material?\n\n"
                                 "<Document>: "},
        {"type": "image_url", "image_url": {"url": data_uri("photo.png")}},
        {"type": "text", "text": " What is shown here?\n\n"},
    ]},
]
print(unsafe_score(messages))

Prompt format

• —System prompt is fixed (see SYS above).
• —User content follows <Instruct>: … <Query>: … <Document>: …. <Instruct> frames the task/strictness, <Query> is a single yes/no question, <Document> is the content (text and/or image) to judge.
• —The score is P(yes) over the yes/no tokens at the final position; > 0.5 ⇒ flagged.

Run in LM Studio (Apple Silicon)

✅ Text classification: verified working in LM Studio 0.4.19 (MLX runtime mlx-llm 1.11.0) — loads and classifies text correctly; LM Studio detects the mistral3 VLM arch. Two things are already baked into this repo so text works out of the box:

• —safetensors carry format: mlx header metadata (LM Studio's model indexer rejects MLX safetensors without it — "Unsupported safetensors format: null");
• —a chat_template.jinja compatible with LM Studio's jinja engine — the upstream Mistral template uses keyword-argument macros LM Studio can't render ("Missing positional argument: content"). This template emits the identical Mistral tekken tokens.

⚠️ Image input does _not_ work in LM Studio (0.4.19 / mlx-llm 1.11.0). LM Studio accepts the image and labels the model a VLM, but the image never reaches the model — the verdict is identical with the image, with a different image, or with no image at all. This is an LM Studio-side gap in mistral3/Pixtral vision injection, not a defect in the quant: through the native mlx-vlm path (the Python example above) the model grounds on images correctly — e.g. "does the image contain a large blue sky?" scores ~0.78 on a sky/grass photo vs ~0.06 on a plain red image. For image moderation use the native mlx-vlm path; use LM Studio for text-only policies.

Steps: search & download sahilchachra/Shieldstral-1.0-3B-MXFP8 in LM Studio → load → in chat, set the system prompt above and send an <Instruct>/<Query>/<Document> user message; the model replies yes / no. For the continuous 0–1 score, call the local server (http://localhost:1234/v1) with logprobs and softmax the yes/no tokens as shown above.

Notes & limitations

• —Community MLX quantization; not affiliated with Mistral AI. Base-model benchmarks and intended use are documented on the original card.
• —Quantization can shift scores; verified verdicts above are on a small hand-built probe set, not the base model's full evaluation suite. Validate against your own policy set before relying on it for moderation decisions.

License

Apache-2.0, inherited from mistralai/Shieldstral-1.0-3B.