ViuAI/ViuAI-500M

🪶 Sarus-500M (Sarus-0.5B-Thinking)

A Lightweight Bilingual Cognitive Reasoning Model Built from Scratch by ViuAI

🌐 ViuAI Studio | 💬 Live Demo | 📖 Model Card | ⚡ Quickstart

---

📌 Introduction

Sarus-500M is a 502M-parameter compact causal language model developed by ViuAI. Named after the graceful and sharp Sarus Crane, the model is engineered from the ground up to achieve state-of-the-art reasoning density per parameter.

Unlike standard small language models that output direct shallow answers, Sarus-500M executes an internal Cognitive Reasoning Monologue ([THINK] ... [/THINK]) before formulating its final response. This enables the model to break down mathematical logic, algorithmic puzzles, factual queries, and natural conversational dialogue across English and Hindi (Hinglish).

---

🌟 Key Features

1. 🧠 Pure Cognitive Reasoning ([THINK] ... [/THINK]):
   • Generates step-by-step internal reasoning traces before answering, significantly improving factual accuracy, multi-step problem solving, and deductive reasoning.
2. 🇮🇳 Native Bilingual Fluency:
   • Pre-trained and fine-tuned on extensive English, Devanagari Hindi, and romanized Hinglish datasets with seamless code-switching capabilities.
3. ⚡ High-Throughput Modern Architecture:
   • Built with Grouped Query Attention (GQA), SwiGLU Activations, Rotary Position Embeddings (RoPE), and RMSNorm for low-latency inference on edge devices, consumer GPUs, and mobile NPUs.
4. 💬 Multi-Turn Conversation Ready:
   • Robust conversation context handling with strict end-of-turn separation tokens (<|user|>, <|assistant|>, <|endofturn|>).

---

📐 Model Specifications

Parameter	Specification	Details
Model Name	Sarus-500M	Sarus-0.5B-Thinking
Total Parameters	502,488,320 (~502M)	Non-embedding params: ~420M
Hidden Dimension ($d_{\text{model}}$)	1280	Model width
Layers ($n_{\text{layers}}$)	24	Transformer blocks
Query Heads ($n_{\text{heads}}$)	20	Attention heads
KV Heads ($n_{\text{kv_heads}}$)	4	Grouped Query Attention (GQA 5:1 ratio)
Intermediate Size	3456	SwiGLU projection dimension
Vocabulary Size	64,009	Custom BPE tokenizer with specialized thinking & turn tokens
Max Context Length	2048 Tokens	Rotary Positional Embeddings ($\theta = 10,000$)
Activation Function	SwiGLU	Gated Linear Unit
Normalization	RMSNorm	Root Mean Square Layer Normalization ($\epsilon = 10^{-6}$)

---

🚀 Quickstart & Inference

1. Installation

pip install torch transformers huggingface_hub

2. Standalone Python Inference

import torch
import torch.nn.functional as F
from huggingface_hub import hf_hub_download
from transformers import PreTrainedTokenizerFast

# 1. Download Architecture & SFT v8 Checkpoint
REPO_ID = "ViuAI/ViuAI-500M"
device = "cuda" if torch.cuda.is_available() else "cpu"

!wget -q https://huggingface.co/ViuAI/ViuAI-500M/resolve/main/code/config.py -O config.py
!wget -q https://huggingface.co/ViuAI/ViuAI-500M/resolve/main/code/model.py -O model.py

from config import ViuAIConfig
from model import ViuAI

# 2. Load Tokenizer & Model
tokenizer_path = hf_hub_download(repo_id=REPO_ID, filename="tokenizer/tokenizer.json")
tokenizer = PreTrainedTokenizerFast(
    tokenizer_file=tokenizer_path,
    pad_token="<pad>", bos_token="<bos>", eos_token="<eos>", unk_token="<unk>",
    additional_special_tokens=["<|user|>", "<|assistant|>", "<|endofturn|>", "[THINK]", "[/THINK]"]
)
eot_id = tokenizer.convert_tokens_to_ids("<|endofturn|>")

ckpt_path = hf_hub_download(repo_id=REPO_ID, filename="sft_checkpoints/sft_v8/sft_ckpt_final.pt")
checkpoint = torch.load(ckpt_path, map_location=device, weights_only=False)

raw_sd = checkpoint.get("model_state_dict", checkpoint.get("model", checkpoint))
state_dict = {k.replace("module._orig_mod.", "").replace("_orig_mod.", "").replace("module.", ""): v for k, v in raw_sd.items()}

config = ViuAIConfig(vocab_size=state_dict['tok_emb.weight'].shape[0], use_checkpoint=False)
model = ViuAI(config).to(device)
model.load_state_dict(state_dict, strict=False)
model.head.weight = model.tok_emb.weight
model.eval()

# 3. Generate Response with Deep Reasoning
prompt = "Explain why the sky is blue in simple terms."
prompt_formatted = f"<|user|>\n{prompt}<|endofturn|>\n<|assistant|>\n"

input_ids = tokenizer(prompt_formatted, add_special_tokens=False, return_tensors="pt").input_ids.to(device)
generated_ids = []

for _ in range(512):
    with torch.no_grad():
        with torch.autocast(device_type=device, dtype=torch.bfloat16 if device == "cuda" else torch.float32):
            logits = model(input_ids)[0][:, -1, :]
    
    # Temperature & Top-P Sampling
    logits = logits / 0.7
    probs = F.softmax(logits, dim=-1)
    idx_next = torch.multinomial(probs, num_samples=1)
    input_ids = torch.cat([input_ids, idx_next], dim=1)
    
    if idx_next.item() == eot_id:
        break
    generated_ids.append(idx_next.item())

print(tokenizer.decode(generated_ids, skip_special_tokens=False))

---

💬 Prompt Template & Chat Format

Sarus-500M adheres to the structured chat turn template:

<|user|>
{user_message}<|endofturn|>
<|assistant|>
[THINK]
{step_by_step_internal_reasoning_monologue}
[/THINK]
{final_comprehensive_response}<|endofturn|>

Example Generation:

Input:

<|user|>
A train travels 150 km in 3 hours. What is its speed in km/h and m/s?<|endofturn|>
<|assistant|>

Output:

[THINK]
1. The user is asking for the speed of a train given distance = 150 km and time = 3 hours.
2. Speed in km/h = Distance / Time = 150 / 3 = 50 km/h.
3. To convert km/h to m/s: multiply by 5/18 or (50 * 1000) / 3600.
4. 50 * (5/18) = 250 / 18 = 13.888... m/s ≈ 13.89 m/s.
5. Formulate a clean, structured response showing both answers clearly.
[/THINK]
Here is the step-by-step calculation:

1. **Speed in km/h**:
   $$\text{Speed} = \frac{\text{Distance}}{\text{Time}} = \frac{150\text{ km}}{3\text{ hours}} = 50\text{ km/h}$$

2. **Speed in m/s**:
   $$\text{Speed in m/s} = 50 \times \frac{5}{18} \approx 13.89\text{ m/s}$$

**Final Answer:** The train travels at **50 km/h** (or approximately **13.89 m/s**).<|endofturn|>

---

📈 Training Details

• Pre-training: Trained from scratch on high-quality multilingual web corpus, code repositories, and mathematical formulations.
• Supervised Fine-Tuning (SFT v8):
  • Dataset: Curated 100,000+ instruction-reasoning pairs spanning math, logic, Python code, world knowledge, and natural Hinglish chit-chat.
  • Hardware: Single NVIDIA H200 (150GB VRAM) & Multi-GPU DDP clusters.
  • Precision: Mixed Precision (bfloat16 with FlashAttention & torch.compile).
  • Optimizer: AdamW ($\beta_1=0.9, \beta_2=0.95$, Weight Decay=0.01).
  • Learning Rate: Cosine decay with warmup ($2\times 10^{-5} \rightarrow 2\times 10^{-6}$).

---

📄 License & Attribution

Sarus-500M is released under the Apache 2.0 License, permitting commercial and non-commercial use, modification, and distribution.

@misc{sarus2026viuai,
  title={Sarus-500M: High-Density Cognitive Reasoning at 0.5B Scale},
  author={ViuAI Team},
  year={2026},
  publisher={Hugging Face},
  howpublished={\url{https://huggingface.co/ViuAI/ViuAI-500M}}
}

Crafted with passion by the ViuAI Team.