NextToken-model-2 — SAM (Small Action Model)

A 272.7M-parameter Hindi-English model fine-tuned into a small action model for banking/NBFC customer service: given a user utterance (English, Hindi, Hinglish, or Hindi-English code-mixed) and a system prompt listing the available tools, it decides whether one applies and extracts the arguments.

This checkpoint (sep18_round2) replaces the earlier gc_v2 checkpoint and is a substantial redesign, not an incremental update — see What changed below. Updated in place — this repository tracks the current best checkpoint; check back for updates.

What it does

System: You are a banking assistant for an Indian NBFC. Today's date is 2026-09-18 (Friday).
Given the user's message, decide whether one of the tools below applies. If it does, call it:
reply with a short natural-language line if appropriate, followed by a JSON object
{"tool": "<name>", "arguments": {...}} on its own line. Only use arguments that are stated in
the conversation or can be computed from today's date -- never invent a value. If no tool
applies, reply in natural language only, or with {"tool": null} if a bare JSON response is expected.

Available tools:
{"name": "get_recent_transactions", "description": "Last N transactions", ...}
{"name": "set_autopay", "description": "Enable/disable automatic EMI deduction", ...}
...

User: home loan ki abhi chinta mat karo -- bas recent transactions dikha do.
Assistant: {"tool": "get_recent_transactions", "arguments": {}}

Unlike the previous checkpoint, the tool schema is shown in the system prompt every time (JSON-Schema-style, one line per tool) — the model is schema-conditioned, not relying on having memorized a fixed 41-tool list from training alone. It also handles multi-turn slot-filling (asking a clarifying question when a required argument is missing) and is trained to suppress hypothetical/past-tense/third-person framings.

What changed vs. the previous checkpoint (gc_v2)

Deliberately narrowed scope for reliability: 5 tools + no-tool (down from 41 — the full 41-tool schema alone is 3,359 tokens, more than this model's useful context budget can absorb well). Alongside the narrowing:

	Previous	This checkpoint
Prompt format	Bare User:/Assistant:, tools never shown	System prompt with tool schemas + reference date, shown every example
Argument values	Teacher free-generated utterance + values together	Canonical values generated first; teacher only rephrases around them
Negative examples	Single literal-string target {"tool": null}	Real natural-language replies, 6 distinct negative types
Train/test split	Row-level random	Stratified by template ID — tests unseen phrasing, not just unseen numbers
Training visibility	Single before/after loss number	Per-epoch train+val loss, best-checkpoint auto-save, automatic overfitting diagnosis

50-test head-to-head result (same 50 prompts, each model in its own native format):

Metric	Previous (41 tools)	This checkpoint (5 tools)
JSON validity	98.0%	100.0%
Action-correct (call vs. no-call)	68.0%	86.0%
Tool-correct	85.7%	91.4%
Args-exact	68.6%	91.4%

Architecture

Same base architecture as somasekhar-dev/NextToken-model-1:

Parameters	272.7M
Attention	GQA — 16 query heads, 4 KV heads
Normalization	RMSNorm + per-head QK-norm
FFN	SwiGLU
Context length	2048
Tokenizer	sarvamai/sarvam-1, vocab 68,096

Usage

Tokenize each message segment separately and concatenate — do not re-tokenize a merged prompt string. A real train/inference tokenization mismatch was found and fixed earlier in this project (merging the string before tokenizing can shift the last context token vs. what the model was actually trained on); the example below replicates the verified-correct convention and was checked against a real training example before publishing:

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

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

system_prompt = (
    "You are a banking assistant for an Indian NBFC. Today's date is 2026-09-18 (Friday).\n"
    "Given the user's message, decide whether one of the tools below applies. If it does, "
    "call it: reply with a short natural-language line if appropriate, followed by a JSON "
    "object {\"tool\": \"<name>\", \"arguments\": {...}} on its own line. Only use arguments "
    "that are stated in the conversation or can be computed from today's date -- never invent "
    "a value. If no tool applies, reply in natural language only, or with {\"tool\": null} if "
    "a bare JSON response is expected.\n\n"
    "Available tools:\n"
    '{"name": "get_recent_transactions", "description": "Last N transactions", '
    '"parameters": {"type": "object", "properties": {"count": {"type": "integer"}}, "required": []}}'
)
user_utterance = "home loan ki abhi chinta mat karo -- bas recent transactions dikha do."

segments = [f"System: {system_prompt}\n", f"User: {user_utterance}\n", "Assistant: "]
ids = []
for seg in segments:
    ids.extend(tokenizer.encode(seg, add_special_tokens=False))
ids_t = torch.tensor([ids])

out = model.generate(ids_t, max_new_tokens=60, do_sample=False,
                      eos_token_id=tokenizer.eos_token_id, pad_token_id=tokenizer.pad_token_id)
print(tokenizer.decode(out[0][ids_t.shape[1]:], skip_special_tokens=True))
# -> {"tool": "get_recent_transactions", "arguments": {}}

Training

Full fine-tune (not LoRA) off the Hindi-English pretrained base. ~4,914 examples (kashyap/task-1/dataset.jsonl) covering 5 tools + no-tool, balanced at ~16.7% each of 6 classes, split by template ID (not by row) so held-out examples test generalization to unseen phrasing. Canonical argument values generated first (entity_gen.py), with the teacher model only rephrasing natural language around them — the teacher can never invent an argument value, which was a real failure mode in the previous checkpoint.

Status and limitations

Being upfront about what's actually verified vs. still open:

• Verified via a same-prompt 50-example head-to-head against the previous checkpoint (see table above) — not yet validated against a larger, independently-generated test set the way the previous checkpoint's 200-scenario comparison was.
• Narrowed to 5 tools deliberately; does not cover the other ~36 tools the previous (worse-performing) checkpoint attempted.
• The system-prompt-with-schema format means correctness depends on the exact tool-schema JSON given in-context matching training formatting closely — schemas very different in style from training are unverified.
• Loss curves alone were actively misleading during this project's own development (a "good" loss looked wrong on real inference; a rising-val-loss curve looked like overfitting but real testing showed the later checkpoint was actually better) — always validate against real generated output, not the reported loss, before trusting a checkpoint.
• Any state-changing action should be confirmed before executing, not trusted blindly from a single model output.
• Conversion to this Qwen3ForCausalLM checkpoint was verified numerically (100% argmax agreement vs. the original training-format checkpoint on random inputs, plus a full save/reload round-trip, plus an independent check here against one real training example).