kw5

Read this before loading

The input embedding and the output head are NOT tied, even though the config this run was launched with said they were. The tie broke when the trainer moved the model to the TPU. nn.Module._apply rebuilds a Parameter per entry when it cannot reuse the source storage, and CPU to XLA cannot, so the run trained two independent tensors.

If you rebuild this model with tied embeddings and call load_state_dict, both keys are written into one storage, the second overwrites the first, and your input embedding becomes the output head. Nothing raises. The shapes match. The model then scores roughly the unigram baseline (7.32 nats instead of 3.16) and generates whitespace. This cost the project a month and ~27 TPU-hours of false conclusions; the shipped config.json sets tie_embeddings: false, which is the truth, and modeling_kw5v2.py honours it.

Quick start (runs in Google Colab as-is)

No GPU required; a free CPU runtime works (a few seconds per line), and a T4 is faster. The repository ships its own modeling code, so there is nothing to clone.

!pip install -q huggingface_hub sentencepiece

import sys, torch, sentencepiece as spm
from huggingface_hub import snapshot_download

path = snapshot_download("regnant-io/kw5-149M")   # ~700 MB, cached
sys.path.insert(0, path)                    # this repo ships modeling_kw5v2.py
from modeling_kw5v2 import KW5V2ForCausalLM

device = "cuda" if torch.cuda.is_available() else "cpu"
model = KW5V2ForCausalLM.from_pretrained(path).to(device)
sp = spm.SentencePieceProcessor(model_file=f"{path}/tokenizer.model")

def swahili(prompt, max_new_tokens=60):
    ids = model.generate(
        sp.encode(prompt),                  # no , see below
        max_new_tokens=max_new_tokens,
        temperature=0.2, top_p=0.9, repetition_penalty=1.3,
    )
    return sp.decode(ids)

print(swahili("Tanzania ni nchi"))

Tanzania ni nchi ya Afrika Mashariki na moja kati ya mataifa yenye uchumi mkubwa duniani.
Tanzania inasifika kwa kuwa na maliasili nyingi, lakini pia ina madini mengi ambayo yanaifanya iwe miongoni mwa nchi zenye utajiri wa rasilimali za asili barani Afrika.
Kwa mujibu wa ripoti iliyotolewa hivi karibuni na Shirika la Utafiti wa Madini Duniani (

More prompts. This is a base model, so give it a prefix to continue rather than an instruction:

for p in ["Mji mkuu wa Tanzania ni", "Kiswahili ni lugha", "Elimu ni muhimu kwa sababu"]:
    print(swahili(p, 40))
    print()

Mji mkuu wa Tanzania ni Dodoma.

Mji huo ulianzishwa mwaka 1964 na wakoloni Waingereza kwa jina la "Tanganyika Territory" (sasa: Tanganyika) katika eneo la Ziwa Nyasa, ambalo sasa lina

Kiswahili ni lugha ya Kibantu nchini Angola inayozungumzwa na Wazambia. Mwaka wa 1983 idadi ya wasemaji wa Kiswahili imehesabiwa kuwa watu 20,500. Kufuatana na uainishaji wa lugha za

Elimu ni muhimu kwa sababu ya umuhimu wake katika maisha yetu. Ni lazima tufahamu kuwa elimu ndiyo msingi wa maendeleo, na hivyo basi tunahitaji kufahamu jinsi tunavyotumia maarifa hayo ili tuweze kufikia malengo tuliyojiwekea maishani mwetu.
Katika makala

Sampling is stochastic, so your output will differ. Decoding defaults live in generation_config.json and were chosen by measurement. See Decoding below.

Read those samples carefully: the capital of Tanzania is right, and "Kiswahili ni lugha ya Kibantu nchini Angola inayozungumzwa na Wazambia" is confidently wrong. That is what 149M parameters buys you: fluency, not knowledge.

Do not prepend . Base packing never prepended BOS: every document began with its own first token and ended with . A leading `` is a format this model has never seen.

A stock AutoModelForCausalLM will not load this correctly: the architecture has Canon layers, which no model in the Llama family has.

Does it follow instructions?

Partly, and it is worth being precise about which part. 12 instructions, each run twice with identical text, once wrapped in the chat format and once as plain continuation, using the same decoder settings:

The format is learned. The chat wrapper makes the model terminate, by emitting `` or a role token, about half the time. Plain continuation never terminates, not once in 12 tries. In v1 those ids were reserved and never trained, and prompting with them returned noise.

The register follows too. Same instruction, two formats:

"Eleza kazi ya mwalimu."          (explain a teacher's job)

chat  -> "Mwalimu ni mtu mwenye ujuzi mkubwa ambaye anaweza kusaidia
          wanafunzi katika masomo yao, kutoa ushauri na mwongozo..."
plain -> "- Tambua aina za maneno na misemo katika sentensi.
          - Andika orodha ya majina, alama na vishazi vya kila moja..."

The chat version answers. The plain version continues someone else's grammar worksheet.

**But constraints are ignored and facts are unreliable.