A LoRA adapter over BAAI/bge-small-en-v1.5, fine-tuned for first-stage retrieval over Indian tax-law judgments. Given a question about tax law, it ranks passages of court judgments by how likely they are to answer it.

On a held-out set of 256 queries searched against a 3,868-passage corpus, it improves Recall@10 by +41.9% relative over the un-tuned base model, and Recall@1 by +62.3%.

This is stage one of a two-stage stack. A matching cross-encoder reranker, ms-marco-MiniLM-L6-v2-indian-tax-law, reorders this model's shortlist for a further +12.4% nDCG@10.

⚠️ Queries need a prefix, passages do not

Like the base BGE model, queries are embedded with a prefix and passages are embedded bare. Every number below was measured that way. Omitting the prefix silently degrades retrieval — it does not raise an error, so this is easy to get wrong.

The prefix is registered on this model as a named prompt, so you do not have to type it:

python

from sentence_transformers import SentenceTransformer

model = SentenceTransformer("FlorianTraugott/bge-small-en-v1.5-indian-tax-law-lora")
model.max_seq_length = 512

query_emb = model.encode(["Can a reassessment notice be issued after four years?"],
                         prompt_name="query",            # <- applies the prefix
                         normalize_embeddings=True)
passage_emb = model.encode(["<judgment passage text>"],  # <- no prompt_name: passages stay bare
                           normalize_embeddings=True)

scores = query_emb @ passage_emb.T   # cosine similarity

The equivalent explicit form, if you prefer to see the string:

python

PREFIX = "Represent this sentence for searching relevant passages: "  # trailing space matters

query_emb = model.encode([PREFIX + "Can a reassessment notice be issued after four years?"],
                         normalize_embeddings=True)

Both produce the same embeddings (verified: cosine 0.99999988, i.e. identical to float32 precision; the same query with no prefix scores 0.959). Note that default_prompt_name is deliberately left unset — a default would also prefix your passages, which is wrong.

The base model's weights are downloaded automatically; this repository contains only the 1.8 MB LoRA adapter.

Results

256 held-out queries, searched against the full 3,868-passage corpus. "Base" is BAAI/bge-small-en-v1.5 un-tuned, measured on the identical harness.

metric	base	this model	relative
Recall@1	0.2070	0.3359	+62.3%
Recall@3	0.3633	0.5117	+40.9%
Recall@5	0.4297	0.5938	+38.2%
Recall@10	0.4844	0.6875	+41.9%
MRR@10	0.2996	0.4431	+47.9%
nDCG@10	0.3444	0.5015	+45.6%

All 9 metrics measured improved. The reported checkpoint is the final training step, not the best-scoring one — nothing selected it, so it carries no selection optimism. The best checkpoint scored 0.6836 Recall@10, within 0.004 of it.

Which checkpoint this is. The published weights are the final-step checkpoint (step 567), the one these numbers describe. The eval-selected best checkpoint (step 300, Recall@10 0.6836) is not published — it was chosen using the evaluation set, so its figures carry selection optimism. If you are comparing against the GitHub repository, note that its local models/bge-small-tax-lora/adapter/ directory is the step-300 checkpoint, not this one.

These are relative improvements on one fixed instrument. The absolute values reflect how hard this corpus is and are not comparable to scores from other benchmarks.

Training


Base model	`BAAI/bge-small-en-v1.5` (MIT, 33.4M params, 384-dim)
Method	LoRA via PEFT — r=16, α=32, dropout 0.05, on query/key/value
Trainable	442,368 params (1.31%); base weights frozen
Loss	`MultipleNegativesRankingLoss` (InfoNCE), scale 20.0
Data	1,132 queries, each with 1 positive + 5 mined hard negatives
Batch	6, `NO_DUPLICATES` sampler, 36 candidates/query (30 hard negatives)
Schedule	3 epochs, 567 steps, lr 2e-4, linear, 57 warmup
Precision	fp32 on Apple Silicon MPS — 49 minutes

Training data

Corpus — 3,868 passages chunked from 400 Indian tax judgments, filtered from opennyaiorg/InJudgements_dataset (Apache-2.0) by Case_Type == 'Tax' plus keyword, quality, boilerplate and near-duplicate gates. The build is seeded and deterministic.
Queries — synthetic, generated from the passages with OpenAI gpt-4.1-mini, with anti-echo measures (distinctive-vocabulary overlap scoring and a hard n-gram rejection loop) so queries do not simply quote their answer.
Hard negatives — 5 per query, mined with the un-tuned base model from ranks past a skip depth of 25 with a 0.78 similarity ceiling, to avoid mining paraphrases of the positive as negatives.
Split — eval positives are excluded from training both as positives and as negatives; both leak checks measure 0.

Full pipeline, provenance manifests and verification scripts: https://github.com/FlorianTraugott/legal-retrieval-finetune

Intended use and limitations

Intended for first-stage retrieval over Indian tax-law judgments, as the retriever in a RAG system, ideally paired with the companion reranker ms-marco-MiniLM-L6-v2-indian-tax-law. Neither model is the full story alone: this one sets the recall ceiling, the reranker reorders within it.

Not intended for other areas of law, other jurisdictions, or general-purpose embedding — it was specialised on a narrow domain and was not evaluated anywhere else. It is a retrieval component, not a legal-advice system, and its output should not be treated as a statement of law.

Honest limitations, in full in docs/results.md and docs/eval_validity.md:

The eval queries are synthetic — LLM-written questions about a known passage, not questions a lawyer actually asked. The measured distribution is friendlier than reality.
Each query has exactly one labelled correct passage, but the corpus often contains others that answer it equally well, so a genuinely good result can score as a miss. The attainable ceiling is below 1.0 and unknown.
Corpus and queries are coupled — every query has a guaranteed answer in the index, which real retrieval does not.
No separate test split. The 256 eval queries served as both validation and test. Reporting the final step rather than the best checkpoint mitigates this.
Training truncated passages to 320 tokens (96.3% of tokens; 47.8% of passages untouched) to fit in memory, while evaluation used the full 512. The model never saw the tail of longer passages during training.
50 of 256 queries got worse, 15 of them materially, against 142 that improved. Aggregate gains are not uniform gains.

Licence

MIT, inherited from BAAI/bge-small-en-v1.5 (MIT). The training corpus derives from opennyaiorg/InJudgements_dataset (Apache-2.0). Synthetic training queries were generated using OpenAI gpt-4.1-mini; if you redistribute or build on this work, review OpenAI's terms for your own use case.

Author: Ayush Padhy (GitHub/HF: @FlorianTraugott)

Citation

bibtex

@software{legal_retrieval_finetune_2026,
  title  = {Fine-tuning a two-stage retrieval stack for Indian tax law},
  year   = {2026},
  url    = {https://github.com/FlorianTraugott/legal-retrieval-finetune}
}