README.md

---
license: other
---
# 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`](https://huggingface.co/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

```python
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.