File size: 4,245 Bytes
29d4e0c 284e40e 29d4e0c 284e40e 5c8f2b2 284e40e 29d4e0c 284e40e 29d4e0c 284e40e 5c8f2b2 284e40e 5c8f2b2 29d4e0c 284e40e 29d4e0c 284e40e 29d4e0c 5c8f2b2 284e40e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 | ---
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.
|