nia-100m-malawi-model

Nia Chat 110M

A 110M-parameter, English, instruction-tuned chat model, trained from scratch by Dev Intern AI — an AI startup based in Malawi — as part of the Nia model family.

Engineering: Lance Muyawa

Model details

• Architecture: decoder-only transformer, RoPE, grouped-query attention, SwiGLU FFN, RMSNorm, tied embeddings
• 10 layers, d_model 768, 12 attention heads, 3072 FFN dim, ~110M params
• Base pretrain: ~4B tokens (Wikipedia / OpenWebText / C4 / Gutenberg / TinyStories mix)
• Continued pretrain ("epoch 2"): factual-clean corpus (Wikipedia + Simple Wikipedia + Cosmopedia + Gutenberg) to reduce confabulation
• SFT: ChatML-formatted mix of Alpaca-style instructions, SQuAD 2.0 + TriviaQA grounded QA, identity, refusal/uncertainty, and conversational-filler examples
• This checkpoint: epoch-2 SFT, step 2024, training loss 1.6722

Usage

Quick start

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

REPO_ID = "devinternmw/nia-chat-110m"

tok = AutoTokenizer.from_pretrained(REPO_ID, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(REPO_ID, trust_remote_code=True)
model.eval()

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

system = "You are Nia, a helpful AI assistant built by Dev Intern AI in Malawi. Answer clearly and honestly, and say when you're not sure."
prompt = f"<|im_start|>system\n{system}\n<|im_end|>\n<|im_start|>user\nWhat is the capital of France?\n<|im_end|>\n<|im_start|>assistant\n"

ids = tok(prompt, return_tensors="pt").input_ids.to(device)
out = model.chat_generate(ids, max_new_tokens=150, temperature=0.65, eos_id=tok.convert_tokens_to_ids("<|im_end|>"))
print(tok.decode(out[0][ids.shape[1]:], skip_special_tokens=True))

trust_remote_code=True is required on both calls above — this is a custom architecture (RoPE + grouped-query attention + SwiGLU), not a stock HF model class.

Important: no model.generate()

This model's forward() is a plain full-sequence pass with no KV cache. Use the included model.chat_generate(...) method for sampling, not HF's standard model.generate(...) — the latter will not produce correct results with this architecture.

Prompt format (ChatML)

<|im_start|>system
{system prompt}
<|im_end|>
<|im_start|>user
{user message}
<|im_end|>
<|im_start|>assistant
{model replies here, then emits <|im_end|>}

Always include a system prompt — the model was trained with one on every example.

Generation parameters (chat_generate)

Performance note

No KV cache means every generated token reruns the full forward pass over the whole sequence so far. This is noticeably slower than a typical cached HF model, especially on CPU. A GPU helps significantly; for longer replies, expect it to be slower than you're used to regardless of hardware.

Standalone CLI script

A ready-to-run chat script (terminal + notebook safe) is included in this repo: chat_with_nia.py.

pip install -r requirements.txt
python chat_with_nia.py                                   # interactive
python chat_with_nia.py "What is the capital of France?"  # one-shot

Inside Jupyter/Colab, import the functions directly instead of running the file as a script:

from chat_with_nia import load_model, ask
tok, model, device = load_model()
print(ask(tok, model, device, "What is the capital of France?"))

Capabilities

• Follows ChatML-style instructions and stops at <|im_end|>
• Answers factual questions with improved (but still imperfect) accuracy after the factual-clean continued pretrain
• Short-form writing, explanations, simple Q&A

Limitations

• Still hallucinates, especially on less common facts
• No arithmetic or multi-step reasoning
• English only
• 110M parameters is small — knowledge retention and generation quality are well below larger models
• No KV-cache in generation (see Performance note above)

Credits

Built by Dev Intern AI, an AI startup based in Malawi 🇲🇼.

Engineering: Lance Muyawa

License

Apache 2.0 — see LICENSE.