Upload Onit Keyboard LM (run7, tok_v2, 40M params)

.gitattributes

@@ -0,0 +1 @@
+*.pt filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,136 @@
+---
+language:
+  - en
+  - fr
+license: apache-2.0
+tags:
+  - keyboard
+  - language-model
+  - mobile
+  - ios
+  - coreml
+  - bilingual
+library_name: pytorch
+pipeline_tag: text-generation
+---
+
+# Onit Keyboard LM
+
+A **41M parameter** bilingual (English + French) language model designed for **mobile keyboard prediction** on iOS.
+
+## Model Description
+
+Onit Keyboard LM is a compact causal language model optimized for next-word prediction in a mobile keyboard context. It supports both English and French, including code-switching between the two languages.
+
+### Architecture
+
+| Component | Value |
+|-----------|-------|
+| Type | Causal LM (decoder-only) |
+| Parameters | ~41M |
+| Vocabulary | 16,384 BPE tokens |
+| Embedding dim | 512 |
+| Layers | 10 |
+| Attention heads | 8 |
+| FFN dim | 1408 (SwiGLU) |
+| Max sequence length | 256 |
+| Positional encoding | RoPE |
+| Normalization | RMSNorm + QK-Norm |
+| Embeddings | Tied (input = output) |
+
+### Key Design Choices
+
+- **SwiGLU FFN** for better parameter efficiency at small scale
+- **QK-Norm** for stable training without careful LR tuning
+- **RoPE** for length generalization
+- **Tied embeddings** to reduce parameter count (critical for mobile)
+- **BPE tokenizer** (16K vocab) trained on the bilingual data mix
+
+## Training
+
+### Dataset (Phase 2)
+
+The model was trained on a diverse bilingual mix:
+
+| Source | Language | Share |
+|--------|----------|-------|
+| OpenSubtitles | FR + EN | ~40% |
+| Wikipedia | FR + EN | ~30% |
+| C4 (web) | FR + EN | ~30% |
+
+Total: ~13.6M sentences, ~2.7 GB of clean text.
+
+### Hyperparameters
+
+| Parameter | Value |
+|-----------|-------|
+| Training steps | 30,000 |
+| Effective batch size | 64 (32 x 2 grad accum) |
+| Learning rate | 6e-5 (cosine decay) |
+| Warmup steps | 1,000 |
+| Precision | bf16 mixed |
+| Optimizer | AdamW |
+
+### Results
+
+| Metric | Value |
+|--------|-------|
+| Training loss (final) | 2.01 |
+| Validation PPL | 58.8 |
+| Tokens seen | 491M |
+
+## Usage
+
+### PyTorch
+
+```python
+import torch
+from tokenizers import Tokenizer
+from keyboard_lm.model import JointUniLM, ModelConfig
+
+# Load
+ckpt = torch.load("model.pt", map_location="cpu", weights_only=False)
+config = ModelConfig(**ckpt["model_config"])
+model = JointUniLM(config)
+model.load_state_dict(ckpt["model_state_dict"])
+model.eval()
+
+tokenizer = Tokenizer.from_file("tokenizer.json")
+
+# Predict next token
+prompt = "I'm going to the"
+ids = [config.bos_token_id] + tokenizer.encode(prompt).ids
+input_ids = torch.tensor([ids])
+
+with torch.no_grad():
+    logits, _ = model(input_ids)
+    probs = torch.softmax(logits[0, -1], dim=-1)
+    top5 = torch.topk(probs, 5)
+
+for prob, idx in zip(top5.values, top5.indices):
+    print(f"  {tokenizer.decode([idx.item()]):>10} ({prob:.1%})")
+```
+
+### CoreML (iOS)
+
+See `scripts/export_coreml.py` in the [GitHub repo](https://github.com/synth-inc/onit-keyboard-lm) for CoreML conversion.
+
+## Files
+
+| File | Description |
+|------|-------------|
+| `model.pt` | Model weights + config (no optimizer) |
+| `checkpoint_full.pt` | Full training checkpoint (with optimizer, for resume) |
+| `config.json` | Model configuration |
+| `tokenizer.json` | BPE tokenizer (v2, trained on Phase 2 mix) |
+
+## Limitations
+
+- Optimized for short text (keyboard input), not long-form generation
+- May produce grammatical errors in French (e.g., double negatives)
+- 256-token context window limits long-range coherence
+- Not suitable for factual Q&A or instruction following
+
+## License
+
+Apache 2.0

checkpoint_full.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1815019697c1ba8ac9c770097bd0234d9ead3a92c8aa74f40c567aef220eab7c
+size 486296722

config.json

@@ -0,0 +1,18 @@
+{
+  "vocab_size": 16384,
+  "dim": 512,
+  "num_layers": 10,
+  "num_heads": 8,
+  "ffn_dim": 1408,
+  "max_seq_len": 256,
+  "dropout": 0.1,
+  "tied_embeddings": true,
+  "qk_norm": true,
+  "rope_base": 10000.0,
+  "rms_norm_eps": 1e-06,
+  "pad_token_id": 0,
+  "unk_token_id": 1,
+  "bos_token_id": 2,
+  "eos_token_id": 3,
+  "mask_token_id": 4
+}

model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d8648eaa471f0baf1229546dfd29bb4b5532b9b37a1338c82bb50ce6be074049
+size 162087274