ai4ug/pearl-chat

Pearl Model (8.5M-causal-v2): Native Luganda Conversational Language Model

This is the Pearl Model (8.5M-causal-v2), a decoder-only GPT-2 style autoregressive language model trained from scratch on native Luganda prose, parallel corpora, and conversational dialogues. It contains 8,473,122 actual trainable parameters (~8.5M parameter class) and is designed to generate coherent Luganda text and respond accurately to conversational prompts, greetings, Ugandan factual questions, and English-to-Luganda translation queries.

This model was engineered entirely in pure JAX using the Flax NNX API as part of the technical workshop "Building Pearl-Chat: Engineering a Native Language Model from Scratch in Pure JAX" for PyCon Africa 2026.

Model overview

• Model name: Pearl Model (8.5M-causal-v2)
• Repository: kambale/pearl-chat
• Architecture: Decoder-only Transformer (GPT-2 style)
• Framework: Pure JAX and Flax NNX
• Target language: Luganda ('lg')
• Total parameters: 8,473,122 actual trainable parameters (~8.5M parameter class)
• Context length: 256 tokens
• Layers: 8 transformer decoder blocks
• Attention heads: 8 (head dimension 32)
• Embedding dimension: 256
• Feed-forward dimension: 1024 (4x expansion with GELU activation)
• Tokenizer: Byte-level Byte-Pair Encoding (BPE), vocabulary size 8,192
• Weight tying: Output projection tied to the 8,192 x 256 token embedding table (saving 2,097,152 parameters)
• Training curriculum: Two-stage curriculum (Stage 1 foundational Luganda pre-training + Stage 2 assistant-masked conversational instruction fine-tuning)

Intended use

The model is designed for:

• Conversational interaction and question answering in Luganda (greetings, identity, geography, education, and culture in Uganda).
• Bilingual English-to-Luganda sentence translation (Vvuunula: / Vvuunula mu Luganda:).
• Generative text continuation in Luganda.
• Research in low-resource African language modeling with pure functional JAX and explicit state management in Flax NNX.

Out-of-scope:

• High-stakes factual retrieval or medical and legal counsel.
• Highly specialized technical jargon absent from the training corpus.

Training details

Datasets and two-stage curriculum

The model is trained across a curated combination of four Hugging Face datasets and native conversational splits:

1. kambale/luganda-english-parallel-corpus: 25,006 unique contemporary Luganda-English parallel sentence pairs used in Stage 1 foundational pre-training and Stage 2 translation instruction tuning.
2. kambale/luganda-english-bible-corpus: 32,291 unique Luganda verses (with leading verse numbers and chapter headers stripped) used in Stage 1 foundational pre-training and BPE tokenizer fitting to instill broad Bantu verb morphology and noun-class concords.
3. wikimedia/wikipedia (20231101.lg split): Encyclopedic articles covering geography, history, and society in Uganda and Buganda.
4. Sunbird/salt (text-all subset, strictly filtered for Luganda and English columns): Parallel English-Luganda sentences and translation pairs.
5. Curated conversational dialogues: Native Luganda question-answer pairs and greetings partitioned into disjoint training, validation, and held-out benchmark splits prior to oversampling.

Infrastructure and hyperparameters

• Hardware: Compatible with 1x NVIDIA T4 / A100, Google Cloud TPU, or multi-core CPU.
• Batch size: 16
• Sequence length: 256 tokens
• Stage 1 (Foundational Pre-Training): 3,000 steps on Luganda prose and parallel corpora
• Stage 2 (Conversational SFT): 1,000 steps with assistant-only loss masking
• Optimizer: Optax AdamW (weight decay 0.01, global gradient norm clipping 1.0)
• Peak learning rate: 6e-4 with linear warmup and cosine decay to 1e-5
• Loss function: Assistant-masked softmax cross-entropy over discrete tokens

Evaluation metrics

• Final validation loss: 0.1818
• Final validation perplexity: 1.20
• Held-out ROUGE-L: 0.9792
• Held-out BLEU: 0.9405

How to use

Load the model using the native Pearl-Chat library or direct Flax NNX definitions:

from huggingface_hub import snapshot_download
from flax import nnx
from pearlchat.tokenizer import LugandaTokenizer
from pearlchat.model import PearlChatModel
from pearlchat.config import ModelConfig
from pearlchat.checkpointing import CheckpointManager
from pearlchat.generate import generate_text

# Download model repository
repo_dir = snapshot_download(repo_id="kambale/pearl-chat")

# Load tokenizer and 8.5M model configuration
tokenizer = LugandaTokenizer.load(repo_dir)
config = ModelConfig(
    vocab_size=tokenizer.vocab_size,
    context_length=256,
    embed_dim=256,
    num_heads=8,
    num_layers=8,
    feed_forward_dim=1024,
    dropout_rate=0.0,
)

# Initialize model and restore weights
rngs = nnx.Rngs(params=0, dropout=0)
model = PearlChatModel(config, rngs)
manager = CheckpointManager(repo_dir)
model, _ = manager.restore_latest_checkpoint(model)

# Generate response to a Luganda prompt
prompt = "Omuntu: Oli otya?
Omuyambi:"
response = generate_text(model, tokenizer, prompt=prompt, max_new_tokens=48, temperature=0.0)
print(f"Response: {response}")

Limitations and bias

• Corpus domain balance: While combining parallel, Bible, Wikipedia, and SALT corpora substantially expands Luganda morphological coverage, low-resource African languages still have smaller digital footprints than high-resource languages.
• Model capacity: At 8,473,122 parameters (~8.5M parameter class), the model balances expressive multi-layer capacity with fast single-GPU and laptop training for live educational workshops.

Ethical considerations

Generative models may reflect biases or domain artifacts present in training sources. Users should exercise critical review of generated outputs in public applications.

Citation

If you use this model or code in your research or educational work, please cite:

@misc{kambale2026pearlchat,
  author = {Wesley Kambale},
  title = {Pearl-Chat: Engineering a Native Language Model from Scratch in Pure JAX},
  year = {2026},
  publisher = {Hugging Face},
  howpublished = {https://huggingface.co/kambale/pearl-chat},
  note = {PyCon Africa 2026}
}

License

Apache 2.0