Views
No views yet
[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).[THINK] ... [/THINK]):
<|user|>, <|assistant|>, <|endofturn|>).| 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}$) |
pip install torch transformers huggingface_hub1import torch
2import torch.nn.functional as F
3from huggingface_hub import hf_hub_download
4from transformers import PreTrainedTokenizerFast
5
6# 1. Download Architecture & SFT v8 Checkpoint
7REPO_ID = "ViuAI/ViuAI-500M"
8device = "cuda" if torch.cuda.is_available() else "cpu"
9
10!wget -q https://huggingface.co/ViuAI/ViuAI-500M/resolve/main/code/config.py -O config.py
11!wget -q https://huggingface.co/ViuAI/ViuAI-500M/resolve/main/code/model.py -O model.py
12
13from config import ViuAIConfig
14from model import ViuAI
15
16# 2. Load Tokenizer & Model
17tokenizer_path = hf_hub_download(repo_id=REPO_ID, filename="tokenizer/tokenizer.json")
18tokenizer = PreTrainedTokenizerFast(
19 tokenizer_file=tokenizer_path,
20 pad_token="<pad>", bos_token="<bos>", eos_token="<eos>", unk_token="<unk>",
21 additional_special_tokens=["<|user|>", "<|assistant|>", "<|endofturn|>", "[THINK]", "[/THINK]"]
22)
23eot_id = tokenizer.convert_tokens_to_ids("<|endofturn|>")
24
25ckpt_path = hf_hub_download(repo_id=REPO_ID, filename="sft_checkpoints/sft_v8/sft_ckpt_final.pt")
26checkpoint = torch.load(ckpt_path, map_location=device, weights_only=False)
27
28raw_sd = checkpoint.get("model_state_dict", checkpoint.get("model", checkpoint))
29state_dict = {k.replace("module._orig_mod.", "").replace("_orig_mod.", "").replace("module.", ""): v for k, v in raw_sd.items()}
30
31config = ViuAIConfig(vocab_size=state_dict['tok_emb.weight'].shape[0], use_checkpoint=False)
32model = ViuAI(config).to(device)
33model.load_state_dict(state_dict, strict=False)
34model.head.weight = model.tok_emb.weight
35model.eval()
36
37# 3. Generate Response with Deep Reasoning
38prompt = "Explain why the sky is blue in simple terms."
39prompt_formatted = f"<|user|>\n{prompt}<|endofturn|>\n<|assistant|>\n"
40
41input_ids = tokenizer(prompt_formatted, add_special_tokens=False, return_tensors="pt").input_ids.to(device)
42generated_ids = []
43
44for _ in range(512):
45 with torch.no_grad():
46 with torch.autocast(device_type=device, dtype=torch.bfloat16 if device == "cuda" else torch.float32):
47 logits = model(input_ids)[0][:, -1, :]
48
49 # Temperature & Top-P Sampling
50 logits = logits / 0.7
51 probs = F.softmax(logits, dim=-1)
52 idx_next = torch.multinomial(probs, num_samples=1)
53 input_ids = torch.cat([input_ids, idx_next], dim=1)
54
55 if idx_next.item() == eot_id:
56 break
57 generated_ids.append(idx_next.item())
58
59print(tokenizer.decode(generated_ids, skip_special_tokens=False))<|user|>
{user_message}<|endofturn|>
<|assistant|>
[THINK]
{step_by_step_internal_reasoning_monologue}
[/THINK]
{final_comprehensive_response}<|endofturn|><|user|>
A train travels 150 km in 3 hours. What is its speed in km/h and m/s?<|endofturn|>
<|assistant|>[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|>bfloat16 with FlashAttention & torch.compile).1@misc{sarus2026viuai,
2 title={Sarus-500M: High-Density Cognitive Reasoning at 0.5B Scale},
3 author={ViuAI Team},
4 year={2026},
5 publisher={Hugging Face},
6 howpublished={\url{https://huggingface.co/ViuAI/ViuAI-500M}}
7}