saitejasathiraju/indian-legal-reranker

Indian Legal Reranker (Qwen3-Reranker-0.6B, LoRA fine-tuned)

A pointwise reranker fine-tuned from Qwen/Qwen3-Reranker-0.6B for Indian criminal law retrieval: BNS 2023 (Bharatiya Nyaya Sanhita), BNSS 2023 (Bharatiya Nagarik Suraksha Sanhita), and BSA 2023 (Bharatiya Sakshya Adhiniyam) — the acts that replaced the IPC, CrPC, and Indian Evidence Act in July 2024.

Given a query and a candidate passage, the model judges relevance as "yes" or "no", and the relative logit between those two tokens gives a continuous relevance score — designed to be dropped into a RAG reranking step over retrieved statute/case-law chunks.

Why this exists

Standalone generation from small legal-domain LLMs (including the base 0.6B/1.5B/3B Qwen models fine-tuned only on QA pairs) reliably hallucinates specific section citations for BNS/BNSS/BSA 2023 — these acts are too recent (in force July 2024) to be well represented in most pretraining data, and QA-pair fine-tuning teaches style, not facts. This reranker is intended as the reranking stage of a retrieval-augmented pipeline (BM25/dense retrieval over the actual bare-act text, reranked by this model, then fed to a generator) rather than a standalone answer engine.

Training

• Base model: Qwen3-Reranker-0.6B (Apache 2.0)
• Method: LoRA (r=16, alpha=32, target modules: q/k/v/o/gate/up/down_proj), merged into the base weights for this repo
• Data: ~11,500 (query, positive passage, negative passage) triplets covering BNS/BNSS/BSA 2023 sections, Supreme/High Court judgment excerpts, and IRAC-style legal analysis
• Objective: binary cross-entropy over the "yes"/"no" token pair at the final position, following the base model's native pointwise scoring convention
• Result: pairwise accuracy on a 400-triplet held-out set improved from 0.9900 (base model, zero-shot) to 1.0000 (fine-tuned)

Usage

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

MODEL_ID = "saitejasathiraju/indian-legal-reranker"

INSTRUCTION = ("Given a legal query about Indian statutes and case law, retrieve the "
               "passage that directly and correctly answers the query")
SYSTEM_PREFIX = ('<|im_start|>system\nJudge whether the Document meets the requirements based on the '
                  'Query and the Instruct provided. Note that the answer can only be "yes" or "no".'
                  '<|im_end|>\n<|im_start|>user\n')
SUFFIX = "<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n"

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, padding_side="left")
model = AutoModelForCausalLM.from_pretrained(MODEL_ID, dtype=torch.bfloat16).to("cuda").eval()

token_true_id = tokenizer.convert_tokens_to_ids("yes")
token_false_id = tokenizer.convert_tokens_to_ids("no")

def score(query, doc):
    body = f"<Instruct>: {INSTRUCTION}\n<Query>: {query}\n<Document>: {doc}"
    ids = tokenizer.encode(SYSTEM_PREFIX + body + SUFFIX, add_special_tokens=False)
    inp = torch.tensor([ids], device=model.device)
    with torch.no_grad():
        logits = model(input_ids=inp, attention_mask=torch.ones_like(inp), logits_to_keep=1).logits[0, -1, :]
    pair = torch.stack([logits[token_false_id], logits[token_true_id]])
    return torch.nn.functional.log_softmax(pair, dim=0)[1].exp().item()

print(score(
    "What does Section 103 of BNS 2023 deal with?",
    "Section 103 of BNS 2023 - Punishment for murder.--(1) Whoever commits murder shall be "
    "punished with death or imprisonment for life, and shall also be liable to fine."
))  # -> ~1.0

For reranking multiple candidates, score each (query, doc) pair and sort by score descending.

Important caveats

• This is a reranker, not a citation-verification tool. It scores whether a retrieved passage answers a query — it does not know statute text it wasn't shown, and will not catch a retriever that missed the right passage entirely.
• Scoring requires reading the raw yes/no logits directly (as in the usage example above), not generic text generation. Serving this through a chat-style API (e.g., Ollama's /api/generate without raw-prompt mode) will silently degrade quality — see the chat_template behavior described in this repo's files before wrapping it in a different serving stack.
• Trained on English-language passages; BNS/BNSS/BSA content in other Indian languages was not part of the training set.
• Not legal advice. Outputs should not be relied upon without review by a qualified legal professional.

License

Apache 2.0, inherited from the base model.