NextToken-model-1

A 272.7M-parameter language model, pretrained from scratch on 11 Indic languages + English (Sangraha corpus), then fine-tuned on grounded multilingual QA for the Indian government-schemes / financial-banking domain (PM-KISAN, Ayushman Bharat, banking products, insurance, savings instruments, etc.).

Updated in place — this repo tracks the current best domain checkpoint, not a fixed snapshot; check back for updates as fine-tuning improves.

Architecture

Built from custom composable primitives, structurally equivalent to Qwen3 (confirmed by direct source comparison during HF conversion) and saved in that format for standard transformers loading:

Parameters	272.7M (69.7M embedding, tied with output head; 203.0M transformer)
Layers	16
Hidden size	1024
Attention	GQA — 16 query heads, 4 KV heads
Normalization	RMSNorm (pre-norm) + per-head QK-norm before RoPE
FFN	SwiGLU, intermediate size 3277
Positional encoding	RoPE (θ=10000)
Context length	2048
Tokenizer	sarvamai/sarvam-1, vocab 68,096

Training

Pretraining: Sangraha corpus, 11 Indic languages + English, ~5.7B tokens/phase across H100 and V100 GPUs (best validation loss checkpoint, step 23,500, loss 2.9318).

Domain SFT (v2, this checkpoint): full fine-tune (not LoRA) on 38,002 grounded QA examples (40,363 total across train/val/test) covering Indian government schemes and banking/financial products, generated across 9 languages (hi/kn/ta/te/mr/bn/gu/pa/en) from 846 scheme/product source documents (57 schemes/products, 1,445 source-document chunks). Prompt format:

Context: {context}
Question: {question}
Answer: {answer}

Loss masked to the Answer: span only. Trained with early stopping on validation loss (patience=3 evals with no improvement) rather than a fixed epoch count — v1 (3 fixed epochs on a smaller, thinner 4,728-example set) overfit badly (train loss kept falling while val loss rose). This run stopped at step 1400/3561 (< 1 epoch over the larger v2 set); the published weights are from step 1100, the actual best-val-loss point (1.886), not the final step.

Usage

from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("somasekhar-dev/NextToken-model-1")
tokenizer = AutoTokenizer.from_pretrained("somasekhar-dev/NextToken-model-1")

prompt = "Context: PM-KISAN provides Rs. 6,000 per year, paid in three installments of Rs. 2,000 each.\nQuestion: How much does PM-KISAN pay per year?\nAnswer:"
inputs = tokenizer(prompt, return_tensors="pt")
out = model.generate(**inputs, max_new_tokens=60, do_sample=False)
print(tokenizer.decode(out[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True))

Status and limitations (read before using for eval/comparison)

• v2 domain SFT — NOT a confirmed improvement over v1 on harder questions. Early stopping worked correctly (textbook val_loss curve: improves then plateaus, caught cleanly at the inflection point), and the training set is ~8x larger and more diverse than v1's. But a spot comparison against v1 on 6 held-out eligibility/edge-case questions (Rythu Bandhu scheme) showed a mixed result, not a clear win: v2 was more coherent on one question, but also flipped one answer to the wrong yes/no direction where v1 had at least gotten the direction right. Both checkpoints still hallucinate, self-contradict, or go off-topic on nuanced eligibility questions.
• Simple direct-fact questions (e.g. "how much does X pay") are answered more reliably than eligibility/edge-case questions requiring combining multiple conditions.
• Training data (both source documents and Q/A pairs) is entirely LLM-generated (grounded against related documents/passages for consistency), not scraped or independently fact-checked at scale.
• Training/dataset source: somasekhar-dev/nexttoken-model-1-dataset-sft (v2).
• Conversion from the original OLM training format to Qwen3ForCausalLM was verified numerically (100% argmax agreement, full round-trip) for the pretrained base this was fine-tuned from.