DJLougen / granite4-tax

Granite 4 Tax

A fine-tuned version of IBM Granite 4.0 Tiny Preview specialized for U.S. tax law reasoning with IRC citation support.

Model Details

• Base model: ibm-granite/granite-4.0-tiny-preview
• Architecture: GraniteMoeHybridForCausalLM (Mamba + Attention hybrid, MoE with 62 experts, 6 active per token)
• Hidden size: 1536
• Layers: 40 (36 Mamba + 4 Attention)
• Context length: 131,072 tokens
• Precision: bfloat16

Training

Fine-tuned using Unsloth with LoRA (rank=16, alpha=16) on synthetic U.S. tax law Q&A data covering:

• Individual taxation
• Business entity taxation
• Estate and gift tax
• International tax (CFCs, GILTI, FDII)
• Tax procedure and compliance

LoRA target modules: q_proj, k_proj, v_proj, o_proj, in_proj, out_proj

The model was trained using SFT (Supervised Fine-Tuning) with structured reasoning in IRAC format (Issue, Rule, Application, Conclusion) and IRC citation grounding.

Usage

from transformers import AutoModelForCausalLM, AutoTokenizer

model_name = "DJLougen/granite4-tax"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype="auto", device_map="auto")

messages = [
    {"role": "system", "content": "You are a tax law assistant. Provide accurate analysis with IRC citations."},
    {"role": "user", "content": "What are the requirements for a corporation to elect S corporation status under IRC § 1362?"}
]

input_ids = tokenizer.apply_chat_template(messages, return_tensors="pt").to(model.device)
output = model.generate(input_ids, max_new_tokens=512)
print(tokenizer.decode(output[0][input_ids.shape[-1]:], skip_special_tokens=True))

• The models were fine-tuned and evaluated using a hybrid RAG pipeline with rule-based section forcing, code-computed tax calculations, disambiguation chunks for complex statutes, and an agentic self-correction loop. Evaluated on complex tax scenarios including SSTB phase-outs, passive loss exceptions, and nonqualified use proration.

Limitations

• Trained on synthetic data; not a substitute for professional tax advice
• Coverage is focused on U.S. federal tax law