DexterP29/gemma-2-2b-it-legal-qa

gemma-2-2b-it-legal-qa

google/gemma-2-2b-it fine-tuned with LoRA for grounded question answering over legal and financial passages: answer only from the supplied context, and reply with the exact sentence "The provided context does not state this." when the context does not contain the answer.

The merged model is at the repo root; the LoRA adapter alone is in adapter/.

Results (500 held-out examples)

F1 (SQuAD-style token overlap with the reference answer) and refusals come from greedy generation on every held-out question. "Any wording" counts answers that open by saying the context does not contain the answer. Only 14 held-out questions are unanswerable, so refusal rates are coarse.

Held-out loss is lowest after epoch 1, while generated answers are marginally better after epoch 2; the F1 gap between the two epochs is within bootstrap noise (95% CI spans zero).

Prompt format

Gemma 2 has no system role; the instruction goes at the top of the user turn:

<bos><start_of_turn>user
You are a legal and financial assistant. Answer the question using ONLY the provided context.

<context>
{context}
</context>

Question: {question}<end_of_turn>
<start_of_turn>model

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

tok = AutoTokenizer.from_pretrained("DexterP29/gemma-2-2b-it-legal-qa")
model = AutoModelForCausalLM.from_pretrained("DexterP29/gemma-2-2b-it-legal-qa", torch_dtype=torch.bfloat16,
                                             attn_implementation="eager", device_map="auto")
system = "You are a legal and financial assistant. Answer the question using ONLY the provided context."
user = f"{system}\n\n<context>\n{context}\n</context>\n\nQuestion: {question}"
ids = tok.apply_chat_template([{"role": "user", "content": user}],
                              add_generation_prompt=True, return_tensors="pt").to(model.device)
out = model.generate(ids, max_new_tokens=128, do_sample=False)
print(tok.decode(out[0, ids.shape[1]:], skip_special_tokens=True))

Training

• —Data: 12,000 grounded Q&A pairs generated by gpt-5-mini from legal/financial passages (case law, SEC filings, some FineWeb-Edu), filtered by verbatim-evidence checks, an LLM judge, exact + MinHash dedup and CaseHOLD/LexGLUE decontamination. About 5% are unanswerable.
• —Tokens: 6.39M per epoch; loss on answer tokens only (357k per epoch).
• —LoRA: r=16, alpha=32, dropout=0.05, on qproj, kproj, vproj, oproj, gateproj, upproj, down_proj (20,766,720 trainable parameters).
• —Optimization: 2 epochs, global batch 32, AdamW, LR 0.0002 cosine to 2e-05, 20 warmup steps, bf16, eager attention.
• —Hardware: 4x NVIDIA H100 on Modal, 4.7 min of training.
• —Merge: adapter merged in fp32 and stored in bf16; held-out answer loss 0.33248 unmerged vs 0.33241 merged.

Limitations

Trained for extraction from a supplied passage, not for closed-book legal knowledge. Answers can still be wrong or incomplete; do not rely on them for legal or financial decisions.

Gemma is provided under and subject to the Gemma Terms of Use found at ai.google.dev/gemma/terms.