bge-m3-lite 0.0.2__tar.gz

bge_m3_lite-0.0.2/PKG-INFO

@@ -0,0 +1,103 @@
+Metadata-Version: 2.4
+Name: bge-m3-lite
+Version: 0.0.2
+Summary: Lightweight CPU inference for BAAI/bge-m3 (dense + sparse + ColBERT) with onnxruntime as the only dependency.
+Keywords: bge-m3,colbert,embedding,onnxruntime,retrieval,sparse
+Author: Allen Chou
+License-Expression: MIT
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: Operating System :: MacOS
+Classifier: Operating System :: POSIX :: Linux
+Classifier: Programming Language :: Python :: 3 :: Only
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3.14
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Dist: onnxruntime>=1.17
+Requires-Dist: onnx>=1.16 ; extra == 'quant'
+Requires-Dist: onnx-ir ; extra == 'quant'
+Requires-Dist: sympy ; extra == 'quant'
+Requires-Python: >=3.11
+Project-URL: Homepage, https://github.com/allen2c/bge-m3-lite
+Provides-Extra: quant
+Description-Content-Type: text/markdown
+
+# bge-m3-lite
+
+CPU inference for [BAAI/bge-m3](https://huggingface.co/BAAI/bge-m3) with
+**`onnxruntime` as the only dependency**. All three BGE-M3 outputs are supported
+and match the official PyTorch implementation (FlagEmbedding) to fp32 precision:
+
+| output | shape | notes |
+|---|---|---|
+| `dense_vecs` | `(n, 1024)` | CLS pooling, L2-normalised |
+| `lexical_weights` | `list[dict[str, float]]` | token-id → weight, max-pooled, specials removed |
+| `colbert_vecs` | `list[(len-1, 1024)]` | per-token vectors without `<s>`, L2-normalised |
+
+Everything except the transformer forward pass is implemented in this package
+from scratch: the XLM-RoBERTa tokenizer (SentencePiece unigram model, the
+`nmt_nfkc` precompiled charsmap, Unicode grapheme segmentation), the torch-free
+loader for the sparse / ColBERT heads, the model downloader and the pooling.
+
+Platforms: Apple Silicon, Linux ARM64, Linux x86_64 (Python 3.11+).
+
+## Install
+
+```bash
+uv add bge-m3-lite        # or: pip install bge-m3-lite
+```
+
+## Use
+
+```python
+from bge_m3_lite import BGEM3Embedder
+
+embedder = BGEM3Embedder()  # first call downloads ~2.3 GB into ~/.cache/bge-m3-lite
+out = embedder.encode(
+    ["What is BGE M3?", "BGE M3 是一個多語言嵌入模型。"],
+    return_dense=True,
+    return_sparse=True,
+    return_colbert_vecs=True,
+)
+out["dense_vecs"].shape  # (2, 1024)
+out["lexical_weights"][0]  # {'4865': 0.08, '83': 0.08, ...}
+out["colbert_vecs"][0].shape  # (7, 1024)
+
+embedder.convert_id_to_token(out["lexical_weights"][0])
+embedder.compute_lexical_matching_score(lw_query, lw_passage)
+embedder.colbert_score(q_vecs, p_vecs)
+```
+
+Passing a single string returns unwrapped values, like FlagEmbedding.
+`BGEM3Embedder(precision="int8")` loads a 4× smaller quantised backbone
+(see `docs/quantization.md` for the accuracy trade-off).
+
+### CLI
+
+```bash
+bge-m3-lite download                     # pre-fetch the model files
+bge-m3-lite info                         # cache state
+echo "hello" | bge-m3-lite encode --sparse --colbert --tokens
+```
+
+### Environment variables
+
+| variable | effect |
+|---|---|
+| `BGE_M3_LITE_CACHE` | cache directory (default `~/.cache/bge-m3-lite/BAAI--bge-m3`) |
+| `HF_ENDPOINT` | Hugging Face mirror, e.g. `https://hf-mirror.com` |
+| `BGE_M3_LITE_OFFLINE=1` | never download, fail if files are missing |
+| `BGE_M3_LITE_THREADS` | onnxruntime intra-op threads (default: physical cores) |
+
+Model files are pinned to a specific Hugging Face revision and verified by
+SHA-256 after download.
+
+## Development
+
+See `AGENTS.md` and `docs/` (architecture, tokenizer, verification, development).
+
+## Status
+
+v0.0.2: fp32 with exact parity with FlagEmbedding, plus an opt-in int8 backbone.

bge_m3_lite-0.0.2/README.md

@@ -0,0 +1,77 @@
+# bge-m3-lite
+
+CPU inference for [BAAI/bge-m3](https://huggingface.co/BAAI/bge-m3) with
+**`onnxruntime` as the only dependency**. All three BGE-M3 outputs are supported
+and match the official PyTorch implementation (FlagEmbedding) to fp32 precision:
+
+| output | shape | notes |
+|---|---|---|
+| `dense_vecs` | `(n, 1024)` | CLS pooling, L2-normalised |
+| `lexical_weights` | `list[dict[str, float]]` | token-id → weight, max-pooled, specials removed |
+| `colbert_vecs` | `list[(len-1, 1024)]` | per-token vectors without `<s>`, L2-normalised |
+
+Everything except the transformer forward pass is implemented in this package
+from scratch: the XLM-RoBERTa tokenizer (SentencePiece unigram model, the
+`nmt_nfkc` precompiled charsmap, Unicode grapheme segmentation), the torch-free
+loader for the sparse / ColBERT heads, the model downloader and the pooling.
+
+Platforms: Apple Silicon, Linux ARM64, Linux x86_64 (Python 3.11+).
+
+## Install
+
+```bash
+uv add bge-m3-lite        # or: pip install bge-m3-lite
+```
+
+## Use
+
+```python
+from bge_m3_lite import BGEM3Embedder
+
+embedder = BGEM3Embedder()  # first call downloads ~2.3 GB into ~/.cache/bge-m3-lite
+out = embedder.encode(
+    ["What is BGE M3?", "BGE M3 是一個多語言嵌入模型。"],
+    return_dense=True,
+    return_sparse=True,
+    return_colbert_vecs=True,
+)
+out["dense_vecs"].shape  # (2, 1024)
+out["lexical_weights"][0]  # {'4865': 0.08, '83': 0.08, ...}
+out["colbert_vecs"][0].shape  # (7, 1024)
+
+embedder.convert_id_to_token(out["lexical_weights"][0])
+embedder.compute_lexical_matching_score(lw_query, lw_passage)
+embedder.colbert_score(q_vecs, p_vecs)
+```
+
+Passing a single string returns unwrapped values, like FlagEmbedding.
+`BGEM3Embedder(precision="int8")` loads a 4× smaller quantised backbone
+(see `docs/quantization.md` for the accuracy trade-off).
+
+### CLI
+
+```bash
+bge-m3-lite download                     # pre-fetch the model files
+bge-m3-lite info                         # cache state
+echo "hello" | bge-m3-lite encode --sparse --colbert --tokens
+```
+
+### Environment variables
+
+| variable | effect |
+|---|---|
+| `BGE_M3_LITE_CACHE` | cache directory (default `~/.cache/bge-m3-lite/BAAI--bge-m3`) |
+| `HF_ENDPOINT` | Hugging Face mirror, e.g. `https://hf-mirror.com` |
+| `BGE_M3_LITE_OFFLINE=1` | never download, fail if files are missing |
+| `BGE_M3_LITE_THREADS` | onnxruntime intra-op threads (default: physical cores) |
+
+Model files are pinned to a specific Hugging Face revision and verified by
+SHA-256 after download.
+
+## Development
+
+See `AGENTS.md` and `docs/` (architecture, tokenizer, verification, development).
+
+## Status
+
+v0.0.2: fp32 with exact parity with FlagEmbedding, plus an opt-in int8 backbone.

bge_m3_lite-0.0.2/bge_m3_lite/__init__.py

@@ -0,0 +1,6 @@
+"""bge-m3-lite: CPU inference for BAAI/bge-m3, onnxruntime is the only dependency."""
+
+from bge_m3_lite.embedder import BGEM3Embedder
+
+__version__ = "0.0.2"
+__all__ = ["BGEM3Embedder", "__version__"]

bge_m3_lite-0.0.2/bge_m3_lite/_grapheme.py

@@ -0,0 +1,128 @@
+"""Extended grapheme cluster segmentation (UAX #29), pure Python.
+
+Needed because Hugging Face's ``Precompiled`` normalizer applies the
+SentencePiece charsmap per grapheme cluster rather than by longest prefix.
+Tables come from ``_grapheme_data`` (generated from the Unicode database).
+"""
+
+from __future__ import annotations
+
+from bisect import bisect_right
+from collections.abc import Iterator
+
+from bge_m3_lite import _grapheme_data as _d
+
+# Grapheme_Cluster_Break values (indices into _d.GCB_NAMES)
+_OTHER, _CR, _LF, _CONTROL, _EXTEND, _ZWJ, _RI, _PREPEND, _SPACINGMARK = range(9)
+_L, _V, _T, _LV, _LVT = range(9, 14)
+# Indic_Conjunct_Break values
+_INCB_NONE, _INCB_CONSONANT, _INCB_EXTEND, _INCB_LINKER = range(4)
+
+
+def _lookup(cp: int, starts, ends, values) -> int:
+    i = bisect_right(starts, cp) - 1
+    if i >= 0 and cp <= ends[i]:
+        return values[i]
+    return 0
+
+
+_gcb_cache: dict[int, int] = {}
+_incb_cache: dict[int, int] = {}
+_ext_cache: dict[int, bool] = {}
+
+
+def gcb(cp: int) -> int:
+    v = _gcb_cache.get(cp)
+    if v is None:
+        v = _lookup(cp, _d.GCB_STARTS, _d.GCB_ENDS, _d.GCB_VALUES)
+        if len(_gcb_cache) < 65536:
+            _gcb_cache[cp] = v
+    return v
+
+
+def incb(cp: int) -> int:
+    v = _incb_cache.get(cp)
+    if v is None:
+        v = _lookup(cp, _d.INCB_STARTS, _d.INCB_ENDS, _d.INCB_VALUES)
+        if len(_incb_cache) < 65536:
+            _incb_cache[cp] = v
+    return v
+
+
+def is_extended_pictographic(cp: int) -> bool:
+    v = _ext_cache.get(cp)
+    if v is None:
+        v = bool(_lookup(cp, _d.EXTPICT_STARTS, _d.EXTPICT_ENDS, _d.EXTPICT_VALUES))
+        if len(_ext_cache) < 65536:
+            _ext_cache[cp] = v
+    return v
+
+
+def grapheme_boundaries(text: str) -> Iterator[int]:
+    """Yield every index ``i`` (0 < i < len) where a grapheme cluster boundary falls."""
+    n = len(text)
+    if n < 2:
+        return
+    cps = [ord(c) for c in text]
+    props = [gcb(cp) for cp in cps]
+    ri_run = 0  # number of consecutive Regional_Indicator chars ending at i-1
+    for i in range(1, n):
+        prev, cur = props[i - 1], props[i]
+        if prev == _RI:
+            ri_run += 1
+        else:
+            ri_run = 0
+        # GB3
+        if prev == _CR and cur == _LF:
+            continue
+        # GB4 / GB5
+        if prev in (_CONTROL, _CR, _LF) or cur in (_CONTROL, _CR, _LF):
+            yield i
+            continue
+        # GB6 / GB7 / GB8 (Hangul)
+        if prev == _L and cur in (_L, _V, _LV, _LVT):
+            continue
+        if prev in (_LV, _V) and cur in (_V, _T):
+            continue
+        if prev in (_LVT, _T) and cur == _T:
+            continue
+        # GB9 / GB9a / GB9b
+        if cur in (_EXTEND, _ZWJ, _SPACINGMARK) or prev == _PREPEND:
+            continue
+        # GB9c: Consonant [Extend Linker]* Linker [Extend Linker]* x Consonant
+        if incb(cps[i]) == _INCB_CONSONANT:
+            j = i - 1
+            seen_linker = False
+            while j >= 0:
+                v = incb(cps[j])
+                if v == _INCB_LINKER:
+                    seen_linker = True
+                elif v != _INCB_EXTEND:
+                    break
+                j -= 1
+            if seen_linker and j >= 0 and incb(cps[j]) == _INCB_CONSONANT:
+                continue
+        # GB11: ExtPict Extend* ZWJ x ExtPict
+        if prev == _ZWJ and is_extended_pictographic(cps[i]):
+            j = i - 2
+            while j >= 0 and props[j] == _EXTEND:
+                j -= 1
+            if j >= 0 and is_extended_pictographic(cps[j]):
+                continue
+        # GB12 / GB13: keep RI pairs together
+        if prev == _RI and cur == _RI and ri_run % 2 == 1:
+            continue
+        # GB999
+        yield i
+
+
+def graphemes(text: str) -> list[str]:
+    """Split ``text`` into extended grapheme clusters."""
+    out: list[str] = []
+    start = 0
+    for b in grapheme_boundaries(text):
+        out.append(text[start:b])
+        start = b
+    if text:
+        out.append(text[start:])
+    return out