bge-small-en-MNN

Published by TokForge for on-device inference.

TokForge

• Website: https://tokforge.ai
• Discord: https://discord.gg/Acv3CBtfVm
• Google Play: https://play.google.com/store/apps/details?id=dev.tokforge
• iOS TestFlight: https://testflight.apple.com/join/jnufjzRr

Runs on-device in the TokForge app.

bge-small-en-v1.5: MNN format (fp16)

MNN-format conversion of BAAI/bge-small-en-v1.5 for use in on-device inference via the MNN runtime (Android / Linux / iOS). Produced for TokForge to consolidate embedding + LLM inference on a single runtime and remove a GGUF/llama.cpp dependency for RAG and memory/reflection embeddings.

Model details

• Base model: BAAI/bge-small-en-v1.5 (BERT encoder, 33M params)
• Architecture: 12-layer BERT, hidden=384, 12 attention heads, intermediate=1536
• Pooling: CLS token (token 0 of last_hidden_state), then L2-normalize
• Max sequence length: 512
• Embedding dim: 384
• Precision: fp16 weights
• File size: 64 MB

Conversion pipeline

1. Downloaded BAAI/bge-small-en-v1.5 from Hugging Face: the repo ships an onnx/model.onnx (opset 11) exported from PyTorch.
2. Ran MNNConvert -f ONNX --fp16 against that ONNX, which fused 25 LayerNorms and produced encoder.mnn.
3. Copied tokenizer files (tokenizer.json, tokenizer_config.json, special_tokens_map.json, vocab.txt) from the source repo.

Reproduce:

MNNConvert -f ONNX \
    --modelFile BAAI/bge-small-en-v1.5/onnx/model.onnx \
    --MNNModel encoder.mnn \
    --fp16 --bizCode bge-small

A/B quality vs GGUF baseline

Synthetic retrieval benchmark over 300 chunks + 20 gold queries drawn from an internal infra documentation corpus. GGUF baseline ran via llama.cpp llama-embedding against the standard BAAI bge-small-en-v1.5.Q*_K_M.gguf. Both models use CLS pooling + L2 normalization and the BGE query prefix "Represent this sentence for searching relevant passages: ".

Cross-model per-chunk cosine similarity: mean 0.9962, 5th-percentile 0.9948. The MNN fp16 conversion produces essentially the same embedding subspace as the PyTorch/ONNX reference (cosine 1.0000 vs sentence-transformers on sanity checks). There is no measurable retrieval regression.

Inputs / outputs

Input tensors (int32):

• input_ids: shape (batch, seq)
• attention_mask: shape (batch, seq)
• token_type_ids: shape (batch, seq), all zeros for single-text embedding

Output tensor (fp32 on CPU):

• last_hidden_state: shape (batch, seq, 384): Caffe/NCHW layout (use Tensor_DimensionType_Caffe when copying to host).

To get a sentence embedding: take position 0 (CLS token), then L2 normalize.

Python usage (reference)

import MNN, numpy as np
from transformers import AutoTokenizer

tok = AutoTokenizer.from_pretrained("darkmaniac7/bge-small-en-MNN")
interp = MNN.Interpreter("encoder.mnn")
sess = interp.createSession({"backend": "CPU", "numThread": 4})
in_ids  = interp.getSessionInput(sess, "input_ids")
in_mask = interp.getSessionInput(sess, "attention_mask")
in_tok  = interp.getSessionInput(sess, "token_type_ids")
out_t   = interp.getSessionOutput(sess, "last_hidden_state")

def embed(text):
    enc = tok(text, padding=False, truncation=True, max_length=512, return_tensors="np")
    seq = enc["input_ids"].shape[1]
    for t, arr in ((in_ids, enc["input_ids"]), (in_mask, enc["attention_mask"]), (in_tok, enc["token_type_ids"])):
        interp.resizeTensor(t, (1, seq))
    interp.resizeSession(sess)
    for t, arr in ((in_ids, enc["input_ids"]), (in_mask, enc["attention_mask"]), (in_tok, enc["token_type_ids"])):
        tmp = MNN.Tensor((1, seq), MNN.Halide_Type_Int,
                         arr.astype(np.int32), MNN.Tensor_DimensionType_Tensorflow)
        t.copyFrom(tmp)
    interp.runSession(sess)
    shape = out_t.getShape()
    host = MNN.Tensor(shape, MNN.Halide_Type_Float,
                      np.zeros(tuple(shape), dtype=np.float32),
                      MNN.Tensor_DimensionType_Caffe)
    out_t.copyToHostTensor(host)
    arr = np.array(host.getData(), dtype=np.float32).reshape(tuple(shape))
    cls = arr[0, 0, :]
    return cls / (np.linalg.norm(cls) + 1e-12)

# For asymmetric retrieval, prefix queries per BGE convention:
qvec = embed("Represent this sentence for searching relevant passages: " + "your query here")
pvec = embed("passage text here")
score = float(np.dot(qvec, pvec))  # cosine (both are unit-norm)

Credits

• bge-small-en-v1.5 by BAAI (Hugging Face). Used under the MIT license. Source weights are Copyright 2023 Beijing Academy of Artificial Intelligence, MIT-licensed.
• The bge_dynamic_int8*.tflite files are TensorFlow Lite INT8 exports of the same BAAI model.
• Converted from the source ONNX graph with MNNConvert from MNN by Alibaba (Apache-2.0).