rwkv-embedder-toolkit

CPU-only inference and fine-tuning toolkit for EmbeddingRWKV (RWKV-7, 144M) — the RNN text embedder. Pure PyTorch, no CUDA kernels, no server: it runs on a laptop CPU, which makes the model usable in sovereign, air-gapped or zero-budget RAG deployments.

Everything here was built and battle-tested for a production French RAG system (Morrigan, strict-RAG / zero-hallucination), and for a $0-compute fine-tuning campaign on Kaggle T4. The full write-up: Fine-tuning a 144M RNN embedder for $0 (FR/EN).

Quickstart

pip install torch huggingface_hub

from rwkv_emb import RWKVEmbedder

emb = RWKVEmbedder()                       # fetches the .pth from HF Hub
docs = emb.encode(["TCP uses a three-way handshake."])          # passages
qry  = emb.encode(["comment TCP établit-il une connexion ?"],
                  kind="query")            # cross-lingual FR -> EN works

Vectors are 768-d, L2-normalized. checkpoint= accepts any compatible .pth — including your own fine-tuned one from the kit below.

What's in the box

Path	What
rwkv_emb/	CPU pure-PyTorch inference (backbone, trie tokenizer, embedder)
finetune/train_lib.py	Self-contained fine-tuning library (contrastive, InfoNCE, hard negatives, retention pairs) — mirrors inference tokenization exactly
finetune/*.ipynb	Kaggle T4 notebooks used for the campaign rounds (free tier)
docs/	Findings: the EOS/truncation trap, why whitening does NOT help, campaign numbers
examples/embed.py	10-line cross-lingual retrieval demo
archive/campaign/	Campaign artifacts kept verbatim for reproducibility (written against our production repo — paths/imports are environment-specific)

Findings you probably want before integrating this model

1. The EOS/truncation trap (silent, catastrophic). The state pooling only reads EOS positions. Append EOS before truncating and every text longer than your token cap returns the SAME constant vector (cosine 1.0000 between different documents) while short-text tests keep passing. Truncate first, then append EOS. rwkv_emb does this correctly, and your harness should assert that two different texts never reach cos == 1.0. Details: docs/eos-truncation-trap.md.

2. The anisotropy is signal — do not whiten. Post-hoc isotropy fixes that help transformer embedders (mean-centering, ABTT, PCA whitening) all degrade this RNN embedder (hit@1 26 → 18/11/15 on our FR→EN eval). Details: docs/anisotropy-is-signal.md.

3. Honest campaign verdict. Four fine-tuning rounds ($0, Kaggle T4) improved the model round over round (+4 hit@1 with hard negatives), but on a full-scale production judge, granite-embedding-107m stayed ahead on ranking — while the RNN kept the best out-of-corpus refusal rate (22/24 vs 17/24), which matters for strict-RAG. Numbers and protocol: docs/finetune-campaign.md.

4. Pool evals overestimate RNN embedders. Cosine ranges are much tighter than transformer embedders' — small-pool metrics looked great, full-scale inverted the verdict. Always judge at production scale, and never compare numbers across two different judges.

Practical numbers (CPU)

• Model: 144M params, ~0.5 GB checkpoint, 768-d output.
• Batch 4 is the CPU sweet spot we measured (state is O(heads·N²) memory).
• Length-sorted batching minimizes padding; instruction on the query side is worth it (hit@3 18 → 26/56 on our eval WITH it).

License & attribution

Apache-2.0. Builds on RWKV (Peng Bo et al.) and howard-hou/EmbeddingRWKV — both Apache-2.0. Maintained by ScarletWolf (sovereign AI, France/Ireland).