taxbrainai-gemma-4-e4b

The advisor-synthesis model behind TaxBrainAI — Gemma-4-E4B fine-tuned to turn a deterministic tax engine's output and retrieved statute into a grounded, cited, bilingual answer. It never authors a number: every figure it writes arrives in its prompt, computed by the engine.

Most "AI tax" tools let a language model do the arithmetic and hope. TaxBrainAI inverts that: a property-tested Decimal rules engine (taxbrainai-compute) computes, an 8-stage hybrid + graph retrieval pipeline grounds, and this model synthesizes — it receives the engine's breakdown and the retrieved passages in its prompt, and its entire trained job is faithful presentation: transcribe the figures exactly, explain the mechanism, cite the sections, answer in the user's language (English or Hindi), refuse what it cannot ground. Correctness of amounts is a guarantee of the engine, not a hope about the model.

Repo contents — the full post-training staircase

Path	What it is
/ (root)	The ship: merged bf16 weights (base + SFT LoRA) — ready for transformers / vLLM
adapters/sft-lora/	The SFT LoRA adapter alone (r=32, α=64) — the checkpoint that ships
adapters/dpo-lora/	DPO adapter — explored branch off the SFT checkpoint (did not ship; see below)
adapters/grpo-lora/	GRPO adapter — explored branch off SFT (did not ship; see below)
LICENSE · NOTICE	Apache-2.0 + base-model attribution

All three fine-tunes are included so the head-to-head comparison below is fully reproducible.

The I/O contract (what the fine-tune actually does)

The backend assembles one prompt per turn — question + router signals (including the deterministic engine's result, when the question is computational) + retrieved statute context:

CURRENT QUESTION: My salary is Rs.14,18,500 with Rs.1,35,000 invested under Section 80C.
What is my tax in the old regime?

SIGNALS:
Computation: recommended regime = old; total tax = Rs.1,89,850
Engine breakdown:
- Gross total income: Rs.13,68,500
- Less: Chapter VI-A deductions: -Rs.1,35,000
- Total income: Rs.12,33,500
- Tax on normal income: Rs.1,82,550
- Health & education cess (4%): Rs.7,302
- Total tax: Rs.1,89,850

CONTEXT: Retrieved passages:
- [IT Dept Tax Tutorials] ... deduction under section 80C ... [Income-tax Act, Section 16(ia)] ...

The trained answer transcribes the engine's table verbatim — headline figure, then the breakdown rows exactly as given — explains the mechanism (here: the 80C deduction reducing total income under the old regime), cites the sections, and matches the question's language (a Hindi question gets the same treatment in Hindi; 26% of training rows are Hindi). Training penalized bare numbers without working, hallucinated citations, and any figure that differs from the engine's — the faithfulness axis below is literally "did the answer carry the engine's numbers."

Routing, tool dispatch, and safety screening live in the backend graph, not in this model's contract. Standalone, the model still behaves as trained — synthesize-from-context or refuse — but it is designed to be fed by the engine + retrieval for grounded operation.

The models in this system — and why

Role	Model	Why
Synthesis (this release)	Gemma-4-E4B, fine-tuned	Small enough to self-host on one GPU scale-to-zero; Apache-2.0; strong bilingual (EN/HI) base; cheap routine tier
Base	unsloth/gemma-4-E4B-it @ 0d5a7f9b	Community bf16 mirror of Google DeepMind's Gemma 4 E4B instruct — same weights, loader-compatible; revision pinned for reproducibility
Teacher (data gen)	a frontier LLM (vendor-neutral by policy; identity pinned in the project's provenance snapshot)	Rejection-sampling distillation: teacher drafts, the deterministic engine verifies, only oracle-correct answers enter training
Safety screen (live tier)	Llama Guard 4	Jailbreak/unsafe screening by a dedicated classifier, never the answering model
Retrieval embeddings	BGE-M3	Dense vectors over the ~18K-chunk statute corpus
Escalation ladder	larger hosted frontier models	Low-confidence turns escalate; the product metric is escalation rate ↓ — the fine-tune earning routine traffic, not "beating" frontier models

Training

Data (all engine-grounded, all synthetic, decontaminated). Every prompt carries the real retrieved statute context and, for computational rows, the real engine result — train-time inputs match serve-time inputs. Generated via teacher distillation with the engine as a hard filter; decontaminated against the held-out eval at train time (exact-question + 13-gram overlap).

Stage	Rows	Composition
SFT	1,269 rows · 2 epochs (160 steps, effective batch 16)	573 compute · 446 conceptual · 129 notice · 121 planning · 26% Hindi
DPO	1,200 preference pairs	chosen/rejected on synthesis quality (grounding, citation, faithfulness)
GRPO	900 prompts	gold engine result pre-baked per row; reward = taxbrain_synthesis_reward

Recipe. QLoRA SFT on the pinned base: r=32, α=64, dropout 0, 4-bit base, 4,096 ctx, completion-only loss masking (the trainer hard-aborts if turn markers can't be detected, so prompt tokens can never silently leak into the loss). Run provenance: trained on a single H100 (Modal), loss 0.28 → 0.07 over 160 steps; DPO and GRPO branched from this same checkpoint, and the head-to-head eval scored all arms on the same volume artifacts.

What was explored beyond SFT — and why it didn't ship. DPO and GRPO (verifiable reward scored by the deterministic engine + statute graph — no learned reward model) were both trained from the SFT checkpoint and evaluated head-to-head:

Arm	Composite score (engine-oracle, 366-row test split)
base (Gemma-4-E4B-it)	0.324
SFT (ships — root of this repo)	0.895
DPO (from SFT — adapters/dpo-lora/)	0.892
GRPO (from SFT — adapters/grpo-lora/)	0.899

Base → SFT is the entire win. DPO and GRPO are statistically comparable to SFT — measured honestly, not an improvement. The negative result is reported instead of hidden.

Why SFT ships (the decision rationale)

1. The rule was pre-registered, then obeyed. Before any RL spend, the promotion rule was fixed: a challenger ships only if its head-to-head win-rate CI clears 0.5 and its reward-delta CI clears zero on the held-out split. GRPO's nominal +0.004 over SFT met neither — promoting it would be noise-mining, not measurement.
2. The task was designed to leave RL little headroom. GRPO amplifies capability that already exists in the policy's sampling distribution — it pays when pass@k far exceeds pass@1 (open-ended math, code). Here the architecture deliberately removes that gap: the engine computes, retrieval grounds, and the model synthesizes.
3. SFT had already distilled the reward. The SFT data was oracle-filtered — every training answer verified by the same engine reward GRPO later optimized. That makes SFT effectively offline RL on the same signal; what remained is dominated by reward noise (ambiguous citation choices, judgment-call refusal boundaries), and RL against a noisy, flat reward buys variance, not signal.
4. Ties break toward simplicity. At identical measured quality, the earliest checkpoint wins: fewer training stages to defend, cleaner provenance, one artifact to serve, reproduce, and audit.

RL earns its place when a reward axis with genuine exploration headroom appears — retrieval-query rewriting scored by gold-section recall, multi-turn clarification scored by downstream compute correctness, escalation-rate reduction. The verifiable-reward + paired-eval infrastructure built here is reusable the day that axis exists.

Evaluation

• Held-out set: 406-row eval bank (126 compute · 145 conceptual · 41 planning · 94 notice · 28% Hindi), with a 40-row probe reserved for go/no-go gating; the remaining 366 rows are the judged test split.
• Scoring: objective engine-oracle composite (engine-figure faithfulness, graph-validated citations, format, refusal correctness, language match) with paired bootstrap CIs; the pre-registered headline comparison is CI-judged, secondary axes Holm-corrected. No LLM judge in the loop — the numbers are verifiable, not vibes.
• Selected axes (SFT, test split): citation F1 0.76 · grounding 0.87 · faithfulness 0.99 (engine-number recall, n=115 compute rows) · format validity 0.99 · Hindi 0.87 · base→SFT win-rate 1.000 (CI [1.00, 1.00]).
• System-level gates (the deployed agent around this model, on its golden harness): compute exactness 100% (architecture-enforced — the engine authors every figure), notice classification 100%, jailbreak resistance 100%, PII leakage 0.

Quickstart

vLLM (how production serves it — bf16, scale-to-zero):

vllm serve harshilprojects/taxbrainai-gemma-4-e4b \
  --served-model-name gemma-4-e4b --dtype bfloat16 --max-model-len 4096

Transformers:

from transformers import AutoModelForCausalLM, AutoTokenizer

repo = "harshilprojects/taxbrainai-gemma-4-e4b"
tok = AutoTokenizer.from_pretrained(repo)
model = AutoModelForCausalLM.from_pretrained(repo, torch_dtype="bfloat16", device_map="auto")

prompt = """CURRENT QUESTION: What's the tax on a ₹15,00,000 salary under the new regime?

SIGNALS:
Computation: recommended regime = new; total tax = Rs.97,500
Engine breakdown:
- Gross total income: Rs.14,25,000
- Total income: Rs.14,25,000
- Tax on normal income: Rs.93,750
- Health & education cess (4%): Rs.3,750
- Total tax: Rs.97,500

CONTEXT: Retrieved passages:
- [Income-tax Act, Section 115BAC] ..."""

messages = [{"role": "user", "content": prompt}]
inputs = tok.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt").to(model.device)
print(tok.decode(model.generate(inputs, max_new_tokens=512)[0], skip_special_tokens=True))

Loading an explored adapter (e.g. to reproduce the head-to-head):

from peft import PeftModel
base = AutoModelForCausalLM.from_pretrained("unsloth/gemma-4-E4B-it", torch_dtype="bfloat16", device_map="auto")
model = PeftModel.from_pretrained(base, "harshilprojects/taxbrainai-gemma-4-e4b", subfolder="adapters/grpo-lora")

Intended use & limitations

• Scope: Indian personal income-tax, AY 2026-27, English + Hindi. Educational / informational — not professional tax advice; verify with a qualified professional before filing.
• It is a component. Figures should arrive in the prompt from the engine. Without the engine + retrieval feeding it, treat any free-text figure as unverified.
• Out of scope → refusal. Other jurisdictions/years, investment-product picks, and tax-evasion requests are trained refusals, in both languages.
• Frozen to its statute snapshot; does not know post-training amendments.

Safety

Refuses evasion and product-recommendation requests (English + Hindi, including keyword-free paraphrases — trained and evaluated). In the deployed system, PII (PAN, Verhoeff-validated Aadhaar, GSTIN, and more) is redacted before the model sees it, a dedicated safety classifier screens inputs, and every answer passes a fail-closed citation verifier with a hash-chained audit trail.

Future

GGUF quantizations, the training-data card, and post-amendment refreshes may follow. FP8 serving was benchmarked but the production endpoint serves bf16 (documented as-is).

License & attribution

Apache-2.0. The base is unsloth/gemma-4-E4B-it (revision 0d5a7f9b) — a community bf16 mirror of Google DeepMind's Gemma 4 E4B instruct, released under Apache-2.0. This derivative retains attribution to both (see NOTICE).

Contact

Questions, issues, or access to the live agent: harshilprojects3@gmail.com · taxbrainai.harshilprojects.com