shivamfet/gemma-2-2b-grounded-qa

Gemma 2 2B-it — Grounded Extractive QA (QLoRA)

A QLoRA adapter on top of google/gemma-2-2b-it that answers questions using only a passage it is shown. It is tuned for grounded, extractive question answering over legal and financial prose: given a passage and a question, it responds with a short answer stated in — or directly inferable from — that passage, and does not narrate the source ("the passage states…").

This is a LoRA adapter (~83 MB), not a merged model. Load it on top of the base at inference time.

Intended use

• Reading-comprehension / grounded QA where the answer must come from a supplied passage.
• Domains it saw in training: SEC filings, U.S. case law, and general web-education text.

Not for open-domain question answering, retrieval, or use without a passage — it is trained to answer from the given context only.

Prompt format

Standard Gemma 2 chat template, with the passage and question in the user turn:

<start_of_turn>user
{passage}

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

Usage

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel

BASE = "google/gemma-2-2b-it"
ADAPTER = "shivamfet/gemma-2-2b-grounded-qa"

tok = AutoTokenizer.from_pretrained(BASE)
model = AutoModelForCausalLM.from_pretrained(
    BASE, torch_dtype=torch.bfloat16, attn_implementation="eager", device_map="cuda"
)
model = PeftModel.from_pretrained(model, ADAPTER)
model.eval()

passage = ("The Court held that the warrantless search of the vehicle violated the "
           "Fourth Amendment because no exigent circumstances existed. The evidence "
           "obtained was therefore suppressed.")
question = "Why was the evidence suppressed?"

msgs = [{"role": "user", "content": f"{passage}\n\nQuestion: {question}"}]
ids = tok.apply_chat_template(msgs, add_generation_prompt=True, return_tensors="pt").to("cuda")
out = model.generate(ids, max_new_tokens=256, do_sample=False,
                     eos_token_id=tok.convert_tokens_to_ids("<end_of_turn>"))
print(tok.decode(out[0][ids.shape[1]:], skip_special_tokens=True).strip())
# -> The evidence was suppressed because the warrantless search of the vehicle
#    violated the Fourth Amendment, as no exigent circumstances existed.

Training data

13,774 grounded question/answer pairs, teacher-generated and LLM-judged for support/correctness/answerability, deduplicated. Every pair is answerable from its passage alone.

Mix	Share
SEC filings	40.1%
U.S. case law	40.1%
FineWeb-Edu (web)	19.7%
— easy / medium / hard	42% / 40% / 19%

Passages were sampled one-per-document for diversity; pairs were generated with a small teacher model and filtered by a stronger judge (support + correctness + answerability + well-formedness), then deduplicated on question embeddings.

Training procedure

Re-tokenized through Gemma's 256k SentencePiece tokenizer (the pairs originated in a project with a 16k-vocab tokenizer) using the chat template above, with the loss masked to the answer span only.

Method	QLoRA — 4-bit NF4 (double-quant), LoRA r=16, α=32, dropout=0.05
Target modules	q, k, v, o, gate, up, down proj
Epochs	1
Effective batch	32 (micro-batch 16 × grad-accum 2)
LR / schedule	2e-4, cosine, 3% warmup, paged AdamW-8bit
Tokens / epoch	8.25M (744k supervised) · mean 599 tok/row
Hardware	1× H100, ~17 min, ~8,050 tok/s
Final train loss	0.42

Limitations

• Answers are only as reliable as the supplied passage; it will not correct or supplement it with outside knowledge (by design).
• Tuned on legal/financial/edu prose — other domains and languages are out of distribution.
• Single epoch of light SFT on an already-instruction-tuned base: it nudges style and grounding, it does not add new knowledge.

Framework versions

• PEFT 0.13.2 · Transformers 4.46.2 · PyTorch 2.5.1