docs: komentar inline + README argumen + bugfix & optimasi (7dd92fc)

README.md

@@ -18,88 +18,386 @@ Model bahasa kecil GPT-style yang dibangun **dari nol** (tanpa library transform
 
 ## Arsitektur
 
-| | Nilai default |
+| Komponen | Nilai default |
 |---|---|
 | Tipe | Decoder-only transformer (pre-LN, SDPA) |
-| Parameter | ~0.81M |
+| Parameter | ~0.81M (default) |
 | Layer / Head | 4 / 4 |
-| Dimensi | 128 |
-| Konteks | 96 token |
-| Tokenizer | karakter atau BPE (`--tokenizer bpe`) |
+| Dimensi embedding | 128 |
+| Konteks (block_size) | 128 token |
+| Tokenizer | `char` (karakter) atau `bpe` (byte-pair encoding) |
+| Weight tying | Ya (output head = token embedding) |
+| Aktivasi | GELU (MLP) |
+| Attention | Multi-head causal self-attention + KV-cache |
 
-## File penting
+## Struktur file
 
 ```
-indigo/model.py        arsitektur GPT (attention kausal + KV-cache)
-indigo/bpe.py          tokenizer BPE byte-level minimal
-indigo/tokenizer.py    tokenizer karakter
-train.py               training (best-checkpoint, resume penuh, split val per-file)
-generate.py            generasi (top-k, top-p, repetition penalty)
-out/indigo_best.safetensors   bobot terbaik + _meta.json
+indigo/
+  model.py          Arsitektur GPT (CausalSelfAttention, MLP, Block, GPT, sampling)
+  common.py         Utilitas bersama (clean_text, meta save/load, morphological analysis)
+  bpe.py            Tokenizer BPE level-byte minimal
+  tokenizer.py      Tokenizer karakter (char-level)
+
+train.py            Training loop (best-checkpoint, resume, split val, cosine LR schedule)
+generate.py         Generasi teks (top-k, top-p, repetition penalty, guard kamus)
+pipeline.py         Pipeline end-to-end (data collection → training → manifest → finalize/push)
+eval.py             Evaluasi batched antar-checkpoint (nats/token, nats/karakter)
+audit_kamus.py      Audit wordlist kamus terhadap KBBI daring
+
+out/
+  indigo_best.safetensors   Bobot terbaik (dipilih berdasarkan validasi)
+  indigo_best_meta.json     Metadata terbaik
+  indigo.safetensors        Bobot checkpoint akhir
+  indigo_meta.json          Metadata akhir
+  indigo_optimizer.pt       State optimizer (untuk resume)
+  run_info.json             Info run terakhir yang di-finalize
+
+data/
+  sample.txt        Teks contoh Indonesia (training default)
+  kamus_id.txt      Wordlist Indonesia (~70K kata, gitignored)
+  prefiks.txt       Daftar prefiks Indonesia untuk morphological guard
+  sufixs.txt        Daftar sufiks Indonesia untuk morphological guard
+```
+
+## Persyaratan
+
+```
+torch>=2.0
+safetensors>=0.4
 ```
 
-## Cara pakai
+Install:
 
 ```bash
 pip install -r requirements.txt
+```
+
+---
+
+## `train.py` — Training Model
+
+Script utama untuk melatih model GPT dari nol.
 
+### Argumen
+
+| Argumen | Tipe | Default | Deskripsi |
+|---|---|---|---|
+| `--data` | str (nargs+) | `data/sample.txt` | Path file/folder teks untuk training. Bisa banyak, spasi-separated. Folder akan dicari rekursif untuk file `.txt`. |
+| `--out` | str | `out` | Folder output checkpoint (`.safetensors` + `_meta.json` + `_optimizer.pt`). |
+| `--steps` | int | `2000` | Jumlah total langkah training. |
+| `--batch-size` | int | `32` | Jumlah sampel per batch. Batch lebih besar → gradient lebih stabil tapi butuh lebih banyak VRAM. |
+| `--block-size` | int | `128` | Panjang konteks token per sampel (sequence length). Semakin besar → model bisa lihat lebih jauh tapi butuh lebih banyak memori. |
+| `--lr` | float | `3e-4` | Learning rate maksimum. Cosine decay dari lr ini ke 10% lr selama training. |
+| `--warmup` | int | `100` | Jumlah langkah warmup linear sebelum cosine decay dimulai. |
+| `--weight-decay` | float | `0.1` | L2 regularization / weight decay untuk AdamW optimizer. |
+| `--dropout` | float | `0.1` | Dropout rate (0.0 = nonaktif). Membantu mencegah overfitting pada data kecil. |
+| `--n-layer` | int | `4` | Jumlah blok transformer bertumpuk. Lebih banyak layer → model lebih dalam tapi lebih lambat. |
+| `--n-head` | int | `4` | Jumlah head per attention layer. Harus habis membagi `--n-embd`. |
+| `--n-embd` | int | `128` | Dimensi embedding / hidden size. Parameter ≈ 12 × n_layer × n_embd². |
+| `--tokenizer` | `char`/`bpe` | `char` | Jenis tokenizer. `char` = cepat, `bpe` = lebih efisien untuk teks panjang. |
+| `--vocab-size` | int | `512` | Ukuran vocab untuk BPE (diabaikan jika `--tokenizer char`). |
+| `--eval-interval` | int | `200` | Evaluasi validasi setiap N langkah. Set `0` untuk skip validasi. |
+| `--eval-iters` | int | `20` | Jumlah batch untuk estimasi loss validasi. |
+| `--seed` | int | `1337` | Seed random untuk reproduktibilitas. |
+| `--val-fraction` | float | `0.1` | Proporsi file untuk validasi (split per-file, bukan per-karakter). |
+| `--init-from` | str | `None` | Path checkpoint untuk melanjutkan training (resume). Muat model + optimizer + step. |
+| `--device` | `auto`/`cpu`/`cuda` | `auto` | Device training. `auto` = CUDA jika tersedia, else CPU. |
+
+### Contoh
+
+```bash
+# Training dasar
 python train.py --data data/sample.txt --steps 2000
-python train.py --data data/tekskamu.txt --tokenizer bpe --vocab-size 512
 
-python generate.py --prompt "Indigo" --max-new 300 \
-  --temperature 0.8 --top-k 40 --top-p 0.9 --repetition-penalty 1.2
+# Training dengan BPE tokenizer
+python train.py --data data/teks.txt --tokenizer bpe --vocab-size 512
+
+# Model lebih besar, training lebih lama
+python train.py --data data/ --n-layer 6 --n-head 8 --n-embd 256 --steps 5000
+
+# Resume dari checkpoint
+python train.py --init-from out/indigo_best.safetensors --steps 1000
+
+# Tanpa validasi (data hanya 1 file)
+python train.py --data data/combined.txt --val-fraction 0 --eval-interval 0
+```
+
+### Output
+
+- `out/indigo_best.safetensors` + `out/indigo_best_meta.json` — checkpoint terbaik (val loss minimum)
+- `out/indigo.safetensors` + `out/indigo_meta.json` — checkpoint akhir
+- `out/indigo_optimizer.pt` — state optimizer (untuk resume)
+
+### Learning Rate Schedule
 
-python train.py --init-from out/indigo_best.safetensors --steps 1000  # lanjutkan training
 ```
+Step 0 → warmup: linear dari 0 → lr_max
+Step warmup → total: cosine decay dari lr_max → 0.1 × lr_max
+```
+
+Formula cosine: `0.1 × lr + 0.45 × lr × (1 + cos(π × progress))`
+
+---
+
+## `generate.py` — Generasi Teks
+
+Generate teks dari checkpoint Indigo dengan berbagai opsi sampling.
 
-Generasi memakai KV-cache sehingga cepat untuk output panjang.
+### Argumen
 
-### Penjaga kamus (kualitas ejaan)
+| Argumen | Tipe | Default | Deskripsi |
+|---|---|---|---|
+| `--ckpt` | str | `out/indigo_best.safetensors` | Path ke file checkpoint model. |
+| `--prompt` | str | `""` | Teks awal (prompt) untuk memulai generasi. |
+| `--max-new` | int | `300` | Jumlah token baru yang akan dihasilkan (bukan panjang total). |
+| `--temperature` | float | `0.8` | Skala randomness: `0.0` ≈ greedy, `0.8` ≈ standar, `>1.0` ≈ random. |
+| `--top-k` | int | `40` | Batasi sampling ke k token teratas. `0` = nonaktif. |
+| `--top-p` | float | `1.0` | Nucleus sampling: batasi kumulatif probabilitas. `1.0` = nonaktif. |
+| `--repetition-penalty` | float | `1.0` | Penalti pengulangan token. `1.0` = nonaktif, `>1.0` = kurangi pengulangan. |
+| `--seed` | int | `None` | Seed random. `None` = tidak ditentukan (random setiap kali). |
+| `--device` | `auto`/`cpu`/`cuda` | `auto` | Device untuk inferensi. |
+| `--guard` | str | `None` | Path file kamus (satu kata per baris). Generate beberapa kandidat → pilih yang rasio kata dikenal tertinggi. |
+| `--guard-prefiks` | str | `None` | Path file prefiks Indonesia. Default: `data/prefiks.txt` bila ada. |
+| `--guard-sufiks` | str | `None` | Path file sufiks Indonesia. Default: `data/sufiks.txt` bila ada. |
+| `--guard-tries` | int | `5` | Jumlah kandidat generate saat `--guard` aktif. |
+| `--guard-min` | float | `0.6` | Rasio kata dikenal minimum. Berhenti generate lebih awal jika tercapai. |
 
-`--guard` memuat kamus kata (satu kata per baris, mis. wordlist KBBI berlisensi MIT dari
-`hf.co/datasets/Deddy/Indonesia-dataset-2023`), lalu memilih kandidat generasi dengan
-rasio kata dikenal tertinggi — menyaring ejaan tak-baku tanpa mengubah model:
+### Contoh
 
 ```bash
+# Generasi dasar
+python generate.py --prompt "Indigo" --max-new 300
+
+# Sampling lebih random
+python generate.py --prompt "hello" --temperature 1.0 --top-k 60 --top-p 0.9
+
+# Dengan guard kamus (pilih output paling koheren)
 python generate.py --prompt "kepekaan" --max-new 120 --guard data/kamus_id.txt
 
-# metrik otomatis per run di manifest.json ("kamus_ratio")
-python pipeline.py --tag run03 --data data/ --steps 2000 --guard data/kamus_id.txt
+# Repetition penalty untuk mengurangi pengulangan
+python generate.py --prompt "cerita" --repetition-penalty 1.2 --max-new 200
 ```
 
-Kata berimbuhan dicek lewat formula morfologi: `data/prefiks.txt` & `data/sufiks.txt`
-(mis. *dipahami* → akar *paham*, termasuk asimilasi *meny*-→*s*, *pem*-→*p*).
-Edit kedua file itu untuk memperluas cakupan tanpa menyentuh kode.
+### Guard Kamus
+
+Saat `--guard` aktif, generate.py akan:
+1. Generate `--guard-tries` kandidat teks (seed berbeda tiap kandidat)
+2. Untuk setiap kandidat, hitung rasio kata yang dikenal kamus
+3. Pilih kandidat dengan rasio tertinggi
+4. Berhenti lebih awal jika rasio >= `--guard-min`
+
+Kata berimbuhan dicek lewat formula morfologi:
+- Hapus sufiks → cek akar
+- Hapus prefiks → cek akar (+ asimilasi: `meny-`→`s`, `pem-`→`p`)
+
+---
 
-### Pipeline
+## `pipeline.py` — Pipeline End-to-End
 
-`pipeline.py` merangkai semuanya: menyalin data lokal dan/atau menarik dataset HF,
-melatih, menyimpan `manifest.json` (statistik + konfigurasi run), lalu opsional
-mempromosikan checkpoint terbaik ke `out/` dan mengunggahnya ke Hub.
+Rangkai semua tahap: kumpul data → training → evaluasi → manifest → finalize → push.
+
+### Argumen
+
+| Argumen | Tipe | Default | Deskripsi |
+|---|---|---|---|
+| `--tag` | str | **wajib** | Nama run. Semua artefak disimpan di `runs/<tag>/`. |
+| `--data` | str (nargs*) | `[]` | File/folder teks lokal tambahan (banyak, spasi-separated). |
+| `--hf-dataset` | str | `None` | Repo dataset HF untuk menarik file teks (mis. `adyoi/indigo`). |
+| `--hf-patterns` | str (nargs*) | `["*.txt", "*.md"]` | Pola file yang diambil dari HF. |
+| `--format-qa` | flag | `False` | Auto-convert Alpaca JSON (`instruction`/`output`) ke `.txt` sebelum training. |
+| `--runs` | str | `runs` | Folder root untuk semua run. |
+| `--device` | str | `None` | Device training (diteruskan ke `train.py`, default: auto). |
+| `--finalize` | flag | `False` | Promosikan checkpoint terbaik ke folder `out/` kanonik. |
+| `--push` | flag | `False` | Upload checkpoint terbaik ke repo HF. |
+| `--repo` | str | `adyoi/indigo` | Repo HF tujuan upload. |
+| `--guard` | str | `None` | File kamus untuk metrik rasio ejaan di manifest. |
+| `--guard-max-new` | int | `120` | Jumlah token generate untuk evaluasi guard. |
+| `--guard-prefiks` | str | `None` | File prefiks Indonesia. |
+| `--guard-sufiks` | str | `None` | File sufiks Indonesia. |
+
+**Hyperparameter training** (diteruskan ke `train.py`):
+
+| Argumen | Tipe | Default | Deskripsi |
+|---|---|---|---|
+| `--steps` | int | `2000` | Jumlah langkah training. |
+| `--batch-size` | int | `32` | Batch size. |
+| `--block-size` | int | `128` | Panjang konteks. |
+| `--n-layer` | int | `4` | Jumlah layer transformer. |
+| `--n-head` | int | `4` | Jumlah head attention. |
+| `--n-embd` | int | `128` | Dimensi embedding. |
+| `--dropout` | float | `0.1` | Dropout rate. |
+| `--lr` | float | `3e-4` | Learning rate. |
+| `--warmup` | int | `100` | Langkah warmup. |
+| `--weight-decay` | float | `0.1` | Weight decay. |
+| `--eval-interval` | int | `200` | Evaluasi setiap N langkah. |
+| `--eval-iters` | int | `20` | Jumlah batch evaluasi. |
+| `--seed` | int | `1337` | Seed random. |
+| `--init-from` | str | `None` | Checkpoint untuk resume. |
+| `--tokenizer` | `char`/`bpe` | `char` | Jenis tokenizer. |
+| `--vocab-size` | int | `512` | Vocab size untuk BPE. |
+| `--val-fraction` | float | `0.1` | Proporsi file validasi. |
+
+### Contoh
 
 ```bash
-# data lokal + dataset HF, tokenizer BPE
-python pipeline.py --tag run01 --data data/tekskamu.txt --hf-dataset adyoi/indigo \
+# Data lokal saja
+python pipeline.py --tag run01 --data data/sample.txt --steps 2000
+
+# Data dari HF + BPE tokenizer
+python pipeline.py --tag run02 --data data/teks.txt --hf-dataset adyoi/indigo \
   --tokenizer bpe --vocab-size 512 --steps 900 --device cpu
 
-# promosikan hasil terbaik ke out/ sekaligus upload ke Hub
-python pipeline.py --tag run02 --data data/ --steps 2000 --finalize --push
+# Dataset Alpaca JSON dari HF
+python pipeline.py --tag qa01 \
+  --hf-dataset rohanrdy/CS-Theory-QA-Dataset \
+  --format-qa --steps 3000 --n-layer 6 --n-embd 256
+
+# Full pipeline: train + finalize + push
+python pipeline.py --tag final01 --data data/ --steps 2000 --finalize --push
+
+# Dengan guard kamus untuk metrik kualitas
+python pipeline.py --tag guarded --data data/ --steps 2000 --guard data/kamus_id.txt
 ```
 
-### Evaluasi standar antar-checkpoint
+### Output
+
+- `runs/<tag>/data/` — data yang digunakan (copy dari sumber)
+- `runs/<tag>/ckpt/` — checkpoint training
+- `runs/<tag>/manifest.json` — statistik + metadata run (termasuk `kamus_ratio` jika `--guard` aktif)
+
+---
 
-`eval.py` menykor checkpoint pada set uji tetap (nats/token & nats/karakter) sehingga
-run dengan tokenizer berbeda tetap sebanding:
+## `eval.py` — Evaluasi Batched
+
+Bandingkan beberapa checkpoint pada set uji tetap. Metrik:
+
+- **nats/token**: loss rata-rata per token (cross-entropy). Lebih rendah = lebih baik.
+- **nats/karakter**: loss per karakter (dikoreksi dengan compression ratio). Memungkinkan perbandingan antar tokenizer.
+
+### Argumen
+
+| Argumen | Tipe | Default | Deskripsi |
+|---|---|---|---|
+| `--ckpt` | str (nargs+) | **wajib** | Path ke satu atau lebih file checkpoint (`.safetensors`). |
+| `--test` | str | `data/sample.txt` | Path ke file teks uji. |
+| `--device` | `cpu`/`cuda` | `cpu` | Device untuk evaluasi. |
+| `--guard` | str | `None` | Path file kamus; generate teks → hitung rasio kata dikenal. |
+| `--guard-max-new` | int | `120` | Jumlah token generate untuk evaluasi guard. |
+| `--seed` | int | `42` | Seed untuk generate saat `--guard` aktif. |
+| `--batch-size` | int | `32` | Batch size untuk evaluasi. |
+
+### Contoh
 
 ```bash
+# Evaluasi satu checkpoint
+python eval.py --ckpt out/indigo_best.safetensors
+
+# Bandingkan beberapa checkpoint
 python eval.py --ckpt out/indigo_best.safetensors runs/*/ckpt/indigo.safetensors \
-  --test data/sample.txt --guard data/kamus_id.txt
+  --test data/sample.txt
+
+# Dengan guard kamus
+python eval.py --ckpt out/indigo_best.safetensors --guard data/kampus_id.txt
+```
+
+### Contoh output
+
+```
+set uji: data/sample.txt (29,123 karakter)
+checkpoint                                 nats/tok  nat/kar   kamus
+run/csqa500/ckpt/indigo_best.safetensors      3.962    3.962     32%
+run/wiki500/ckpt/indigo_best.safetensors      3.018    3.018     25%
+```
+
+---
+
+## `audit_kamus.py` — Audit Wordlist Kamus
+
+Verifikasi kata-kata dalam file kamus terhadap KBBI daring (sampling acak).
+
+### Argumen
+
+| Argumen | Tipe | Default | Deskripsi |
+|---|---|---|---|
+| `--kamus` | str | `data/kamus_id.txt` | Path ke file kamus (satu kata per baris). |
+| `--n` | int | `200` | Jumlah kata sampel yang diambil untuk verifikasi. |
+| `--seed` | int | `1337` | Seed random untuk sampling. |
+| `--delay` | float | `0.8` | Jeda antar-permintaan web (detik). Hindari rate limiting. |
+| `--backend` | `auto`/`pypi`/`web` | `auto` | Backend verifikasi. `pypi` = library kbbi, `web` = scraping kbbi.web.id. |
+| `--apply` | flag | `False` | Hapus kata TIDAK-DIKENAL dari kamus (backup otomatis ke `.bak.txt`). |
+| `--out` | str | `None` | Path output laporan JSON. Default: `runs/kamus_audit_<timestamp>.json`. |
+
+### Contoh
+
+```bash
+# Audit dasar
+python audit_kamus.py --kamus data/kamus_id.txt --n 200
+
+# Hapus kata tak-valid (backup otomatis)
+python audit_kamus.py --kamus data/kamus_id.txt --apply
+
+# Backend web saja, delay lebih lama
+python audit_kamus.py --backend web --delay 1.5 --n 100
 ```
 
+### Output
+
+Laporan JSON berisi:
+- `ada`: jumlah kata yang ditemukan di KBBI
+- `tidak_dikenal`: jumlah kata yang tidak ditemukan
+- `gagal_ceks`: jumlah verifikasi yang gagal (timeout, dll)
+- `estimasi_entri_tidak_valid`: proyeksi total entri tak-valid di seluruh kamus
+- `kata_ditolak`: daftar kata yang tidak dikenal (akan dihapus jika `--apply`)
+
+---
+
+## Guard Kamus (Sistem Kualitas Ejaan)
+
+Sistem guard membantu menjaga kualitas output model tanpa mengubah bobot:
+
+1. **Wordlist**: `data/kamus_id.txt` — ~70K kata Indonesia (satu kata per baris)
+2. **Morfologi**: `data/prefiks.txt` + `data/sufiks.txt` — pendeteksi imbuhan
+3. **Asimilasi**: `meny-`→`s`, `peny-`→`s`, `pem-`→`p` — menangani perubahan bunyi
+
+Cara kerja:
+- **generate.py**: generate beberapa kandidat → pilih yang rasio kata dikenal tertinggi
+- **pipeline.py**: generate sekali → simpan `kamus_ratio` di manifest.json
+- **eval.py**: generate sekali → tampilkan rasio di tabel perbandingan
+
+Edit `data/prefiks.txt` dan `data/sufiks.txt` untuk memperluas cakupan tanpa menyentuh kode.
+
+---
+
+## Format Checkpoint
+
+| File | Isi |
+|---|---|
+| `*.safetensors` | Bobot model (format aman, tanpa pickle) |
+| `*_meta.json` | Metadata: config, vocab, step, val_loss, backend, tokenizer info |
+| `*_optimizer.pt` | State optimizer AdamW (untuk resume training) |
+
+Format `.safetensors` tidak mengeksekusi kode saat dimuat — lebih aman dari format `.pt` lama.
+
+---
+
+## Tips Training
+
+| Skenario | Rekomendasi |
+|---|---|
+| Data < 100 KB | `--tokenizer char`, 4L/4H/128E, 2000 steps |
+| Data 100 KB–1 MB | `--tokenizer char` atau `bpe` (vocab 512), 4L/8H/256E, 3000+ steps |
+| Data > 1 MB | `--tokenizer char` (BPE training sangat lambat di CPU >1MB) |
+| CPU-only | Model max ~1M params (4L/4H/128E), 5-8 detik/step |
+| Resume training | `--init-from runs/xxx/ckpt/indigo_best.safetensors` |
+
+---
+
 ## Batasan
 
 - Dilatih pada data sangat kecil → output belum koheren; cocok untuk edukasi bukan produksi.
 - Gunakan checkpoint `indigo_best` (terpilih berdasarkan validasi), bukan checkpoint akhir.
+- BPE training sangat lambat di CPU untuk data > 1 MB — gunakan `--tokenizer char`.
 
 ## Keamanan
 

audit_kamus.py

@@ -1,3 +1,22 @@
+"""
+Script audit wordlist kamus terhadap KBBI daring.
+
+Mengambil sampel random dari file kamus (satu kata per baris), lalu memverifikasi
+setiap kata menggunakan library kbbi (pypi) atau scraping kbbi.web.id.
+
+Fitur:
+- Verifikasi via pypi package `kbbi` (lebih stabil)
+- Fallback ke web scraping kbbi.web.id (heuristik, rapuh)
+- Estimasi jumlah entri tak-valid di seluruh kamus
+- Mode --apply: hapus kata tak-valid + backup otomatis
+- Laporan JSON lengkap (timestamp, statistik, daftar kata ditolak)
+
+Cara pakai:
+    python audit_kamus.py --kamus data/kamus_id.txt --n 200
+    python audit_kamus.py --kamus data/kamus_id.txt --apply
+    python audit_kamus.py --backend web --delay 1.0
+"""
+
 import argparse
 import json
 import os
@@ -11,6 +30,14 @@ UA = {"User-Agent": "Mozilla/5.0 (audit-kamus-indigo)"}
 
 
 def cek_pypi(word):
+    """Verifikasi kata menggunakan library kbbi dari pypi.
+
+    Args:
+        word: Kata yang akan diverifikasi.
+
+    Returns:
+        True jika kata ditemukan di KBBI, False jika tidak, None jika error.
+    """
     try:
         from kbbi import KBBI, TidakDitemukan
     except ImportError:
@@ -28,6 +55,21 @@ _SESI = None
 
 
 def cek_web(word):
+    """Verifikasi kata dengan scraping kbbi.web.id (fallback).
+
+    Metode heuristik: hitung jumlah marker "Maaf, tidak ditemukan" di halaman.
+    - 0 marker → kata ditemukan
+    - < 3 marker → kemungkinan ada (ambigu)
+    - >= 3 marker → tidak ditemukan
+
+    Peringatan: metode ini rapuh (bergantung pada struktur HTML kbbi.web.id).
+
+    Args:
+        word: Kata yang akan diverifikasi.
+
+    Returns:
+        True/False/None — sama seperti cek_pypi.
+    """
     global _SESI
     import requests
 
@@ -53,23 +95,30 @@ def main():
     ap = argparse.ArgumentParser(
         description="Audit wordlist kamus terhadap KBBI daring (sampling acak)"
     )
-    ap.add_argument("--kamus", default=str(ROOT / "data" / "kamus_id.txt"))
-    ap.add_argument("--n", type=int, default=200, help="jumlah kata yang diambil sampel")
-    ap.add_argument("--seed", type=int, default=1337)
-    ap.add_argument("--delay", type=float, default=0.8, help="jeda antar-permintaan (detik)")
+    ap.add_argument("--kamus", default=str(ROOT / "data" / "kamus_id.txt"),
+                    help="path ke file kamus (satu kata per baris, default: data/kamus_id.txt)")
+    ap.add_argument("--n", type=int, default=200,
+                    help="jumlah kata sampel yang diambil (default: 200)")
+    ap.add_argument("--seed", type=int, default=1337,
+                    help="seed random untuk sampling (default: 1337)")
+    ap.add_argument("--delay", type=float, default=0.8,
+                    help="jeda antar-permintaan web dalam detik (default: 0.8)")
     ap.add_argument("--backend", choices=["auto", "pypi", "web"], default="auto",
-                    help="web: kbbi.web.id (heuristik jumlah marker; rapuh)")
+                    help="backend verifikasi: auto/pypi/web (default: auto)")
     ap.add_argument("--apply", action="store_true",
-                    help="hapus kata TIDAK-DIKENAL dari kamus (backup otomatis)")
-    ap.add_argument("--out", default=None, help="path laporan json")
+                    help="hapus kata TIDAK-DIKENAL dari kamus (backup otomatis ke .bak.txt)")
+    ap.add_argument("--out", default=None,
+                    help="path output laporan JSON (default: runs/kamus_audit_<timestamp>.json)")
     args = ap.parse_args()
 
+    # --- Muat kamus ---
     kamus_path = Path(args.kamus)
     kata_semua = [w.strip() for w in kamus_path.read_text(encoding="utf-8").splitlines() if w.strip()]
     rng = random.Random(args.seed)
     sampel = sorted(rng.sample(kata_semua, min(args.n, len(kata_semua))))
     print(f"kamus={kamus_path.name} ({len(kata_semua):,} kata) | sampel={len(sampel)} | seed={args.seed}")
 
+    # --- Verifikasi setiap kata sampel ---
     hasil = {}
     n_ada = n_tidak = n_gagal = 0
     t0 = time.time()
@@ -90,6 +139,7 @@ def main():
             print(f"  {i}/{len(sampel)} | ada={n_ada} tidak={n_tidak} gagal={n_gagal} | {time.time() - t0:.0f}s")
         time.sleep(args.delay)
 
+    # --- Buat laporan ---
     ditolak = sorted(w for w, s in hasil.items() if s is False)
     laporan = {
         "timestamp": datetime.now().isoformat(timespec="seconds"),
@@ -110,6 +160,7 @@ def main():
     print(f"ringkasan: ada={n_ada} | tidak-dikenal={n_tidak} | gagal={n_gagal} "
           f"| estimasi entri tak-valid di kamus: ~{laporan['estimasi_entri_tidak_valid']:,}")
 
+    # --- Hapus kata tak-valid (jika --apply) ---
     if args.apply and ditolak:
         backup = kamus_path.with_suffix(".bak.txt")
         backup.write_text("\n".join(sorted(kata_semua)) + "\n", encoding="utf-8")

eval.py

@@ -1,3 +1,24 @@
+"""
+Script evaluasi batched untuk membandingkan checkpoint Indigo.
+
+Menghitung metrik per-token (nats/token) dan per-karakter (nats/karakter) pada
+set uji tetap, sehingga run dengan tokenizer berbeda (char vs BPE) tetap sebanding.
+
+Alur kerja:
+1. Muat teks uji (data/sample.txt atau custom)
+2. Untuk setiap checkpoint:
+   a. Muat model + tokenizer
+   b. Encode teks uji → pecah menjadi jendela (block_size)
+   c. Hitung cross-entropy loss per-token secara batched
+   d. Konversi ke nats/karakter menggunakan compression ratio
+3. (Opsional) Generate teks → hitung rasio kata dikenal (kamus guard)
+
+Cara pakai:
+    python eval.py --ckpt out/indigo_best.safetensors
+    python eval.py --ckpt runs/*/ckpt/indigo.safetensors --test data/sample.txt
+    python eval.py --ckpt out/indigo_best.safetensors --guard data/kamus_id.txt
+"""
+
 import argparse
 from pathlib import Path
 
@@ -14,6 +35,14 @@ from indigo.model import GPT, GPTConfig
 
 
 def muat(path):
+    """Muat model GPT + tokenizer dari checkpoint .safetensors.
+
+    Args:
+        path: Path ke file .safetensors.
+
+    Returns:
+        Tuple (model, tokenizer, meta).
+    """
     from safetensors.torch import load_file
 
     meta = load_meta(path)
@@ -28,30 +57,58 @@ def muat(path):
 
 @torch.no_grad()
 def nats_per_token(model, ids, block_size, device, batch_size=32):
+    """Hitung loss rata-rata (nats per token) pada seluruh sequence.
+
+    Algoritma:
+    1. Pecah sequence panjang menjadi jendela-jendela sepanjang block_size
+    2. Pad jendela ke panjang yang sama dalam batch (zero-padding + mask)
+    3. Forward pass batched → hitung cross-entropy per token → rata-rata
+
+    Args:
+        model:      Model GPT.
+        ids:        List of int — token IDs dari teks uji.
+        block_size: Int — panjang konteks model.
+        device:     Str — "cpu" atau "cuda".
+        batch_size: Int — jumlah jendela per batch (default: 32).
+
+    Returns:
+        Tuple (nats_per_token, total_tokens).
+    """
+    # Pecah sequence menjadi jendela-jendela block_size
     jendela = []
     for i in range(0, max(0, len(ids) - 1), block_size):
-        potongan = ids[i : i + block_size + 1]
+        potongan = ids[i : i + block_size + 1]  # +1 untuk target
         if len(potongan) >= 2:
             jendela.append(potongan)
+
+    # Proses batched
     total_nll = 0.0
     total_tok = 0
     for k in range(0, len(jendela), batch_size):
         kelompok = jendela[k : k + batch_size]
-        L = max(len(w) - 1 for w in kelompok)
+        L = max(len(w) - 1 for w in kelompok)  # panjang terpanjang dalam batch
+
+        # Buat tensor x (input), y (target), dan mask (ignore padding)
         x = torch.zeros(len(kelompok), L, dtype=torch.long)
         y = torch.zeros(len(kelompok), L, dtype=torch.long)
         mask = torch.zeros(len(kelompok), L, dtype=torch.bool)
         for r, w in enumerate(kelompok):
             n = len(w) - 1
-            x[r, :n] = torch.tensor(w[:-1], dtype=torch.long)
-            y[r, :n] = torch.tensor(w[1:], dtype=torch.long)
-            mask[r, :n] = True
+            x[r, :n] = torch.tensor(w[:-1], dtype=torch.long)  # input: semua kecuali terakhir
+            y[r, :n] = torch.tensor(w[1:], dtype=torch.long)   # target: semua kecuali pertama
+            mask[r, :n] = True  # hanya hitung posisi yang ada isinya
+
         x, y, mask = x.to(device), y.to(device), mask.to(device)
+
+        # Forward pass → log probability → negative log-likelihood
         logits, _ = model(x)
         logp = torch.log_softmax(logits.float(), dim=-1)
         nll = -logp.gather(2, y.unsqueeze(2)).squeeze(2)
+
+        # Akumulasi (hanya hitung posisi yang dimask)
         total_nll += float(nll[mask].sum())
         total_tok += int(mask.sum())
+
     return total_nll / max(1, total_tok), total_tok
 
 
@@ -59,16 +116,27 @@ def main():
     ap = argparse.ArgumentParser(
         description="Skor checkpoint pada set uji tetap agar antar-run dapat dibandingkan"
     )
-    ap.add_argument("--ckpt", nargs="+", required=True)
-    ap.add_argument("--test", default="data/sample.txt")
-    ap.add_argument("--device", default="cpu", choices=["cpu", "cuda"])
-    ap.add_argument("--guard", default=None, help="kamus opsional untuk metrik rasio ejaan")
-    ap.add_argument("--seed", type=int, default=42)
-    ap.add_argument("--batch-size", type=int, default=32)
+    ap.add_argument("--ckpt", nargs="+", required=True,
+                    help="path ke satu atau lebih file checkpoint (.safetensors)")
+    ap.add_argument("--test", default="data/sample.txt",
+                    help="path ke file teks uji (default: data/sample.txt)")
+    ap.add_argument("--device", default="cpu", choices=["cpu", "cuda"],
+                    help="device untuk evaluasi (default: cpu)")
+    ap.add_argument("--guard", default=None,
+                    help="path ke file kamus; generate teks → hitung rasio kata dikenal")
+    ap.add_argument("--guard-max-new", type=int, default=120,
+                    help="jumlah token generate untuk evaluasi guard (default: 120)")
+    ap.add_argument("--seed", type=int, default=42,
+                    help="seed untuk generate saat --guard aktif (default: 42)")
+    ap.add_argument("--batch-size", type=int, default=32,
+                    help="batch size untuk evaluasi (default: 32)")
     args = ap.parse_args()
 
+    # --- Muat teks uji ---
     teks = read_clean(args.test)
     n_karakter = len(teks.encode("utf-8"))
+
+    # --- Muat kamus (jika --guard) ---
     wordset = prefiks = sufiks = None
     if args.guard:
         root = Path(__file__).resolve().parent
@@ -77,28 +145,39 @@ def main():
         prefiks = load_wordlist(str(p)) if p.exists() else None
         sufiks = load_wordlist(str(s)) if s.exists() else None
 
+    # --- Header tabel ---
     print(f"set uji: {args.test} ({n_karakter:,} karakter)")
     print(f"{'checkpoint':44s} {'nats/tok':>9s} {'nat/kar':>8s} {'kamus':>7s}")
     baris = []
+
+    # --- Evaluasi setiap checkpoint ---
     for path in args.ckpt:
         model, tokenizer, meta = muat(path)
+
+        # Encode teks uji → hitung nats per token
         ids = tokenizer.encode(teks)
         npt, n_tok = nats_per_token(
             model, ids, meta["config"]["block_size"], args.device, args.batch_size
         )
+
+        # Konversi: nats/token → nats/karakter (menggunakan compression ratio)
         kompresi = n_karakter / max(1, n_tok)
         npc = npt / kompresi
+
+        # (Opsional) hitung rasio kata dikenal via generate
         rasio = ""
         if wordset:
             torch.manual_seed(args.seed)
             out = model.generate(
                 torch.tensor([[0]], dtype=torch.long, device=args.device),
-                120,
+                args.guard_max_new,
                 temperature=0.8,
                 top_k=40,
             )
             teks_out = tokenizer.decode(out[0].tolist())
             rasio = f"{word_known_ratio(teks_out, wordset, prefiks, sufiks):6.0%}"
+
+        # Format nama checkpoint yang pendek (runs/xxx/ckpt/file.safetensors)
         bagian = str(Path(path)).replace("\\", "/").split("/")
         nama = "/".join(bagian[-3:-1] + [bagian[-1]]) if len(bagian) >= 3 else bagian[-1]
         print(f"{nama:44s} {npt:9.3f} {npc:8.3f} {rasio:>7s}")

generate.py

@@ -1,3 +1,18 @@
+"""
+Script inferensi / generasi teks dari checkpoint Indigo.
+
+Mendukung:
+- Top-k, top-p, temperature, repetition penalty sampling
+- KV-cache untuk generasi cepat (token-by-token)
+- Guard kamus: generate beberapa kandidat → pilih yang rasio kata dikenal tertinggi
+- Guard morfologi: cek imbuhan Indonesia (prefiks + sufiks + asimilasi)
+
+Cara pakai:
+    python generate.py --prompt "Indigo" --max-new 300
+    python generate.py --prompt "hello" --temperature 0.8 --top-k 40 --top-p 0.9
+    python generate.py --prompt "kepekaan" --guard data/kamus_id.txt --guard-min 0.6
+"""
+
 import sys
 import torch
 import argparse
@@ -9,6 +24,18 @@ from indigo.model import GPT, GPTConfig
 
 
 def load_model(path):
+    """Muat model GPT + tokenizer dari file checkpoint.
+
+    Mendukung dua format:
+    1. .safetensors: format utama Indigo
+    2. .pt: format PyTorch lama
+
+    Args:
+        path: Path ke file checkpoint.
+
+    Returns:
+        Tuple (model, tokenizer).
+    """
     if path.endswith(".safetensors"):
         from safetensors.torch import load_file
 
@@ -31,29 +58,55 @@ def load_model(path):
 
 def main():
     parser = argparse.ArgumentParser(description="Generate teks dari checkpoint Indigo")
-    parser.add_argument("--ckpt", default="out/indigo_best.safetensors")
-    parser.add_argument("--prompt", default="")
-    parser.add_argument("--max-new", type=int, default=300)
-    parser.add_argument("--temperature", type=float, default=0.8)
-    parser.add_argument("--top-k", type=int, default=40)
-    parser.add_argument("--top-p", type=float, default=1.0)
-    parser.add_argument("--repetition-penalty", type=float, default=1.0)
-    parser.add_argument("--seed", type=int, default=None)
-    parser.add_argument("--device", default="auto", choices=["auto", "cpu", "cuda"])
-    parser.add_argument("--guard", default=None, help="file kamus (satu kata per baris); pilih kandidat dengan ejaan terbaik")
-    parser.add_argument("--guard-prefiks", default=None, help="daftar prefiks (default data/prefiks.txt bila ada)")
-    parser.add_argument("--guard-sufiks", default=None, help="daftar sufiks (default data/sufiks.txt bila ada)")
-    parser.add_argument("--guard-tries", type=int, default=5, help="jumlah kandidat saat --guard aktif")
-    parser.add_argument("--guard-min", type=float, default=0.6, help="rasio kata dikenal minimum (info saja)")
+
+    # --- Checkpoint ---
+    parser.add_argument("--ckpt", default="out/indigo_best.safetensors",
+                        help="path ke file checkpoint model (default: out/indigo_best.safetensors)")
+
+    # --- Prompt & Generasi ---
+    parser.add_argument("--prompt", default="",
+                        help="teks awal (prompt) untuk memulai generasi (default: kosong)")
+    parser.add_argument("--max-new", type=int, default=300,
+                        help="jumlah token baru yang akan dihasilkan (default: 300)")
+    parser.add_argument("--temperature", type=float, default=0.8,
+                        help="skala randomness: 0.0 ≈ greedy, 0.8 ≈ standar, >1.0 ≈ random (default: 0.8)")
+    parser.add_argument("--top-k", type=int, default=40,
+                        help="batasi sampling ke k token teratas (0 = nonaktif, default: 40)")
+    parser.add_argument("--top-p", type=float, default=1.0,
+                        help="nucleus sampling: batasi kumulatif probabilitas (1.0 = nonaktif, default: 1.0)")
+    parser.add_argument("--repetition-penalty", type=float, default=1.0,
+                        help="penalti pengulangan token (>1.0 = aktif, 1.0 = nonaktif, default: 1.0)")
+    parser.add_argument("--seed", type=int, default=None,
+                        help="seed random (None = tidak ditentukan, default: None)")
+
+    # --- Device ---
+    parser.add_argument("--device", default="auto", choices=["auto", "cpu", "cuda"],
+                        help="device: auto/cpu/cuda (default: auto)")
+
+    # --- Guard Kamus ---
+    parser.add_argument("--guard", default=None,
+                        help="path file kamus (satu kata per baris); generate beberapa kandidat → pilih terbaik")
+    parser.add_argument("--guard-prefiks", default=None,
+                        help="path file prefiks Indonesia (default: data/prefiks.txt bila ada)")
+    parser.add_argument("--guard-sufiks", default=None,
+                        help="path file sufiks Indonesia (default: data/sufiks.txt bila ada)")
+    parser.add_argument("--guard-tries", type=int, default=5,
+                        help="jumlah kandidat generate saat --guard aktif (default: 5)")
+    parser.add_argument("--guard-min", type=float, default=0.6,
+                        help="rasio kata dikenal minimum — berhenti generate jika tercapai (default: 0.6)")
+
     args = parser.parse_args()
 
+    # --- Setup seed & device ---
     if args.seed is not None:
         torch.manual_seed(args.seed)
     device = "cuda" if torch.cuda.is_available() else "cpu" if args.device == "auto" else args.device
 
+    # --- Muat model ---
     model, tokenizer = load_model(args.ckpt)
     model = model.to(device)
 
+    # --- Muat kamus (jika --guard aktif) ---
     wordset = None
     pref_set = suf_set = None
     if args.guard:
@@ -63,10 +116,12 @@ def main():
 
         wordset = load_wordlist(args.guard)
         p_def, s_def = _Path("data/prefiks.txt"), _Path("data/sufiks.txt")
+        # Muat prefiks: prioritaskan argumen CLI → default path
         if args.guard_prefiks and _Path(args.guard_prefiks).exists():
             pref_set = load_wordlist(args.guard_prefiks)
         elif not args.guard_prefiks and p_def.exists():
             pref_set = load_wordlist(str(p_def))
+        # Muat sufiks: prioritaskan argumen CLI → default path
         if args.guard_sufiks and _Path(args.guard_sufiks).exists():
             suf_set = load_wordlist(args.guard_sufiks)
         elif not args.guard_sufiks and s_def.exists():
@@ -74,10 +129,16 @@ def main():
         mode = "dengan formula afiks" if pref_set and suf_set else "kata persis"
         print(f"[guard] kamus: {len(wordset):,} kata ({mode}) | target rasio >= {args.guard_min:.0%}")
 
+    # --- Encode prompt → token IDs ---
     ids = tokenizer.encode(args.prompt) or [0]
     idx = torch.tensor([ids], dtype=torch.long, device=device)
 
     def sample():
+        """Generate satu kandidat teks dari model.
+
+        Returns:
+            Tuple (text, ratio) — teks hasil generate dan rasio kata dikenal.
+        """
         out = model.generate(
             idx,
             args.max_new,
@@ -90,20 +151,22 @@ def main():
         ratio = word_known_ratio(text, wordset, pref_set, suf_set) if wordset else 1.0
         return text, ratio
 
+    # --- Tanpa guard: langsung generate & print ---
     if wordset is None:
         text, _ = sample()
         print(text)
         return
 
+    # --- Dengan guard: generate beberapa kandidat → pilih yang terbaik ---
     best_text, best_ratio = "", -1.0
     for t in range(args.guard_tries):
-        torch.manual_seed((args.seed or 0) + t * 1013)
+        torch.manual_seed((args.seed or 0) + t * 1013)  # seed berbeda tiap kandidat
         text, ratio = sample()
         mark = f"  [kandidat {t + 1}: {ratio:.0%}]"
         if ratio > best_ratio:
             best_text, best_ratio = text, ratio
         if best_ratio >= args.guard_min:
-            break
+            break  # sudah cukup bagus, tidak perlu generate lagi
     print(best_text)
     print(f"[guard] rasio kata dikenal: {best_ratio:.0%}")
 

indigo/bpe.py

@@ -1,4 +1,27 @@
+"""
+Tokenizer BPE (Byte Pair Encoding) level-byte untuk Indigo.
+
+BPE adalah algoritma tokenisasi yang belajar merge pasangan byte yang paling sering muncul
+secara berulang, menghasilkan subword units yang lebih efisien dari karakter tunggal.
+
+Alur kerja:
+1. Training: iteratively cari pasangan byte paling sering → merge → tambah ke vocabulary
+2. Encoding: teks → daftar byte → terapkan merge berdasarkan rank (prioritas merge)
+3. Decoding: daftar indeks → lookup byte sequences → gabungkan → dekode UTF-8
+
+Referensi: Sennrich et al., 2016, "Neural Machine Translation of Rare Words with Subword Units"
+"""
+
+
 def _get_stats(ids):
+    """Hitung frekuensi setiap pasangan byte yang berdekatan dalam sequence.
+
+    Args:
+        ids: Daftar integer (byte values) — representasi numerik dari teks.
+
+    Returns:
+        Dict {(id_a, id_b): count} — frekuensi setiap pasangan.
+    """
     stats = {}
     for pair in zip(ids, ids[1:]):
         stats[pair] = stats.get(pair, 0) + 1
@@ -6,6 +29,19 @@ def _get_stats(ids):
 
 
 def _merge(ids, pair, idx):
+    """Gabungkan semua kemunculan pasangan (pair[0], pair[1]) menjadi idx tunggal.
+
+    Contoh: ids=[1,2,3,1,2], pair=(1,2), idx=256 → [256,3,256]
+    Algoritma: scan linear dari kiri ke kanan, merge saat pasangan cocok.
+
+    Args:
+        ids:  Daftar integer (token IDs).
+        pair: Tuple (id_a, id_b) — pasangan yang akan digabung.
+        idx:  Integer — ID baru untuk hasil merge.
+
+    Returns:
+        Daftar integer baru dengan merge sudah diterapkan.
+    """
     out = []
     i = 0
     while i < len(ids):
@@ -19,15 +55,50 @@ def _merge(ids, pair, idx):
 
 
 class BPETokenizer:
+    """Tokenizer BPE level-byte untuk teks UTF-8.
+
+    Vocabulary dasar: 256 token pertama = semua kemungkinan byte (0–255).
+    Token 256 ke atas = hasil merge pasangan byte selama training.
+
+    Atribut:
+        merges: Daftar tuple (id_a, id_b) — urutan merge yang dipelajari (indeks = rank).
+        ranks:  Dict {(id_a, id_b): rank} — lookup cepat rank untuk pasangan.
+        vocab:  Daftar bytes — indeks → representasi byte gabungan.
+    """
+
     def __init__(self, merges=None):
+        """
+        Args:
+            merges: Daftar pasangan [(id_a, id_b), ...] atau None (kosong).
+        """
         self.merges = [tuple(p) for p in (merges or [])]
         self.ranks = {pair: i for i, pair in enumerate(self.merges)}
+        # Vocabulary dasar: 256 byte individual (0x00–0xFF)
         self.vocab = [bytes([i]) for i in range(256)]
+        # Tambahkan byte gabungan dari setiap merge
         for a, b in self.merges:
             self.vocab.append(self.vocab[a] + self.vocab[b])
 
     @classmethod
     def train(cls, text, vocab_size):
+        """Latih tokenizer BPE dari teks.
+
+        Algoritma:
+        1. Konversi teks ke daftar byte (UTF-8)
+        2. Ulangi sampai vocab_size tercapai atau tidak ada pasangan yang layak:
+           a. Hitung frekuensi semua pasangan byte
+           b. Pilih pasangan paling sering
+           c. Jika frekuensi < 2, berhenti (pasangan terlalu jarang)
+           d. Merge pasangan tersebut → tambah ID baru ke vocabulary
+           e. Catat merge di daftar merges
+
+        Args:
+            text:       Teks training (string UTF-8).
+            vocab_size: Target jumlah token di vocabulary (termasuk 256 byte dasar).
+
+        Returns:
+            BPETokenizer yang sudah terlatih.
+        """
         tok = cls()
         ids = list(text.encode("utf-8"))
         next_id = 256
@@ -35,7 +106,7 @@ class BPETokenizer:
             stats = _get_stats(ids)
             pair = max(stats, key=stats.get)
             if stats[pair] < 2:
-                break
+                break  # Tidak ada pasangan yang cukup sering muncul
             ids = _merge(ids, pair, next_id)
             tok.ranks[pair] = len(tok.merges)
             tok.merges.append(pair)
@@ -45,9 +116,23 @@ class BPETokenizer:
 
     @property
     def vocab_size(self):
+        """Jumlah total token di vocabulary (256 byte dasar + jumlah merge)."""
         return 256 + len(self.merges)
 
     def _encode_chunk(self, ids):
+        """Terapkan merge secara berulang sampai tidak ada lagi pasangan yang bisa di-merge.
+
+        Strategi: pada setiap iterasi, cari pasangan dengan rank terendah (prioritas merge
+        paling awal = paling penting), lalu merge. Ulangi sampai tidak ada pasangan lagi.
+
+        Optimasi: break lebih awal jika panjang sequence sudah sangat pendek.
+
+        Args:
+            ids: Daftar integer (byte values) dari satu chunk teks.
+
+        Returns:
+            Daftar integer (token IDs) setelah semua merge diterapkan.
+        """
         while len(ids) >= 2:
             best = None
             best_rank = None
@@ -56,20 +141,51 @@ class BPETokenizer:
                 if rank is not None and (best_rank is None or rank < best_rank):
                     best = pair
                     best_rank = rank
+                    if best_rank == 0:
+                        break  # rank 0 = prioritas tertinggi, tidak perlu cari lagi
             if best is None:
                 break
             ids = _merge(ids, best, 256 + best_rank)
         return ids
 
     def encode(self, text):
+        """Enkode teks UTF-8 menjadi daftar token ID.
+
+        Args:
+            text: String teks.
+
+        Returns:
+            Daftar integer (token IDs).
+        """
         return self._encode_chunk(list(text.encode("utf-8")))
 
     def decode(self, ids):
+        """Dekode daftar token ID kembali menjadi teks UTF-8.
+
+        Args:
+            ids: Daftar integer (token IDs).
+
+        Returns:
+            String teks (error UTF-8 diganti dengan karakter replacement).
+        """
         return b"".join(self.vocab[i] for i in ids).decode("utf-8", errors="replace")
 
     def state(self):
+        """Ekspor state tokenizer untuk disimpan ke metadata JSON.
+
+        Returns:
+            Dict {"type": "bpe", "merges": [[id_a, id_b], ...]}.
+        """
         return {"type": "bpe", "merges": [list(p) for p in self.merges]}
 
     @classmethod
     def from_state(cls, state):
+        """Muat ulang tokenizer BPE dari state (metadata JSON).
+
+        Args:
+            state: Dict {"type": "bpe", "merges": [[id_a, id_b], ...]}.
+
+        Returns:
+            BPETokenizer yang sudah dipulihkan.
+        """
         return cls(state["merges"])

indigo/common.py

@@ -1,11 +1,37 @@
+"""
+Modul utilitas bersama untuk ekosistem Indigo.
+
+Berisi fungsi-fungsi yang dipakai oleh train.py, generate.py, eval.py, dan pipeline.py:
+- Pembersihan teks (menghapus dekorasi markdown)
+- Pengumpulan file teks dari direktori secara rekursif
+- Penyimpanan/pemuatan metadata checkpoint (JSON)
+- Pembangunan tokenizer (char atau BPE) dari metadata
+- Analisis morfologi Indonesia (imbuhan + akar kata) untuk guard kamus
+"""
+
 import json
 import os
 import re
 
+# Pola regex untuk mendeteksi baris dekorasi markdown (garis-garis ===, ---, ~~~, ***)
+# Minimal 4 karakter berulang — ini bukan konten bermakna untuk training
 DECOR_LINE = re.compile(r"^[\s=\-_~*#.]{4,}$")
 
 
 def clean_text(text):
+    """Bersihkan teks mentah agar lebih cocok untuk training.
+
+    Langkah:
+    1. Hapus baris dekorasi markdown (====, ----, ~~~~, dst.)
+    2. Gabungkan 3+ baris kosong berturut menjadi 2 baris kosong (satu paragraf kosong)
+    3. Strip spasi di awal/akhir, tambahkan newline di akhir
+
+    Args:
+        text: Teks mentah dari file.
+
+    Returns:
+        Teks yang sudah dibersihkan.
+    """
     lines = [ln for ln in text.splitlines() if not DECOR_LINE.match(ln)]
     text = "\n".join(lines)
     text = re.sub(r"\n{3,}", "\n\n", text)
@@ -13,6 +39,17 @@ def clean_text(text):
 
 
 def collect_text_files(paths):
+    """Kumpulkan semua file .txt dari daftar path (file atau direktori).
+
+    Jika path adalah direktori, cari rekursif semua file .txt di dalamnya.
+    Jika path adalah file, langsung tambahkan (tidak harus .txt).
+
+    Args:
+        paths: Daftar string path (file atau direktori).
+
+    Returns:
+        Daftar string path file yang sudah diurutkan.
+    """
     files = []
     for p in paths:
         if os.path.isdir(p):
@@ -24,11 +61,32 @@ def collect_text_files(paths):
 
 
 def read_clean(path):
+    """Baca file teks dan bersihkan isinya dengan clean_text().
+
+    Args:
+        path: Path ke file teks (encoding UTF-8).
+
+    Returns:
+        Teks yang sudah dibersihkan.
+    """
     with open(path, encoding="utf-8") as f:
         return clean_text(f.read())
 
 
 def save_meta(base_path, config, vocab, step, val_loss, backend, tokenizer=None):
+    """Simpan metadata checkpoint ke file JSON.
+
+    File yang dihasilkan: ganti suffix base_path menjadi _meta.json.
+    Contoh: out/indigo_best.safetensors → out/indigo_best_meta.json
+
+    Isi metadata:
+        - config:    Dict hyperparameter model (vocab_size, n_layer, n_embd, dst.)
+        - vocab:     Daftar karakter (untuk char tokenizer) atau None (untuk BPE)
+        - step:      Nomor step training saat checkpoint disimpan
+        - val_loss:  Loss validasi terbaik (atau None jika tidak ada validasi)
+        - backend:   "pytorch" atau "tensorflow"
+        - tokenizer: Dict info tokenizer {"type": "char"} atau {"type": "bpe", "merges": [...]}
+    """
     meta = {
         "config": config,
         "vocab": vocab,
@@ -42,11 +100,35 @@ def save_meta(base_path, config, vocab, step, val_loss, backend, tokenizer=None)
 
 
 def load_meta(path):
+    """Muat metadata checkpoint dari file _meta.json.
+
+    File yang dicari: ganti suffix path menjadi _meta.json.
+    Contoh: out/indigo_best.safetensors → cari out/indigo_best_meta.json
+
+    Args:
+        path: Path ke file checkpoint (.safetensors atau .pt).
+
+    Returns:
+        Dict metadata berisi config, vocab, step, val_loss, backend, tokenizer.
+    """
     with open(os.path.splitext(path)[0] + "_meta.json", encoding="utf-8") as f:
         return json.load(f)
 
 
 def build_tokenizer(tokenizer_info, vocab):
+    """Bangun objek tokenizer dari metadata checkpoint.
+
+    Dispatch berdasarkan type:
+    - "bpe":  Muat BPETokenizer dari daftar merges
+    - lainnya: Muat CharTokenizer dari daftar karakter (vocab)
+
+    Args:
+        tokenizer_info: Dict {"type": "char"} atau {"type": "bpe", "merges": [...]}.
+        vocab:          Daftar karakter (list of str) atau None.
+
+    Returns:
+        Objek tokenizer (CharTokenizer atau BPETokenizer).
+    """
     if tokenizer_info.get("type") == "bpe":
         from indigo.bpe import BPETokenizer
 
@@ -56,25 +138,60 @@ def build_tokenizer(tokenizer_info, vocab):
     return CharTokenizer(vocab)
 
 
+# --- Analisis morfologi Indonesia untuk guard kamus ---
+
+# Pola regex untuk mengekstrak kata alfanumerik dari teks
+# Mencakup karakter Latin + aksen (á, é, dst.) + tanda petik/strip tengah
 WORD_RE = re.compile(r"[a-zA-Z\xC0-\xFF][a-zA-Z\xC0-\xFF'-]+")
 
 
 def load_wordlist(path):
+    """Muat daftar kata dari file (satu kata per baris).
+
+    Mengembalikan frozenset kata lowercase untuk pencarian O(1).
+
+    Args:
+        path: Path ke file teks (satu kata per baris, UTF-8).
+
+    Returns:
+        frozenset of str (kata lowercase).
+    """
     with open(path, encoding="utf-8") as f:
         return frozenset(w.strip().lower() for w in f if w.strip())
 
 
+# Peta asimilasi prefiks: beberapa prefiks berubah bunyi saat ditambahkan ke akar
+# Contoh: "menyapu" → hapus "meny" → "apu", tapi akarnya "sapu" (ada asimilasi s→s)
 PREFIKS_ASIMILASI = {"meny": "s", "peny": "s", "pem": "p"}
 
 
 def _kandidat_akar(word, prefiks, sufiks):
+    """Cari kemungkinan akar kata dari kata berimbuhan (stems candidates).
+
+    Algoritma iteratif (max 3 putaran):
+    1. Coba hapus setiap sufiks → simpan hasilnya
+    2. Coba hapus setiap prefiks → simpan hasilnya (+ asimilasi)
+    3. Ulangi pada hasil baru sampai tidak ada kandidat baru
+
+    Kandidat hanya disimpan jika panjang sisa >= 3 karakter (mencegah akar terlalu pendek).
+
+    Args:
+        word:    Kata yang akan dianalisis (string).
+        prefiks: Set of str — daftar prefiks (mis. {"men", "mem", "di", "ter"}).
+        sufiks:  Set of str — daftar sufiks (mis. {"kan", "an", "i", "nya"}).
+
+    Returns:
+        Set of str — semua kemungkinan akar kata.
+    """
     sisa = {word}
     for _ in range(3):
         lanjut = set()
         for w in sisa:
+            # Coba hapus sufiks
             for suf in sorted(sufiks, key=len, reverse=True):
                 if w.endswith(suf) and len(w) - len(suf) >= 3:
                     lanjut.add(w[: -len(suf)])
+            # Coba hapus prefiks (+ asimilasi)
             for pref in sorted(prefiks, key=len, reverse=True):
                 if w.startswith(pref) and len(w) - len(pref) >= 3:
                     akar = PREFIKS_ASIMILASI.get(pref, "") + w[len(pref):]
@@ -88,6 +205,24 @@ def _kandidat_akar(word, prefiks, sufiks):
 
 
 def word_known_ratio(text, wordset, prefiks=None, sufiks=None):
+    """Hitung rasio kata yang dikenal kamus dari sebuah teks.
+
+    Tahapan:
+    1. Ekstrak semua kata alfanumerik dari teks (min 2 huruf, hanya alpha)
+    2. Untuk setiap kata:
+       - Cek langsung di wordset → jika ada, dikenal
+       - Jika tidak ada dan afiks tersedia, cek semua kandidat akar → jika ada, dikenal
+    3. return jumlah_dikenal / total_kata
+
+    Args:
+        text:    Teks yang akan dievaluasi.
+        wordset: frozenset of str — daftar kata yang dikenal (lowercase).
+        prefiks: Set of str atau None — daftar prefiks Indonesia.
+        sufiks:  Set of str atau None — daftar sufiks Indonesia.
+
+    Returns:
+        Float 0.0–1.0 — rasio kata yang dikenal kamus.
+    """
     words = [m.group(0).strip("'-").lower() for m in WORD_RE.finditer(text)]
     words = [w for w in words if len(w) >= 2 and w.isalpha()]
     if not words:

indigo/model.py

@@ -1,3 +1,15 @@
+"""
+Arsitektur GPT decoder-only transformer untuk Indigo.
+
+Komponen utama:
+- CausalSelfAttention: self-attention kausal dengan SDPA (Scaled Dot-Product Attention)
+  dan KV-cache untuk generasi cepat token-by-token.
+- MLP: feed-forward network dengan aktivasi GELU.
+- Block: satu blok transformer = LayerNorm + Attention + LayerNorm + MLP (pre-LN).
+- GPT: model lengkap = token embedding + positional embedding + blok-blok + output head.
+- Fungsi sampling: top-k, top-p, temperature, repetition penalty.
+"""
+
 import math
 from dataclasses import dataclass
 
@@ -8,6 +20,17 @@ import torch.nn.functional as F
 
 @dataclass
 class GPTConfig:
+    """Konfigurasi hyperparameter model GPT.
+
+    Attributes:
+        vocab_size:  Jumlah token unik di vocabulary (termasuk semua byte untuk BPE).
+        block_size:  Panjang konteks maksimum (jumlah token yang dilihat model sekaligus).
+        n_layer:     Jumlah blok transformer bertumpuk.
+        n_head:      Jumlah head per attention layer (harus habis membagi n_embd).
+        n_embd:      Dimensi embedding / hidden size setiap layer.
+        dropout:     Rate dropout (0.0 = tanpa dropout).
+        bias:        Gunakan bias di Linear dan LayerNorm (GPT-2 asli pakai False).
+    """
     vocab_size: int = 256
     block_size: int = 128
     n_layer: int = 4
@@ -18,43 +41,79 @@ class GPTConfig:
 
 
 class CausalSelfAttention(nn.Module):
+    """Multi-head causal self-attention dengan dua mode:
+
+    1. Forward pass biasa (kv=None): menggunakan SDPA dari PyTorch untuk
+       training yang efisien — mask kausal diterapkan otomatis.
+    2. Forward pass dengan KV-cache (kv=(k,v)): menggabungkan K/V baru dengan
+       K/V dari token sebelumnya, lalu menghitung attention secara manual
+       dengan mask kausal — digunakan saat generasi token-by-token.
+    """
+
     def __init__(self, config):
         super().__init__()
-        assert config.n_embd % config.n_head == 0
+        assert config.n_embd % config.n_head == 0, "n_embd harus habis dibagi n_head"
         self.n_head = config.n_head
         self.n_embd = config.n_embd
+        # Q, K, V diproyeksikan sekaligus dalam satu Linear untuk efisiensi
         self.qkv = nn.Linear(config.n_embd, 3 * config.n_embd, bias=config.bias)
+        # Proyeksi output setelah attention
         self.proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias)
         self.attn_dropout = nn.Dropout(config.dropout)
         self.resid_dropout = nn.Dropout(config.dropout)
 
     def forward(self, x, kv=None):
+        """
+        Args:
+            x:   Tensor (B, T, C) — batch, sequence length, channels.
+            kv:  Tuple (K_prev, V_prev) dari token sebelumnya, atau None.
+                 Saat None → training mode (SDPA). Saat ada → generasi (manual attention).
+
+        Returns:
+            y:        Tensor (B, T, C) — output setelah attention + dropout.
+            (k, v):   Tuple K, V terbaru untuk disimpan ke cache.
+        """
         B, T, C = x.shape
+        # Proyeksi Q, K, V dari input
         q, k, v = self.qkv(x).split(self.n_embd, dim=2)
+        # Reshape ke (B, n_head, T, head_dim) untuk multi-head attention
         q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
         k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
         v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
+
         if kv is None:
+            # Mode training: SDPA dari PyTorch — mask kausal otomatis
             y = F.scaled_dot_product_attention(
                 q, k, v,
                 dropout_p=self.attn_dropout.p if self.training else 0.0,
                 is_causal=True,
             )
         else:
+            # Mode generasi: gabungkan K/V baru dengan cache sebelumnya
             pk, pv = kv
-            k = torch.cat((pk, k), dim=2)
+            k = torch.cat((pk, k), dim=2)  # (B, n_head, T_prev+T, head_dim)
             v = torch.cat((pv, v), dim=2)
+            # Buat mask kausal manual: query baru hanya boleh melihat key yang valid
             Tq, Tk = q.size(2), k.size(2)
             mask = torch.ones(Tq, Tk, dtype=torch.bool, device=x.device).tril(diagonal=Tk - Tq)
+            # Hitung attention score secara manual
             att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
             att = att.masked_fill(~mask, float("-inf"))
             att = F.softmax(att, dim=-1)
             y = att @ v
+
+        # Gabungkan kembali head-head: (B, n_head, T, head_dim) → (B, T, C)
         y = y.transpose(1, 2).contiguous().view(B, T, C)
         return self.resid_dropout(self.proj(y)), (k, v)
 
 
 class MLP(nn.Module):
+    """Feed-forward network per blok transformer.
+
+    Arsitektur: Linear(C → 4C) → GELU → Linear(4C → C) → Dropout.
+    Rasio 4x adalah standar GPT-2 (expansion factor).
+    """
+
     def __init__(self, config):
         super().__init__()
         self.fc = nn.Linear(config.n_embd, 4 * config.n_embd, bias=config.bias)
@@ -66,6 +125,14 @@ class MLP(nn.Module):
 
 
 class Block(nn.Module):
+    """Satu blok transformer (pre-LayerNorm).
+
+    Arsitektur:
+        x → LayerNorm → Attention → (+ residual) → LayerNorm → MLP → (+ residual) → output
+
+    Pre-LN (LayerNorm di depan) lebih stabil untuk training dari nol dibanding Post-LN.
+    """
+
     def __init__(self, config):
         super().__init__()
         self.ln1 = nn.LayerNorm(config.n_embd, bias=config.bias)
@@ -74,13 +141,39 @@ class Block(nn.Module):
         self.mlp = MLP(config)
 
     def forward(self, x, kv=None):
+        """
+        Args:
+            x:   Input tensor (B, T, C).
+            kv:  KV-cache dari blok sebelumnya (opsional).
+
+        Returns:
+            x:       Output tensor (B, T, C).
+            kv_new:  (K, V) baru untuk disimpan ke cache.
+        """
         a, kv_new = self.attn(self.ln1(x), kv)
-        x = x + a
-        x = x + self.mlp(self.ln2(x))
+        x = x + a          # residual connection (attention)
+        x = x + self.mlp(self.ln2(x))  # residual connection (MLP)
         return x, kv_new
 
 
 def _sample_token(logits, temperature, top_k=None, top_p=None):
+    """Sampling satu token dari distribusi logit.
+
+    Pipeline sampling:
+    1. Scaling temperature: logits / T (T kecil → lebih deterministik, T besar → lebih random)
+    2. Top-k: buang semua kandidat kecuali k teratas
+    3. Top-p (nucleus): buang kandidat yang kumulatif probabilitasnya > p
+    4. Softmax → multinomial sampling
+
+    Args:
+        logits:      Tensor (B, vocab_size) — skor sebelum softmax.
+        temperature: Float — skala randomness (0.0 ≈ greedy, 1.0 ≈ standar, >1.0 ≈ random).
+        top_k:       Int atau None — batasi ke k token teratas.
+        top_p:       Float atau None — batasi kumulatif probabilitas sampai p.
+
+    Returns:
+        Token sampled: Tensor (B, 1).
+    """
     logits = logits / max(temperature, 1e-8)
     if top_k is not None:
         v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
@@ -97,6 +190,23 @@ def _sample_token(logits, temperature, top_k=None, top_p=None):
 
 
 def _apply_repetition_penalty(logits, tokens, penalty, window):
+    """Terapkan repetition penalty pada token yang baru saja muncul.
+
+    Token yang sudah muncul dalam `window` token terakhir akan didenda:
+    - Jika logit positif → dibagi penalty (menjadi lebih kecil)
+    - Jika logit negatif → dikali penalty (menjadi lebih negatif)
+
+    Efek: mengurangi kemungkinan model mengulang kata/frasa yang sama.
+
+    Args:
+        logits:  Tensor (1, vocab_size) — logit kandidat.
+        tokens:  Tensor (1, T) — token yang sudah dihasilkan.
+        penalty: Float > 1.0 — semakin besar, semakin keras penaltinya.
+        window:  Int — berapa token terakhir yang diperiksa.
+
+    Returns:
+        logits yang sudah diubah in-place.
+    """
     recent = set(tokens[0, -window:].tolist())
     for t in recent:
         val = logits[0, t]
@@ -105,19 +215,48 @@ def _apply_repetition_penalty(logits, tokens, penalty, window):
 
 
 class GPT(nn.Module):
+    """Model GPT decoder-only transformer.
+
+    Komponen:
+        - tok_emb:   Embedding token (vocab_size → n_embd)
+        - pos_emb:   Embedding posisi (block_size → n_embd) — learnable, bukan sinusoidal
+        - drop:      Dropout gabungan (token + posisi)
+        - blocks:    N blok transformer (pre-LN)
+        - ln_f:      LayerNorm final
+        - head:      Linear(n_embd → vocab_size) — bobot ditied dengan tok_emb (weight tying)
+
+    Forward pass:
+        1. Embed token + posisi → x
+        2. x melewati setiap blok (attention + MLP) → logits
+        3. Hitung loss (cross-entropy) jika targets tersedia
+
+    Generasi (generate):
+        1. Forward pass pertama: proses seluruh prompt sekaligus, simpan KV-cache
+        2. Forward pass selanjutnya: proses hanya 1 token baru, gunakan KV-cache
+        3. Cache di-reset jika sudah mencapai block_size (context window penuh)
+    """
+
     def __init__(self, config):
         super().__init__()
         self.config = config
+        # Token embedding: mengubah indeks token menjadi vektor
         self.tok_emb = nn.Embedding(config.vocab_size, config.n_embd)
+        # Positional embedding: memberi tahu model posisi setiap token dalam sequence
         self.pos_emb = nn.Embedding(config.block_size, config.n_embd)
         self.drop = nn.Dropout(config.dropout)
+        # Blok-blok transformer bertumpuk
         self.blocks = nn.ModuleList([Block(config) for _ in range(config.n_layer)])
+        # LayerNorm final sebelum output head
         self.ln_f = nn.LayerNorm(config.n_embd, bias=config.bias)
+        # Output head: proyeksi dari hidden state ke vocabulary
         self.head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
+        # Weight tying: bobot output head = bobot token embedding (hemat parameter)
         self.head.weight = self.tok_emb.weight
+        # Inisialisasi bobot dengan distribusi normal kecil
         self.apply(self._init_weights)
 
     def _init_weights(self, module):
+        """Inisialisasi bobot: Normal(0, 0.02) untuk Linear dan Embedding."""
         if isinstance(module, nn.Linear):
             nn.init.normal_(module.weight, mean=0.0, std=0.02)
             if module.bias is not None:
@@ -126,14 +265,31 @@ class GPT(nn.Module):
             nn.init.normal_(module.weight, mean=0.0, std=0.02)
 
     def forward(self, idx, targets=None, caches=None, return_caches=False):
+        """
+        Forward pass.
+
+        Args:
+            idx:           Tensor (B, T) — indeks token input.
+            targets:       Tensor (B, T) atau None — indeks token target (untuk hitung loss).
+            caches:        List of (K, V) per blok atau None — KV-cache dari token sebelumnya.
+            return_caches: bool — jika True, return (logits, loss, new_caches).
+
+        Returns:
+            logits: Tensor (B, T, vocab_size)
+            loss:   Float atau None (None jika targets tidak diberikan)
+        """
         B, T = idx.shape
+        # Posisi dimulai dari panjang cache (offset untuk generasi autoregressive)
         start = caches[0][0].size(2) if caches else 0
         pos = torch.arange(start, start + T, device=idx.device)
+        # Embedding = token embedding + positional embedding, lalu dropout
         x = self.drop(self.tok_emb(idx) + self.pos_emb(pos))
+        # Forward pass melalui setiap blok, kumpulkan KV-cache baru
         new_caches = []
         for i, block in enumerate(self.blocks):
             x, c = block(x, caches[i] if caches is not None else None)
             new_caches.append(c)
+        # LayerNorm final → output logits
         logits = self.head(self.ln_f(x))
         loss = None
         if targets is not None:
@@ -152,27 +308,52 @@ class GPT(nn.Module):
         top_p=None,
         repetition_penalty=1.0,
     ):
+        """Generate token baru secara autoregressive menggunakan KV-cache.
+
+        Proses:
+        1. Forward pass awal: proses seluruh idx (prompt) sekaligus → simpan KV-cache.
+        2. Loop max_new_tokens kali: ambil logit terakhir → sample → tambah ke sequence.
+        3. Jika cache penuh (>= block_size), reset cache dan proses ulang block_size terakhir.
+
+        Args:
+            idx:                Tensor (B, T) — token prompt awal.
+            max_new_tokens:     Int — jumlah token baru yang akan dihasilkan.
+            temperature:        Float — skala randomness sampling.
+            top_k:              Int atau None — batasi kandidat ke k token teratas.
+            top_p:              Float atau None — batasi kandidat berdasarkan kumulatif probabilitas.
+            repetition_penalty: Float — penalti pengulangan (>1.0 = aktif, 1.0 = nonaktif).
+
+        Returns:
+            Tensor (B, T + max_new_tokens) — prompt + token hasil generate.
+        """
         self.eval()
+        # Potong prompt jika lebih panjang dari block_size
         idx_cond = idx[:, -self.config.block_size:]
+        # Forward pass pertama: proses seluruh prompt, simpan KV-cache
         logits, _, caches = self(idx_cond, return_caches=True)
         tokens = idx_cond
         for i in range(max_new_tokens):
+            # Ambil logit untuk prediksi token berikutnya (posisi terakhir)
             next_logits = logits[:, -1, :]
             if repetition_penalty != 1.0:
                 next_logits = _apply_repetition_penalty(
                     next_logits.clone(), tokens, repetition_penalty, self.config.block_size
                 )
+            # Sample satu token berdasarkan logit + temperature + top-k + top-p
             next_id = _sample_token(next_logits, temperature, top_k, top_p)
             tokens = torch.cat((tokens, next_id), dim=1)
             if i == max_new_tokens - 1:
                 break
+            # Jika cache penuh, reset dan proses ulang block_size terakhir
             if caches is not None and caches[0][0].size(2) >= self.config.block_size:
                 caches = None
                 idx_in = tokens[:, -self.config.block_size:]
             else:
+                # Forward pass untuk 1 token baru saja, gunakan cache sebelumnya
                 idx_in = next_id
             logits, _, caches = self(idx_in, caches=caches, return_caches=True)
         return tokens
 
     def num_params(self):
+        """Hitung total parameter model (termasuk yang tidak di-train seperti embedding tied)."""
         return sum(p.numel() for p in self.parameters())

indigo/tokenizer.py

@@ -1,8 +1,27 @@
+"""
+Tokenizer karakter (char-level) untuk Indigo.
+
+Setiap karakter unik dalam teks training mendapatkan satu indeks token.
+Sederhana dan cepat, tapi kurang efisien untuk teks panjang (tiap karakter = 1 token).
+Cocok untuk dataset kecil atau percobaan cepat.
+"""
+
 import json
 
 
 class CharTokenizer:
+    """Tokenizer level karakter: mapping satu-satu antara karakter dan indeks.
+
+    Atribut:
+        itos: List of str — indeks → karakter (integer-to-string).
+        stoi: Dict {karakter: indeks} — lookup cepat encode (string-to-integer).
+    """
+
     def __init__(self, vocab=None):
+        """
+        Args:
+            vocab: Daftar karakter (list of str) atau None (tokenizer kosong).
+        """
         if vocab is None:
             self.itos = []
         else:
@@ -11,23 +30,65 @@ class CharTokenizer:
 
     @classmethod
     def from_text(cls, text):
+        """Bangun tokenizer dari teks: setiap karakter unik jadi satu token.
+
+        Karakter diurutkan agar mapping deterministik (tidak bergantung urutan kemunculan).
+
+        Args:
+            text: Teks training (string).
+
+        Returns:
+            CharTokenizer yang sudah terlatih.
+        """
         return cls(sorted(set(text)))
 
     @property
     def vocab_size(self):
+        """Jumlah total karakter unik di vocabulary."""
         return len(self.itos)
 
     def encode(self, text):
+        """Enkode teks menjadi daftar indeks token.
+
+        Karakter yang tidak ada di vocabulary diabaikan (tidak di-encode).
+
+        Args:
+            text: String teks.
+
+        Returns:
+            Daftar integer (token IDs).
+        """
         return [self.stoi[ch] for ch in text if ch in self.stoi]
 
     def decode(self, ids):
+        """Dekode daftar indeks token kembali menjadi teks.
+
+        Args:
+            ids: Daftar integer (token IDs).
+
+        Returns:
+            String teks.
+        """
         return "".join(self.itos[i] for i in ids)
 
     def save(self, path):
+        """Simpan vocabulary (daftar karakter) ke file JSON.
+
+        Args:
+            path: Path output file JSON.
+        """
         with open(path, "w", encoding="utf-8") as f:
             json.dump(self.itos, f, ensure_ascii=False)
 
     @classmethod
     def load(cls, path):
+        """Muat vocabulary dari file JSON.
+
+        Args:
+            path: Path ke file JSON berisi daftar karakter.
+
+        Returns:
+            CharTokenizer yang sudah dipulihkan.
+        """
         with open(path, encoding="utf-8") as f:
             return cls(json.load(f))

pipeline.py

@@ -1,6 +1,24 @@
+"""
+Pipeline end-to-end untuk training Indigo.
+
+Alur kerja:
+1. Kumpulkan data: copy file lokal + tarik dataset dari Hugging Face
+2. Konversi JSON Alpaca ke .txt (opsional, --format-qa)
+3. Jalankan training (train.py) dengan semua hyperparameter
+4. Evaluasi guard kamus (opsional, --guard)
+5. Simpan manifest.json (statistik + metadata run)
+6. Finalize: promosikan checkpoint terbaik ke out/ (opsional, --finalize)
+7. Push ke Hugging Face Hub (opsional, --push)
+
+Cara pakai:
+    python pipeline.py --tag run01 --data data/sample.txt --steps 2000
+    python pipeline.py --tag run02 --hf-dataset adyoi/indigo --steps 900 --device cpu
+    python pipeline.py --tag run03 --data data/ --steps 2000 --finalize --push
+    python pipeline.py --tag qa01 --hf-dataset rohanrdy/CS-Theory-QA-Dataset --format-qa --steps 3000
+"""
+
 import argparse
 import datetime
-import glob
 import json
 import os
 import shutil
@@ -15,6 +33,25 @@ import train
 
 
 def preprocess_alpaca_json(data_dir, patterns=("*.json",)):
+    """Konversi file JSON Alpaca format ke .txt flat.
+
+    Format Alpaca: [{"instruction": "...", "input": "...", "output": "..."}, ...]
+    Format output:
+        ### instruction
+        (input jika ada)
+
+        output
+
+    Fungsi ini auto-detect: hanya mengkonversi file JSON yang punya field "instruction" dan "output".
+    File yang sudah ada .txt-nya tidak ditimpa.
+
+    Args:
+        data_dir:  Direktori yang berisi file JSON (akan dicari rekursif).
+        patterns:  Pola file JSON (default: "*.json").
+
+    Returns:
+        Int — jumlah file JSON yang berhasil dikonversi.
+    """
     converted = 0
     for pat in patterns:
         for jf in Path(data_dir).rglob(pat):
@@ -27,6 +64,7 @@ def preprocess_alpaca_json(data_dir, patterns=("*.json",)):
                 continue
             if not isinstance(data, list) or not data:
                 continue
+            # Cek apakah ini format Alpaca (punya field instruction + output)
             keys = set(data[0].keys()) if isinstance(data[0], dict) else set()
             if not {"instruction", "output"}.issubset(keys):
                 continue
@@ -41,6 +79,7 @@ def preprocess_alpaca_json(data_dir, patterns=("*.json",)):
                 lines.append(f"### {q}\n\n{out}\n\n")
             if not lines:
                 continue
+            # Simpan .txt di sebelah file JSON asli
             txt_path = jf.with_suffix(".txt")
             txt_path.write_text("".join(lines), encoding="utf-8")
             print(f"[pipeline] alpaca: {jf.name} -> {txt_path.name} ({len(lines)} QA)")
@@ -49,6 +88,7 @@ def preprocess_alpaca_json(data_dir, patterns=("*.json",)):
 
 
 def git_short():
+    """Ambil hash commit git pendek (7 karakter) atau "nogit" jika bukan repo git."""
     try:
         return subprocess.check_output(
             ["git", "rev-parse", "--short", "HEAD"], cwd=str(ROOT), text=True
@@ -58,6 +98,19 @@ def git_short():
 
 
 def pull_hf(repo_id, dest, patterns):
+    """Tarik dataset dari Hugging Face Hub ke direktori lokal.
+
+    Menggunakan snapshot_download dengan allow_patterns untuk mengunduh
+    hanya file yang sesuai pola (mis. *.txt, *.md, *.json).
+
+    Args:
+        repo_id:  ID repo HF (mis. "adyoi/indigo" atau "rohanrdy/CS-Theory-QA-Dataset").
+        dest:     Direktori tujuan download.
+        patterns: Pola file yang akan diunduh.
+
+    Returns:
+        Str — path ke direktori lokal yang berisi file yang diunduh.
+    """
     from huggingface_hub import snapshot_download
 
     path = snapshot_download(
@@ -79,47 +132,87 @@ def main():
     ap = argparse.ArgumentParser(
         description="Pipeline latih Indigo: kumpul data -> tokenizer -> training -> manifest -> finalize/push"
     )
-    ap.add_argument("--tag", required=True, help="nama run; artefak di <runs>/<tag>/")
-    ap.add_argument("--data", nargs="*", default=[], help="file/folder teks lokal tambahan")
-    ap.add_argument("--hf-dataset", default=None, help="repo dataset HF untuk menarik file teks")
-    ap.add_argument("--hf-patterns", nargs="*", default=["*.txt", "*.md"])
+
+    # --- Identitas Run ---
+    ap.add_argument("--tag", required=True,
+                    help="nama run; semua artefak disimpan di runs/<tag>/")
+
+    # --- Sumber Data ---
+    ap.add_argument("--data", nargs="*", default=[],
+                    help="file/folder teks lokal tambahan (banyak, spasi-separated)")
+    ap.add_argument("--hf-dataset", default=None,
+                    help="repo dataset HF untuk menarik file teks (mis. adyoi/indigo)")
+    ap.add_argument("--hf-patterns", nargs="*", default=["*.txt", "*.md"],
+                    help="pola file yang diambil dari HF (default: *.txt *.md)")
     ap.add_argument("--format-qa", action="store_true",
-                    help="auto-convert Alpaca JSON (instruction/output) to .txt sebelum training")
-    ap.add_argument("--runs", default="runs")
-    ap.add_argument("--device", default=None, help="diteruskan ke train.py (default auto)")
+                    help="auto-convert Alpaca JSON (instruction/output) ke .txt sebelum training")
+
+    # --- Output & Finalize ---
+    ap.add_argument("--runs", default="runs",
+                    help="folder root untuk semua run (default: runs)")
+    ap.add_argument("--device", default=None,
+                    help="device untuk training (diteruskan ke train.py, default: auto)")
     ap.add_argument("--finalize", action="store_true",
                     help="promosikan checkpoint terbaik run ini ke folder out/ kanonik")
-    ap.add_argument("--push", action="store_true", help="upload checkpoint terbaik ke repo HF")
-    ap.add_argument("--repo", default="adyoi/indigo")
+    ap.add_argument("--push", action="store_true",
+                    help="upload checkpoint terbaik ke repo HF")
+    ap.add_argument("--repo", default="adyoi/indigo",
+                    help="repo HF tujuan upload (default: adyoi/indigo)")
+
+    # --- Guard Kamus ---
     ap.add_argument("--guard", default=None,
                     help="file kamus kata (satu kata per baris); tambahkan metrik rasio ejaan ke manifest")
-    ap.add_argument("--guard-max-new", type=int, default=120)
-    ap.add_argument("--guard-prefiks", default=None, help="default data/prefiks.txt bila ada")
-    ap.add_argument("--guard-sufiks", default=None, help="default data/sufiks.txt bila ada")
+    ap.add_argument("--guard-max-new", type=int, default=120,
+                    help="jumlah token generate untuk evaluasi guard (default: 120)")
+    ap.add_argument("--guard-prefiks", default=None,
+                    help="file prefiks Indonesia (default: data/prefiks.txt bila ada)")
+    ap.add_argument("--guard-sufiks", default=None,
+                    help="file sufiks Indonesia (default: data/sufiks.txt bila ada)")
+
+    # --- Hyperparameter Training (diteruskan ke train.py) ---
     t = ap.add_argument_group("hyperparameter training (diteruskan ke train.py)")
-    t.add_argument("--steps", type=int, default=2000)
-    t.add_argument("--batch-size", type=int, default=32)
-    t.add_argument("--block-size", type=int, default=128)
-    t.add_argument("--n-layer", type=int, default=4)
-    t.add_argument("--n-head", type=int, default=4)
-    t.add_argument("--n-embd", type=int, default=128)
-    t.add_argument("--dropout", type=float, default=0.1)
-    t.add_argument("--lr", type=float, default=3e-4)
-    t.add_argument("--warmup", type=int, default=100)
-    t.add_argument("--weight-decay", type=float, default=0.1)
-    t.add_argument("--eval-interval", type=int, default=200)
-    t.add_argument("--eval-iters", type=int, default=20)
-    t.add_argument("--seed", type=int, default=1337)
-    t.add_argument("--init-from", default=None)
-    t.add_argument("--tokenizer", default="char", choices=["char", "bpe"])
-    t.add_argument("--vocab-size", type=int, default=512)
-    t.add_argument("--val-fraction", type=float, default=0.1)
+    t.add_argument("--steps", type=int, default=2000,
+                   help="jumlah langkah training (default: 2000)")
+    t.add_argument("--batch-size", type=int, default=32,
+                   help="batch size (default: 32)")
+    t.add_argument("--block-size", type=int, default=128,
+                   help="panjang konteks (default: 128)")
+    t.add_argument("--n-layer", type=int, default=4,
+                   help="jumlah layer transformer (default: 4)")
+    t.add_argument("--n-head", type=int, default=4,
+                   help="jumlah head attention (default: 4)")
+    t.add_argument("--n-embd", type=int, default=128,
+                   help="dimensi embedding (default: 128)")
+    t.add_argument("--dropout", type=float, default=0.1,
+                   help="dropout rate (default: 0.1)")
+    t.add_argument("--lr", type=float, default=3e-4,
+                   help="learning rate (default: 3e-4)")
+    t.add_argument("--warmup", type=int, default=100,
+                   help="langkah warmup (default: 100)")
+    t.add_argument("--weight-decay", type=float, default=0.1,
+                   help="weight decay (default: 0.1)")
+    t.add_argument("--eval-interval", type=int, default=200,
+                   help="evaluasi setiap N langkah (default: 200)")
+    t.add_argument("--eval-iters", type=int, default=20,
+                   help="jumlah batch evaluasi (default: 20)")
+    t.add_argument("--seed", type=int, default=1337,
+                   help="seed random (default: 1337)")
+    t.add_argument("--init-from", default=None,
+                   help="checkpoint untuk resume training")
+    t.add_argument("--tokenizer", default="char", choices=["char", "bpe"],
+                   help="jenis tokenizer (default: char)")
+    t.add_argument("--vocab-size", type=int, default=512,
+                   help="vocab size untuk BPE (default: 512)")
+    t.add_argument("--val-fraction", type=float, default=0.1,
+                   help="proporsi file validasi (default: 0.1)")
     args = ap.parse_args()
 
+    # --- Siapkan direktori run ---
     run_dir = ROOT / args.runs / args.tag
     data_dir = run_dir / "data"
     data_dir.mkdir(parents=True, exist_ok=True)
 
+    # --- Kumpulkan data lokal ---
     sources = []
     for p in args.data:
         src = Path(p)
@@ -131,13 +224,15 @@ def main():
             shutil.copy2(src, data_dir / src.name)
         sources.append(str(src))
 
+    # --- Tarik dataset dari HF ---
     if args.hf_dataset:
         patterns = list(args.hf_patterns)
         if args.format_qa and "*.json" not in patterns:
-            patterns.append("*.json")
+            patterns.append("*.json")  # butuh JSON untuk Alpaca conversion
         pull_hf(args.hf_dataset, data_dir / "_hf", patterns)
         sources.append(f"hf:{args.hf_dataset}")
 
+    # --- Konversi JSON Alpaca → .txt ---
     if args.format_qa:
         n = preprocess_alpaca_json(data_dir)
         if n:
@@ -145,6 +240,7 @@ def main():
         else:
             print("[pipeline] peringatan: tidak ada file Alpaca JSON ditemukan")
 
+    # --- Jalankan training ---
     out_dir = run_dir / "ckpt"
     argv = [
         "--data", str(data_dir),
@@ -174,6 +270,7 @@ def main():
     print(f"[pipeline] run='{args.tag}' | sumber={sources or ['<kosong>' ]}")
     stats = train.main(argv) or {}
 
+    # --- Evaluasi guard kamus (opsional) ---
     if args.guard:
         import torch
 
@@ -181,6 +278,7 @@ def main():
         from indigo.model import GPT, GPTConfig
         from safetensors.torch import load_file
 
+        # Cari checkpoint: prioritaskan best, fallback ke final
         probe_path = out_dir / "indigo_best.safetensors"
         if not probe_path.exists():
             probe_path = out_dir / "indigo.safetensors"
@@ -188,6 +286,7 @@ def main():
         p_def, s_def = ROOT / "data" / "prefiks.txt", ROOT / "data" / "sufiks.txt"
 
         def muat_afiks(flag, default_path):
+            """Muat file afiks: prioritaskan flag CLI → default path."""
             if flag:
                 return load_wordlist(flag) if Path(flag).exists() else None
             return load_wordlist(str(default_path)) if default_path.exists() else None
@@ -199,6 +298,7 @@ def main():
         mdl = GPT(GPTConfig(**meta["config"]))
         mdl.load_state_dict(load_file(str(probe_path)), strict=False)
         mdl.eval()
+        # Generate teks dari prompt newline (netral)
         ids = tok.encode("\n") or [0]
         idx = torch.tensor([ids], dtype=torch.long)
         with torch.no_grad():
@@ -206,6 +306,8 @@ def main():
         text = tok.decode(out[0].tolist())
         stats["kamus_ratio"] = round(word_known_ratio(text, wordset, prefiks, sufiks), 4)
         print(f"[pipeline] kamus_ratio={stats['kamus_ratio']:.2%}")
+
+    # --- Simpan manifest.json ---
     stats["pipeline"] = {
         "tag": args.tag,
         "timestamp": datetime.datetime.now().isoformat(timespec="seconds"),
@@ -220,6 +322,7 @@ def main():
         json.dump(stats, f, ensure_ascii=False, indent=2)
     print(f"[pipeline] manifest -> {manifest}")
 
+    # --- Finalize: promosikan checkpoint terbaik ke out/ ---
     best = out_dir / "indigo_best.safetensors"
     if args.finalize:
         has_best = (out_dir / "indigo_best.safetensors").exists()
@@ -252,6 +355,7 @@ def main():
         )
         print("[pipeline] silsilah -> out/run_info.json")
 
+    # --- Push ke Hugging Face Hub ---
     if args.push:
         from huggingface_hub import HfApi
 

tests/test_core.py

@@ -0,0 +1,65 @@
+import sys
+from pathlib import Path
+
+ROOT = Path(__file__).resolve().parent.parent
+sys.path.insert(0, str(ROOT))
+
+from indigo.bpe import BPETokenizer
+from indigo.common import (
+    _kandidat_akar,
+    load_wordlist,
+    save_meta,
+    load_meta,
+    word_known_ratio,
+)
+from indigo.tokenizer import CharTokenizer
+
+
+def test_bpe_roundtrip():
+    teks = "halo dunia, halo lagi! belajar bahasa indonesia bersama."
+    tok = BPETokenizer.train(teks * 5, 280)
+    ids = tok.encode(teks)
+    assert tok.decode(ids) == teks
+    assert 256 <= tok.vocab_size <= 280
+
+
+def test_char_roundtrip():
+    tok = CharTokenizer.from_text("abcba")
+    assert tok.decode(tok.encode("abc")) == "abc"
+
+
+def test_meta_save_load(tmp_path):
+    base = str(tmp_path / "ck.safetensors")
+    (tmp_path / "ck.safetensors").write_bytes(b"")
+    cfg = {"vocab_size": 10, "block_size": 8}
+    save_meta(base, cfg, None, 5, 1.23, backend="pytorch", tokenizer={"type": "char"})
+    meta = load_meta(base)
+    assert meta["config"] == cfg and meta["step"] == 5 and meta["val_loss"] == 1.23
+
+
+def test_afiks_asimilasi():
+    pref = {"meng", "meny", "men", "mem", "pem", "di", "ter"}
+    suf = {"kan", "annya", "nya", "an", "i", "lah"}
+    assert "sapu" in _kandidat_akar("menyapu", pref, suf)
+    assert "pukul" in _kandidat_akar("pemukul", pref, suf)
+    assert "ambil" in _kandidat_akar("mengambil", pref, suf)
+    assert "terima" in _kandidat_akar("diterima", pref, suf)
+
+
+def test_word_known_ratio_naik_dengan_afiks(tmp_path):
+    kamus = tmp_path / "kamus.txt"
+    kamus.write_text("paham\nterima\nringan\n", encoding="utf-8")
+    ws = load_wordlist(str(kamus))
+    pref = {"di"}
+    suf = {"i", "nya"}
+    teks = "paham dipahami ringannya"
+    r0 = word_known_ratio(teks, ws)
+    r1 = word_known_ratio(teks, ws, pref, suf)
+    assert r0 < r1 <= 1.0
+
+
+def test_load_wordlist_normalisasi(tmp_path):
+    f = tmp_path / "w.txt"
+    f.write_text("Apa\nBEBEK\n\nbebek\n", encoding="utf-8")
+    ws = load_wordlist(str(f))
+    assert ws == {"apa", "bebek"}

train.py

@@ -1,3 +1,21 @@
+"""
+Script training model Indigo GPT dari nol.
+
+Fitur:
+- Training loop standar dengan AdamW optimizer
+- Learning rate schedule: warmup linear → cosine decay
+- Best checkpoint otomatis berdasarkan validasi
+- Resume training dari checkpoint sebelumnya (--init-from)
+- Dukungan tokenizer char dan BPE
+- Gradient clipping untuk stabilitas
+- Statistik ringkasan di akhir run
+
+Cara pakai:
+    python train.py --data data/sample.txt --steps 2000
+    python train.py --data data/teks.txt --tokenizer bpe --vocab-size 512
+    python train.py --init-from out/indigo_best.safetensors --steps 1000
+"""
+
 import os
 import time
 import math
@@ -19,11 +37,24 @@ from indigo.tokenizer import CharTokenizer
 
 
 def load_init(path):
+    """Muat checkpoint untuk melanjutkan training (resume).
+
+    Mendukung dua format:
+    1. .safetensors: format utama Indigo (safetensors + _meta.json + optimizer.pt)
+    2. .pt: format PyTorch checkpoint lama (model, config, vocab, optimizer dalam 1 file)
+
+    Args:
+        path: Path ke file checkpoint (.safetensors atau .pt).
+
+    Returns:
+        Tuple (state_dict, meta_dict, optimizer_state atau None).
+    """
     if path.endswith(".safetensors"):
         from safetensors.torch import load_file
 
         state = load_file(path)
         meta = load_meta(path)
+        # Cari file optimizer (suffix _best dihapus untuk file optimizer)
         opt_path = os.path.splitext(path)[0].replace("_best", "") + "_optimizer.pt"
         opt = None
         if os.path.exists(opt_path):
@@ -32,6 +63,7 @@ def load_init(path):
             except Exception as e:
                 print(f"optimizer state dilewati: {e}")
         return state, meta, opt
+    # Format .pt lama
     ckpt = torch.load(path, map_location="cpu", weights_only=True)
     meta = {
         "config": ckpt["config"],
@@ -42,10 +74,32 @@ def load_init(path):
     return ckpt["model"], meta, ckpt.get("optimizer")
 
 
+# Cache arange tensor per (block_size, device) untuk menghindari alokasi berulang
+# saat get_batch dipanggil ribuan kali — menghemat ~11x waktu.
 _ARANGE_CACHE = {}
 
 
 def get_batch(data, block_size, batch_size, device):
+    """Ambil batch data latih secara random (vectorized).
+
+    Proses:
+    1. Pilih batch_size posisi awal secara acak dari data
+    2. Untuk setiap posisi, ambil potongan sepanjang block_size (input) dan block_size (target)
+    3. Target = input bergeser 1 posisi ke kanan (next-token prediction)
+
+    Menggunakan fancy indexing dan arange cache untuk efisiensi:
+    - ix: posisi awal random untuk setiap sampel dalam batch
+    - idx: matriks posisi (batch_size × block_size) dengan offset arange
+
+    Args:
+        data:        Tensor 1D — seluruh data training (token IDs).
+        block_size:  Int — panjang konteks per sampel.
+        batch_size:  Int — jumlah sampel per batch.
+        device:      Str — "cpu" atau "cuda".
+
+    Returns:
+        Tuple (x, y) — x: input (B, T), y: target (B, T).
+    """
     ix = torch.randint(len(data) - block_size - 1, (batch_size,))
     arange = _ARANGE_CACHE.get((block_size, device))
     if arange is None:
@@ -59,6 +113,20 @@ def get_batch(data, block_size, batch_size, device):
 
 @torch.no_grad()
 def estimate_loss(model, data, args, device):
+    """Estimasi loss validasi dengan averaging beberapa batch.
+
+    Model dipindahkan ke mode eval (tanpa dropout), lalu dihitung loss rata-rata
+    dari eval_iters batch random. Hasilnya lebih stabil daripada single batch.
+
+    Args:
+        model:  Model GPT.
+        data:   Tensor 1D — data validasi (token IDs).
+        args:   Namespace — harus punya block_size, batch_size, eval_iters.
+        device: Str — "cpu" atau "cuda".
+
+    Returns:
+        Float — loss rata-rata (cross-entropy, nats per token).
+    """
     model.eval()
     losses = []
     for _ in range(args.eval_iters):
@@ -70,29 +138,83 @@ def estimate_loss(model, data, args, device):
 
 
 def main(argv=None):
+    """Fungsi utama training — bisa dipanggil dari CLI atau dari pipeline.py.
+
+    Pipeline training:
+    1. Parse argumen → setup device & seed
+    2. Kumpulkan file teks → split train/val
+    3. Bangun atau muat tokenizer → encode teks ke token IDs
+    4. Bangun atau muat model GPT
+    5. Setup optimizer (AdamW) + learning rate schedule
+    6. Loop training: forward → loss → backward → clip grad → step optimizer
+    7. Setiap eval_interval langkah: hitung val loss → save best checkpoint
+    8. Simpan checkpoint final + optimizer state + statistik
+
+    Args:
+        argv: List argumen CLI (atau None untuk pakai sys.argv).
+
+    Returns:
+        Dict statistik training (dipakai oleh pipeline.py untuk manifest.json).
+    """
     parser = argparse.ArgumentParser(description="Latih model Indigo dari scratch")
-    parser.add_argument("--data", nargs="+", default=["data/sample.txt"], help="path file/folder teks untuk training")
-    parser.add_argument("--out", default="out", help="folder output checkpoint")
-    parser.add_argument("--steps", type=int, default=2000)
-    parser.add_argument("--batch-size", type=int, default=32)
-    parser.add_argument("--block-size", type=int, default=128)
-    parser.add_argument("--n-layer", type=int, default=4)
-    parser.add_argument("--n-head", type=int, default=4)
-    parser.add_argument("--n-embd", type=int, default=128)
-    parser.add_argument("--dropout", type=float, default=0.1)
-    parser.add_argument("--lr", type=float, default=3e-4)
-    parser.add_argument("--warmup", type=int, default=100)
-    parser.add_argument("--weight-decay", type=float, default=0.1)
-    parser.add_argument("--eval-interval", type=int, default=200)
-    parser.add_argument("--eval-iters", type=int, default=20)
-    parser.add_argument("--seed", type=int, default=1337)
-    parser.add_argument("--init-from", default=None, help="checkpoint untuk melanjutkan training")
-    parser.add_argument("--device", default="auto", choices=["auto", "cpu", "cuda"])
-    parser.add_argument("--tokenizer", default="char", choices=["char", "bpe"])
-    parser.add_argument("--vocab-size", type=int, default=512, help="ukuran vocab untuk tokenizer bpe")
-    parser.add_argument("--val-fraction", type=float, default=0.1, help="proporsi file untuk validasi")
+
+    # --- Data ---
+    parser.add_argument("--data", nargs="+", default=["data/sample.txt"],
+                        help="path file/folder teks untuk training (bisa banyak, spasi-separated)")
+
+    # --- Output ---
+    parser.add_argument("--out", default="out",
+                        help="folder output checkpoint (.safetensors + _meta.json + _optimizer.pt)")
+
+    # --- Hyperparameter Training ---
+    parser.add_argument("--steps", type=int, default=2000,
+                        help="jumlah total langkah training (default: 2000)")
+    parser.add_argument("--batch-size", type=int, default=32,
+                        help="jumlah sampel per batch (default: 32)")
+    parser.add_argument("--block-size", type=int, default=128,
+                        help="panjang konteks token per sampel (default: 128)")
+    parser.add_argument("--lr", type=float, default=3e-4,
+                        help="learning rate maksimum (default: 3e-4)")
+    parser.add_argument("--warmup", type=int, default=100,
+                        help="jumlah langkah warmup linear sebelum cosine decay (default: 100)")
+    parser.add_argument("--weight-decay", type=float, default=0.1,
+                        help="L2 regularization / weight decay (default: 0.1)")
+    parser.add_argument("--dropout", type=float, default=0.1,
+                        help="dropout rate (0.0 = nonaktif, default: 0.1)")
+
+    # --- Arsitektur Model ---
+    parser.add_argument("--n-layer", type=int, default=4,
+                        help="jumlah blok transformer (default: 4)")
+    parser.add_argument("--n-head", type=int, default=4,
+                        help="jumlah head per attention layer (default: 4)")
+    parser.add_argument("--n-embd", type=int, default=128,
+                        help="dimensi embedding / hidden size (default: 128)")
+
+    # --- Tokenizer ---
+    parser.add_argument("--tokenizer", default="char", choices=["char", "bpe"],
+                        help="jenis tokenizer: 'char' (karakter) atau 'bpe' (subword, default: char)")
+    parser.add_argument("--vocab-size", type=int, default=512,
+                        help="ukuran vocab untuk BPE (diabaikan jika --tokenizer char, default: 512)")
+
+    # --- Evaluasi & Seed ---
+    parser.add_argument("--eval-interval", type=int, default=200,
+                        help="evaluasi validasi setiap N langkah (0 = tidak ada validasi, default: 200)")
+    parser.add_argument("--eval-iters", type=int, default=20,
+                        help="jumlah batch untuk estimasi loss validasi (default: 20)")
+    parser.add_argument("--seed", type=int, default=1337,
+                        help="seed random untuk reproduktibilitas (default: 1337)")
+
+    # --- Validasi & Resume ---
+    parser.add_argument("--val-fraction", type=float, default=0.1,
+                        help="proporsi file untuk validasi (default: 0.1 = 10%%)")
+    parser.add_argument("--init-from", default=None,
+                        help="path checkpoint untuk melanjutkan training (resume)")
+    parser.add_argument("--device", default="auto", choices=["auto", "cpu", "cuda"],
+                        help="device training: auto/cpu/cuda (default: auto)")
+
     args = parser.parse_args(argv)
 
+    # --- Setup seed & device ---
     torch.manual_seed(args.seed)
     if args.device == "auto":
         device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -100,10 +222,15 @@ def main(argv=None):
         device = args.device
     os.makedirs(args.out, exist_ok=True)
 
+    # --- Kumpulkan & split data ---
+    # collect_text_files: jika path adalah direktori, cari .txt rekursif
     paths = collect_text_files(args.data)
     if not paths:
         raise SystemExit("tidak ada file teks ditemukan")
 
+    # Acak urutan file → split: n_val file untuk validasi, sisanya untuk training
+    # Split dilakukan per-file (bukan per-karakter), sehingga satu file kecil
+    # bisa menghabiskan seluruh kuota validasi
     files = sorted(paths)
     rng = random.Random(args.seed)
     rng.shuffle(files)
@@ -112,15 +239,18 @@ def main(argv=None):
 
     train_text = "".join(read_clean(p) for p in files[n_val:])
     val_text = "".join(read_clean(p) for p in files[:n_val])
-    all_text = train_text + val_text
+    all_text = train_text + val_text  # dibutuhkan untuk training tokenizer BPE
 
+    # --- Setup model & tokenizer ---
     init_state = None
     init_opt = None
     start_step = 0
     init_meta = None
     config = None
     comp_ratio = 1.0
+
     if args.init_from:
+        # Resume dari checkpoint: muat model, tokenizer, dan optimizer
         init_state, init_meta, init_opt = load_init(args.init_from)
         config = GPTConfig(**init_meta["config"])
         start_step = init_meta.get("step", 0)
@@ -128,6 +258,7 @@ def main(argv=None):
         tokenizer = build_tokenizer(init_meta.get("tokenizer") or {"type": "char"}, init_meta["vocab"])
         tinfo = init_meta.get("tokenizer") or {"type": "char"}
     else:
+        # Training dari nol: bangun tokenizer baru
         if args.tokenizer == "bpe":
             from indigo.bpe import BPETokenizer
 
@@ -142,6 +273,7 @@ def main(argv=None):
         else:
             tokenizer = CharTokenizer.from_text(all_text)
             tinfo = {"type": "char"}
+        # Bangun config model baru dari argumen CLI
         if config is None:
             config = GPTConfig(
                 vocab_size=tokenizer.vocab_size,
@@ -151,11 +283,14 @@ def main(argv=None):
                 n_embd=args.n_embd,
                 dropout=args.dropout,
             )
+
+    # Validasi: vocab size model harus cocok dengan tokenizer
     if config.vocab_size != tokenizer.vocab_size:
         raise SystemExit(
             f"vocab tidak cocok: checkpoint={config.vocab_size}, tokenizer={tokenizer.vocab_size}"
         )
 
+    # --- Encode teks ke token IDs ---
     train_data = torch.tensor(tokenizer.encode(train_text), dtype=torch.long)
     val_data = torch.tensor(tokenizer.encode(val_text), dtype=torch.long)
     if len(train_data) < args.block_size * 2:
@@ -164,6 +299,7 @@ def main(argv=None):
         f"tokens latih={len(train_data):,} | tokens validasi={len(val_data):,}"
     )
 
+    # --- Inisialisasi model ---
     model = GPT(config)
     if init_state is not None:
         missing, unexpected = model.load_state_dict(init_state, strict=False)
@@ -176,6 +312,7 @@ def main(argv=None):
         f"vocab={tokenizer.vocab_size} | total_steps={total_steps}"
     )
 
+    # --- Setup optimizer: AdamW dengan betas=(0.9, 0.95) ---
     optimizer = torch.optim.AdamW(
         model.parameters(), lr=args.lr, betas=(0.9, 0.95), weight_decay=args.weight_decay
     )
@@ -187,6 +324,7 @@ def main(argv=None):
             print(f"optimizer state dilewati: {e}")
 
     def save_model(base_path, val_loss):
+        """Simpan checkpoint model + metadata ke file .safetensors + _meta.json."""
         tensors = {k: v.detach().clone().contiguous() for k, v in model.state_dict().items()}
         save_file(tensors, base_path)
         save_meta(
@@ -200,27 +338,43 @@ def main(argv=None):
         )
 
     def lr_at(step):
+        """Hitung learning rate pada step tertentu.
+
+        Schedule:
+        - Warmup (step < warmup): linear naik dari 0 ke lr maks
+        - Setelah warmup: cosine decay dari lr maks ke 10% lr maks
+        - Formula cosine: 0.1*lr + 0.45*lr * (1 + cos(pi * progress))
+
+        Jika warmup=0, langsung masuk cosine decay dari step 0.
+        """
         if step < args.warmup:
-            return args.lr * (step + 1) / args.warmup
+            return args.lr * (step + 1) / max(1, args.warmup)
         progress = (step - args.warmup) / max(1, total_steps - args.warmup)
         return 0.1 * args.lr + 0.45 * args.lr * (1 + math.cos(math.pi * progress))
 
+    # --- Training loop ---
     best_val = float("inf")
     last_val = None
     model.train()
     t0 = time.time()
     for step in range(start_step, total_steps):
+        # Update learning rate sesuai schedule
         lr = lr_at(step)
         for g in optimizer.param_groups:
             g["lr"] = lr
+
+        # Forward pass: ambil batch → hitung loss
         x, y = get_batch(train_data, config.block_size, args.batch_size, device)
         _, loss = model(x, y)
+
+        # Backward pass: zero grad → backward → clip grad → step optimizer
         optimizer.zero_grad(set_to_none=True)
         loss.backward()
-        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)  # mencegah gradient explosion
         optimizer.step()
 
-        if step % args.eval_interval == 0 or step == total_steps - 1:
+        # Evaluasi validasi + simpan best checkpoint
+        if args.eval_interval > 0 and (step % args.eval_interval == 0 or step == total_steps - 1):
             if len(val_data) > config.block_size + 1:
                 val_loss = estimate_loss(model, val_data, args, device)
                 marker = ""
@@ -237,11 +391,13 @@ def main(argv=None):
                 f"loss {loss.item():.4f} | val {val_str} | {time.time() - t0:.1f}s"
             )
 
+    # --- Simpan checkpoint final (bukan best) ---
     final_path = os.path.join(args.out, "indigo.safetensors")
     save_model(final_path, last_val)
     torch.save(optimizer.state_dict(), os.path.join(args.out, "indigo_optimizer.pt"))
     print(f"model tersimpan di {final_path} (+_meta.json, indigo_optimizer.pt)")
 
+    # --- Ringkasan statistik ---
     stats = {
         "out": args.out,
         "device": device,
@@ -258,7 +414,8 @@ def main(argv=None):
         "best_val": best_val if best_val != float("inf") else None,
         "last_val": last_val,
         "nats_per_char_best": (
-            round(best_val / comp_ratio, 4) if best_val != float("inf") and comp_ratio else None
+            round(best_val / comp_ratio, 4)
+            if best_val != float("inf") and comp_ratio else None
         ),
         "params_million": round(model.num_params() / 1e6, 4),
         "config": config.__dict__,