TaxRAG Embedding — Granite 97m Multilingual

A 97M-parameter bilingual (English/Hindi) sentence-embedding model fine-tuned for dense retrieval over Indian Income-Tax Act statute text.

Overview

• What it is: a 97M-param multilingual sentence-embedding model, contrastively fine-tuned from ibm-granite/granite-embedding-97m-multilingual-r2 for retrieval over Indian Income-Tax Act statute in English and Hindi (with Hinglish query support).
• Why it exists: to fit small deployment envelopes (Streamlit Cloud ~1 GB tier) while retaining Hindi coverage.
• Size: ~211 MB on disk (bf16), ~250-300 MB peak RAM at load, 384-dim output.
• Output: 384-dim embeddings, Matryoshka-nested at [384, 256, 128] — use truncate_dim= at inference to shrink.

Intended use

• Dense retrieval / semantic search over legal/tax corpora in EN, HI, and Latin-transliterated HI (Hinglish).
• Suitable for hybrid retrieval pipelines (paired with BM25).
• Not intended for classification, generation, or non-legal open-domain retrieval — use a general-purpose encoder for those.

How to use

from sentence_transformers import SentenceTransformer

model = SentenceTransformer("vivekkopthsd/taxrag-embedding-granite-97m-multilingual")

queries = [
    "What deductions are allowed for life insurance premium?",
    "जीवन बीमा प्रीमियम के लिए क्या कटौती अनुमेय है?",
]
docs = [
    "Section 123: Deduction for life insurance premia, deferred annuity...",
    "...",
]

q_emb = model.encode(queries, normalize_embeddings=True)
d_emb = model.encode(docs, normalize_embeddings=True)
scores = q_emb @ d_emb.T

Matryoshka truncation: the model was trained with MatryoshkaLoss at dims [384, 256, 128]. For lighter indexes, pass truncate_dim=256 or truncate_dim=128 to SentenceTransformer(...) (or slice the 384-dim vector and re-normalize) — quality degrades gracefully. Always call .encode(..., normalize_embeddings=True) since the model was trained and evaluated with normalized (cosine) embeddings.

Training

• Base: ibm-granite/granite-embedding-97m-multilingual-r2 (Apache-2.0).
• Dataset: 29,440 anchor/positive/negative triples — Tier 1 bilingual statute pairs (Income-Tax Act 2025, FA 2026, 1,237 EN sections + Hindi translations, deterministic query paraphrases per section with BM25-mined hard negatives) + Tier 2 general finance breadth from BeIR/fiqa train qrels.
• Loss: CachedMultipleNegativesRankingLoss (InfoNCE, scale 20) wrapped in MatryoshkaLoss at dims [384, 256, 128] with hardness-weighted negatives.
• Optimizer: AdamW, LR 1e-5, warmup ratio 0.1, weight decay 0.01, max grad norm 1.0.
• Precision: bf16 weights (Granite-native).
• Batch: per-device 4 × grad-accum 16 = effective 64.
• Epochs: 3 (1,380 total steps).
• Model selection: load_best_model_at_end=True on eval_mean_recall@5 (best at step 750/1380).
• Seed: 7.
• Hardware: one NVIDIA RTX 3090 (24 GB). Total training wall time ~24 minutes.

Evaluation

Two evaluations were run. All numbers below are computed by the released ft_report.json (link at bottom).

In-domain: TaxRAG bilingual slice eval — 100 EN + 50 HI + 30 Hinglish gold queries over 2,474 bilingual chunks. Recall@5 / MRR@5:

Out-of-distribution generalization gate: FiQA-2018 test — 648 queries, financial Q&A domain the model was NOT fine-tuned on. Comparison to the stock base:

The fine-tune transfers positively to out-of-distribution finance retrieval (Δ R@5 = +0.014). This is unusual for a small domain-fine-tune and suggests the multilingual+Matryoshka+hard-negatives training regime does not collapse general retrieval capacity.

Known limitations & honest caveats

1. The in-domain eval is easier than a real user Q&A benchmark. The TaxRAG eval queries are constructed by deterministic paraphrases of section headings (e.g. "Deduction for life insurance premia..." → "What is the tax benefit available for life insurance premia?"). Many gold-section nouns appear verbatim in the query, so retrieval reduces partly to lexical overlap. Real user queries that avoid section-heading vocabulary will score lower. The Hinglish slice (0.50 R@5) — where transliteration destroys the exact-noun-match shortcut — is likely the closer estimate of realistic difficulty. Read the EN/HI numbers as an upper bound, not a deployment estimate.

2. Training and eval share sections. All 1,237 sections appear in both training pair anchors and eval queries. Anchor-string dedup was applied, but not section-level dedup — so the model has seen (a different paraphrase of) every eval section during training. A stricter held-out-sections evaluation would produce lower numbers.

3. Hinglish is under-served. Training Hinglish anchors come from a simple deterministic Devanagari → Latin transliterator, not natural Hinglish. Real Hinglish user queries — with code-switching, English intrusion, dialectal variation — will differ.

4. Domain scope. The model is fine-tuned on the Indian Income-Tax Act. Do not expect strong performance on GST, corporate law, criminal law, or non-Indian tax jurisdictions without further training.

5. Not evaluated on standard multilingual retrieval benchmarks. No MIRACL / Mr.TyDi / MTEB numbers are reported here — the fine-tune's target domain is narrow, and general-benchmark numbers would not be meaningful.

6. Best checkpoint at 54% of training (step 750/1380). The remaining epochs did not improve eval_mean_recall@5 — consistent with the small corpus size (~30k pairs). For future work, shortening to ~1.5 epochs or increasing dropout may reduce compute without hurting quality.

Deployment notes

• Streamlit Community Cloud: fits the free tier — ~211 MB on disk, ~250-300 MB peak RAM at load with bf16. Combined with a small BM25 sidecar and app code, stays under the 1 GB container ceiling. This was the primary motivation for the base-model swap from google/embeddinggemma-300m (>1 GB fp32, OOMs the free tier).
• Matryoshka truncation: load once, embed at 384-dim, then .truncate to 256 or 128 at query time for lighter downstream indexes. Retrieval quality degrades gracefully; 256-dim usually keeps >95% of the 384-dim quality.
• Hybrid retrieval strongly recommended: the model is a re-ranker-friendly candidate but for production, pair it with BM25 (e.g., rank_bm25 over the same corpus) to catch the exact-phrase queries that dense retrievers can miss.

Reproducibility

• Full training script + docs: TaxRAG GitHub repo — TODO: user to fill
• Training config JSON: attached as ft_report.json in this repo.
• Base model: ibm-granite/granite-embedding-97m-multilingual-r2

License

Apache 2.0 (inherited from base). See LICENSE file in the base model repo.

Citation

If you use this model, please cite the IBM Granite Embedding paper (base model) and note the fine-tune in your acknowledgements.

@misc{granite-embedding-97m-r2,
  title  = {Granite Embedding 97M Multilingual R2},
  author = {IBM Research},
  year   = {2026},
  url    = {https://huggingface.co/ibm-granite/granite-embedding-97m-multilingual-r2}
}