Upload folder using huggingface_hub

README.md

@@ -14,49 +14,50 @@ datasets:
 
 # Indigo
 
-Model bahasa kecil GPT-style yang dibangun **dari nol** (tanpa library transformers) sebagai proyek pembelajaran. Dual-backend: **PyTorch** dan **TensorFlow/Keras**, dengan format bobot aman `.safetensors`.
+Model bahasa kecil GPT-style yang dibangun **dari nol** (tanpa library transformers) sebagai proyek pembelajaran. Backend **PyTorch**; tersedia juga varian [TensorFlow/Keras terpisah](https://huggingface.co/adyoi/indigo-tf). Bobot disimpan dalam format aman `.safetensors`.
 
 ## Arsitektur
 
 | | Nilai default |
 |---|---|
-| Tipe | Decoder-only transformer (pre-LN) |
+| Tipe | Decoder-only transformer (pre-LN, SDPA) |
 | Parameter | ~0.81M |
 | Layer / Head | 4 / 4 |
 | Dimensi | 128 |
 | Konteks | 96 token |
-| Tokenizer | Level karakter (~90 vocab) |
+| Tokenizer | karakter atau BPE (`--tokenizer bpe`) |
 
 ## File penting
 
 ```
-indigo/model.py        arsitektur PyTorch (SDPA)
-indigo/model_keras.py  arsitektur Keras 3
-train.py / train_tf.py training per backend
-generate.py / generate_tf.py  generasi teks
+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
 ```
 
 ## Cara pakai
 
 ```bash
-pip install -r requirements.txt   # torch (+ tensorflow opsional)
+pip install -r requirements.txt
 
-# PyTorch
 python train.py --data data/sample.txt --steps 2000
-python generate.py --prompt "Indigo" --max-new 300
+python train.py --data data/tekskamu.txt --tokenizer bpe --vocab-size 512
 
-# TensorFlow/Keras
-python train_tf.py --data data/sample.txt --out out_tf
-python generate_tf.py --prompt "Indigo"
+python generate.py --prompt "Indigo" --max-new 300 \
+  --temperature 0.8 --top-k 40 --top-p 0.9 --repetition-penalty 1.2
+
+python train.py --init-from out/indigo_best.safetensors --steps 1000  # lanjutkan training
 ```
 
-Fitur: pembersihan teks otomatis, split validasi per-file, best-checkpoint, resume (`--init-from`), cosine LR + warmup, top-k sampling, pilihan `--device`.
+Generasi memakai KV-cache sehingga cepat untuk output panjang.
 
 ## Batasan
 
-- Dilatih pada data sangat kecil (~16 ribu token) → output belum koheren, cocok untuk edukasi bukan produksi.
-- Checkpoint berlabel `_best` dipilih berdasarkan validasi; gunakan itu, bukan checkpoint akhir.
+- Dilatih pada data sangat kecil → output belum koheren; cocok untuk edukasi bukan produksi.
+- Gunakan checkpoint `indigo_best` (terpilih berdasarkan validasi), bukan checkpoint akhir.
 
 ## Keamanan
 

generate.py

@@ -1,11 +1,8 @@
-import argparse
-import json
-import os
-
 import torch
+import argparse
 
+from indigo.common import load_meta, build_tokenizer
 from indigo.model import GPT, GPTConfig
-from indigo.tokenizer import CharTokenizer
 
 
 def load_model(path):
@@ -13,11 +10,20 @@ def load_model(path):
         from safetensors.torch import load_file
 
         state = load_file(path)
-        with open(os.path.splitext(path)[0] + "_meta.json", encoding="utf-8") as f:
-            meta = json.load(f)
-        return state, meta["config"], meta["vocab"]
-    ckpt = torch.load(path, map_location="cpu", weights_only=True)
-    return ckpt["model"], ckpt["config"], ckpt["vocab"]
+        meta = load_meta(path)
+        config_d = meta["config"]
+        tinfo = meta.get("tokenizer") or {"type": "char"}
+        vocab = meta.get("vocab")
+    else:
+        ckpt = torch.load(path, map_location="cpu", weights_only=True)
+        state = ckpt["model"]
+        config_d = ckpt["config"]
+        tinfo = ckpt.get("tokenizer") or {"type": "char"}
+        vocab = ckpt.get("vocab")
+    tokenizer = build_tokenizer(tinfo, vocab)
+    model = GPT(GPTConfig(**config_d))
+    model.load_state_dict(state, strict=False)
+    return model, tokenizer
 
 
 def main():
@@ -27,6 +33,8 @@ def main():
     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"])
     args = parser.parse_args()
@@ -35,15 +43,19 @@ def main():
         torch.manual_seed(args.seed)
     device = "cuda" if torch.cuda.is_available() else "cpu" if args.device == "auto" else args.device
 
-    state, config_d, vocab = load_model(args.ckpt)
-    model = GPT(GPTConfig(**config_d))
-    model.load_state_dict(state, strict=False)
+    model, tokenizer = load_model(args.ckpt)
     model = model.to(device)
-    tokenizer = CharTokenizer(vocab)
 
     ids = tokenizer.encode(args.prompt) or [0]
     idx = torch.tensor([ids], dtype=torch.long, device=device)
-    out = model.generate(idx, args.max_new, temperature=args.temperature, top_k=args.top_k)
+    out = model.generate(
+        idx,
+        args.max_new,
+        temperature=args.temperature,
+        top_k=args.top_k,
+        top_p=args.top_p,
+        repetition_penalty=args.repetition_penalty,
+    )
     print(tokenizer.decode(out[0].tolist()))
 
 

indigo/bpe.py

@@ -0,0 +1,75 @@
+def _get_stats(ids):
+    stats = {}
+    for pair in zip(ids, ids[1:]):
+        stats[pair] = stats.get(pair, 0) + 1
+    return stats
+
+
+def _merge(ids, pair, idx):
+    out = []
+    i = 0
+    while i < len(ids):
+        if i < len(ids) - 1 and ids[i] == pair[0] and ids[i + 1] == pair[1]:
+            out.append(idx)
+            i += 2
+        else:
+            out.append(ids[i])
+            i += 1
+    return out
+
+
+class BPETokenizer:
+    def __init__(self, merges=None):
+        self.merges = [tuple(p) for p in (merges or [])]
+        self.ranks = {pair: i for i, pair in enumerate(self.merges)}
+        self.vocab = [bytes([i]) for i in range(256)]
+        for a, b in self.merges:
+            self.vocab.append(self.vocab[a] + self.vocab[b])
+
+    @classmethod
+    def train(cls, text, vocab_size):
+        tok = cls()
+        ids = list(text.encode("utf-8"))
+        next_id = 256
+        while next_id < vocab_size and len(ids) > 1:
+            stats = _get_stats(ids)
+            pair = max(stats, key=stats.get)
+            if stats[pair] < 2:
+                break
+            ids = _merge(ids, pair, next_id)
+            tok.ranks[pair] = len(tok.merges)
+            tok.merges.append(pair)
+            tok.vocab.append(tok.vocab[pair[0]] + tok.vocab[pair[1]])
+            next_id += 1
+        return tok
+
+    @property
+    def vocab_size(self):
+        return 256 + len(self.merges)
+
+    def _encode_chunk(self, ids):
+        while len(ids) >= 2:
+            best = None
+            best_rank = None
+            for pair in zip(ids, ids[1:]):
+                rank = self.ranks.get(pair)
+                if rank is not None and (best_rank is None or rank < best_rank):
+                    best = pair
+                    best_rank = rank
+            if best is None:
+                break
+            ids = _merge(ids, best, 256 + best_rank)
+        return ids
+
+    def encode(self, text):
+        return self._encode_chunk(list(text.encode("utf-8")))
+
+    def decode(self, ids):
+        return b"".join(self.vocab[i] for i in ids).decode("utf-8", errors="replace")
+
+    def state(self):
+        return {"type": "bpe", "merges": [list(p) for p in self.merges]}
+
+    @classmethod
+    def from_state(cls, state):
+        return cls(state["merges"])

indigo/common.py

@@ -28,13 +28,14 @@ def read_clean(path):
         return clean_text(f.read())
 
 
-def save_meta(base_path, config, vocab, step, val_loss, backend):
+def save_meta(base_path, config, vocab, step, val_loss, backend, tokenizer=None):
     meta = {
         "config": config,
         "vocab": vocab,
         "step": step,
         "val_loss": val_loss,
         "backend": backend,
+        "tokenizer": tokenizer or {"type": "char"},
     }
     with open(os.path.splitext(base_path)[0] + "_meta.json", "w", encoding="utf-8") as f:
         json.dump(meta, f, ensure_ascii=False)
@@ -43,3 +44,13 @@ def save_meta(base_path, config, vocab, step, val_loss, backend):
 def load_meta(path):
     with open(os.path.splitext(path)[0] + "_meta.json", encoding="utf-8") as f:
         return json.load(f)
+
+
+def build_tokenizer(tokenizer_info, vocab):
+    if tokenizer_info.get("type") == "bpe":
+        from indigo.bpe import BPETokenizer
+
+        return BPETokenizer.from_state(tokenizer_info)
+    from indigo.tokenizer import CharTokenizer
+
+    return CharTokenizer(vocab)

indigo/model.py

@@ -1,3 +1,4 @@
+import math
 from dataclasses import dataclass
 
 import torch
@@ -27,19 +28,30 @@ class CausalSelfAttention(nn.Module):
         self.attn_dropout = nn.Dropout(config.dropout)
         self.resid_dropout = nn.Dropout(config.dropout)
 
-    def forward(self, x):
+    def forward(self, x, kv=None):
         B, T, C = x.shape
         q, k, v = self.qkv(x).split(self.n_embd, dim=2)
         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)
-        y = F.scaled_dot_product_attention(
-            q, k, v,
-            dropout_p=self.attn_dropout.p if self.training else 0.0,
-            is_causal=True,
-        )
+        if kv is None:
+            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:
+            pk, pv = kv
+            k = torch.cat((pk, k), dim=2)
+            v = torch.cat((pv, v), dim=2)
+            Tq, Tk = q.size(2), k.size(2)
+            mask = torch.ones(Tq, Tk, dtype=torch.bool, device=x.device).tril(diagonal=Tk - Tq)
+            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
         y = y.transpose(1, 2).contiguous().view(B, T, C)
-        return self.resid_dropout(self.proj(y))
+        return self.resid_dropout(self.proj(y)), (k, v)
 
 
 class MLP(nn.Module):
@@ -61,10 +73,35 @@ class Block(nn.Module):
         self.ln2 = nn.LayerNorm(config.n_embd, bias=config.bias)
         self.mlp = MLP(config)
 
-    def forward(self, x):
-        x = x + self.attn(self.ln1(x))
+    def forward(self, x, kv=None):
+        a, kv_new = self.attn(self.ln1(x), kv)
+        x = x + a
         x = x + self.mlp(self.ln2(x))
-        return x
+        return x, kv_new
+
+
+def _sample_token(logits, temperature, top_k=None, top_p=None):
+    logits = logits / max(temperature, 1e-8)
+    if top_k is not None:
+        v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
+        logits[logits < v[:, [-1]]] = float("-inf")
+    if top_p is not None and top_p < 1.0:
+        sorted_logits, sorted_idx = torch.sort(logits, descending=True, dim=-1)
+        probs = F.softmax(sorted_logits, dim=-1)
+        cum = torch.cumsum(probs, dim=-1)
+        remove = (cum - probs) > top_p
+        sorted_logits[remove] = float("-inf")
+        logits = torch.full_like(logits, float("-inf")).scatter_(1, sorted_idx, sorted_logits)
+    probs = F.softmax(logits, dim=-1)
+    return torch.multinomial(probs, num_samples=1)
+
+
+def _apply_repetition_penalty(logits, tokens, penalty, window):
+    recent = set(tokens[0, -window:].tolist())
+    for t in recent:
+        val = logits[0, t]
+        logits[0, t] = torch.where(val > 0, val / penalty, val * penalty)
+    return logits
 
 
 class GPT(nn.Module):
@@ -88,32 +125,49 @@ class GPT(nn.Module):
         elif isinstance(module, nn.Embedding):
             nn.init.normal_(module.weight, mean=0.0, std=0.02)
 
-    def forward(self, idx, targets=None):
+    def forward(self, idx, targets=None, caches=None, return_caches=False):
         B, T = idx.shape
-        pos = torch.arange(T, device=idx.device)
+        start = caches[0][0].size(2) if caches else 0
+        pos = torch.arange(start, start + T, device=idx.device)
         x = self.drop(self.tok_emb(idx) + self.pos_emb(pos))
-        for block in self.blocks:
-            x = block(x)
+        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)
         logits = self.head(self.ln_f(x))
         loss = None
         if targets is not None:
             loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), targets.reshape(-1))
+        if return_caches:
+            return logits, loss, new_caches
         return logits, loss
 
     @torch.no_grad()
-    def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None):
+    def generate(
+        self,
+        idx,
+        max_new_tokens,
+        temperature=1.0,
+        top_k=None,
+        top_p=None,
+        repetition_penalty=1.0,
+    ):
         self.eval()
-        for _ in range(max_new_tokens):
-            idx_cond = idx[:, -self.config.block_size:]
-            logits, _ = self(idx_cond)
-            logits = logits[:, -1, :] / max(temperature, 1e-8)
-            if top_k is not None:
-                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
-                logits[logits < v[:, [-1]]] = float("-inf")
-            probs = F.softmax(logits, dim=-1)
-            next_id = torch.multinomial(probs, num_samples=1)
-            idx = torch.cat((idx, next_id), dim=1)
-        return idx
+        idx_cond = idx[:, -self.config.block_size:]
+        logits, _, caches = self(idx_cond, return_caches=True)
+        tokens = idx_cond
+        for i in range(max_new_tokens):
+            next_logits = logits[:, -1, :]
+            if repetition_penalty != 1.0:
+                next_logits = _apply_repetition_penalty(
+                    next_logits.clone(), tokens, repetition_penalty, self.config.block_size
+                )
+            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
+            logits, _, caches = self(next_id, caches=caches, return_caches=True)
+        return tokens
 
     def num_params(self):
         return sum(p.numel() for p in self.parameters())

out/indigo.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:8438242a855ba259d59d07b7e8a9a29461311c2f008993cb20d5e9f5828dbc72
-size 3294144
+oid sha256:1ff2d9c4acdba3bb129caf2163ca2b66885821121a749e0d15d4c282a2d0f71d
+size 3726272

out/indigo_best.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:c2f67a103c33e00f21bb78c13716b70ff97c0de33352a27f0b6fff8c01c480ab
-size 3294144
+oid sha256:6f6f8306cadd2b1130775c53814c35e7fe49e049fe86889e08a73a002a276f2d
+size 3726272

out/indigo_best_meta.json

@@ -1 +1 @@
-{"config": {"vocab_size": 90, "block_size": 96, "n_layer": 4, "n_head": 4, "n_embd": 128, "dropout": 0.1, "bias": false}, "vocab": ["\n", " ", "!", "\"", "#", "%", "&", "'", "(", ")", "+", ",", "-", ".", "/", "0", "1", "2", "3", "4", "5", "6", "7", "8", "9", ":", ";", "<", ">", "?", "A", "B", "C", "D", "E", "F", "G", "H", "I", "J", "K", "L", "M", "N", "O", "P", "Q", "R", "S", "T", "U", "V", "W", "X", "Y", "Z", "[", "]", "_", "`", "a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l", "m", "n", "o", "p", "q", "r", "s", "t", "u", "v", "w", "x", "y", "z", "{", "|", "}", "°"], "step": 900, "val_loss": 3.390139579772949}
+{"config": {"vocab_size": 512, "block_size": 96, "n_layer": 4, "n_head": 4, "n_embd": 128, "dropout": 0.1, "bias": false}, "vocab": null, "step": 900, "val_loss": 5.525297069549561, "backend": "pytorch", "tokenizer": {"type": "bpe", "merges": [[97, 110], [101, 114], [97, 32], [97, 116], [101, 110], [256, 32], [105, 110], [97, 108], [105, 32], [58, 32], [97, 114], [116, 105], [111, 114], [111, 110], [101, 115], [44, 32], [97, 109], [97, 104], [256, 103], [101, 32], [101, 109], [101, 107], [97, 107], [117, 110], [117, 32], [101, 108], [97, 115], [97, 100], [263, 32], [260, 103], [105, 115], [32, 109], [262, 103], [32, 40], [32, 115], [32, 100], [105, 116], [101, 116], [32, 65], [101, 99], [117, 104], [117, 116], [46, 10], [32, 83], [257, 277], [117, 115], [116, 257], [117, 114], [105, 108], [97, 112], [121, 274], [41, 32], [105, 99], [259, 32], [98, 257], [34, 10], [114, 111], [46, 32], [260, 116], [300, 258], [105, 107], [41, 10], [84, 73], [32, 80], [100, 105], [100, 261], [107, 101], [32, 112], [97, 99], [117, 107], [32, 116], [32, 67], [97, 103], [115, 116], [117, 108], [82, 69], [279, 103], [105, 100], [265, 34], [84, 69], [267, 269], [121, 258], [32, 69], [259, 105], [101, 100], [261, 109], [108, 111], [32, 38], [270, 115], [256, 100], [108, 273], [84, 65], [111, 109], [115, 101], [65, 76], [339, 269], [272, 264], [83, 73], [79, 78], [105, 109], [121, 32], [76, 76], [306, 32], [44, 10], [105, 104], [32, 68], [65, 82], [32, 107], [107, 261], [58, 10], [259, 258], [266, 264], [117, 100], [32, 77], [65, 121], [370, 309], [69, 82], [287, 285], [97, 98], [112, 297], [114, 105], [115, 281], [109, 315], [99, 104], [65, 73], [266, 258], [269, 116], [271, 321], [267, 100], [32, 78], [117, 112], [112, 283], [121, 97], [311, 371], [117, 109], [32, 99], [259, 101], [32, 98], [65, 78], [32, 75], [384, 278], [103, 262], [97, 280], [69, 83], [276, 98], [268, 274], [32, 73], [115, 105], [101, 112], [282, 264], [111, 108], [97, 119], [260, 267], [265, 39], [377, 102], [69, 78], [32, 265], [295, 116], [105, 114], [116, 104], [101, 102], [288, 103], [116, 114], [276, 112], [267, 118], [304, 108], [421, 266], [262, 116], [117, 259], [101, 120], [110, 397], [105, 103], [278, 261], [101, 98], [100, 264], [76, 69], [273, 32], [342, 103], [111, 301], [279, 116], [111, 100], [267, 99], [107, 259], [268, 32], [291, 261], [103, 117], [73, 78], [294, 67], [270, 32], [40, 77], [322, 387], [285, 107], [84, 296], [272, 280], [68, 261], [46, 112], [32, 66], [267, 102], [256, 99], [277, 115], [99, 105], [263, 272], [268, 275], [256, 116], [112, 257], [262, 100], [97, 262], [103, 296], [389, 49], [305, 258], [300, 97], [408, 116], [301, 257], [288, 32], [114, 270], [32, 70], [327, 79], [108, 346], [45, 78], [272, 109], [332, 463], [46, 39], [114, 101], [65, 83], [435, 325], [262, 302], [121, 110], [82, 73], [97, 121], [109, 314], [265, 77], [32, 84], [32, 10], [276, 268], [278, 341], [79, 84], [296, 475], [492, 283], [110, 337], [447, 398], [75, 352], [91, 380], [77, 80], [468, 95], [84, 85], [85, 357], [272, 258], [116, 345], [266, 100], [283, 258], [396, 32], [267, 107], [99, 344], [111, 103], [110, 105], [112, 269]]}}

out/indigo_meta.json

@@ -1 +1 @@
-{"config": {"vocab_size": 90, "block_size": 96, "n_layer": 4, "n_head": 4, "n_embd": 128, "dropout": 0.1, "bias": false}, "vocab": ["\n", " ", "!", "\"", "#", "%", "&", "'", "(", ")", "+", ",", "-", ".", "/", "0", "1", "2", "3", "4", "5", "6", "7", "8", "9", ":", ";", "<", ">", "?", "A", "B", "C", "D", "E", "F", "G", "H", "I", "J", "K", "L", "M", "N", "O", "P", "Q", "R", "S", "T", "U", "V", "W", "X", "Y", "Z", "[", "]", "_", "`", "a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l", "m", "n", "o", "p", "q", "r", "s", "t", "u", "v", "w", "x", "y", "z", "{", "|", "}", "°"], "step": 900, "val_loss": 3.423954129219055}
+{"config": {"vocab_size": 512, "block_size": 96, "n_layer": 4, "n_head": 4, "n_embd": 128, "dropout": 0.1, "bias": false}, "vocab": null, "step": 900, "val_loss": 5.809607744216919, "backend": "pytorch", "tokenizer": {"type": "bpe", "merges": [[97, 110], [101, 114], [97, 32], [97, 116], [101, 110], [256, 32], [105, 110], [97, 108], [105, 32], [58, 32], [97, 114], [116, 105], [111, 114], [111, 110], [101, 115], [44, 32], [97, 109], [97, 104], [256, 103], [101, 32], [101, 109], [101, 107], [97, 107], [117, 110], [117, 32], [101, 108], [97, 115], [97, 100], [263, 32], [260, 103], [105, 115], [32, 109], [262, 103], [32, 40], [32, 115], [32, 100], [105, 116], [101, 116], [32, 65], [101, 99], [117, 104], [117, 116], [46, 10], [32, 83], [257, 277], [117, 115], [116, 257], [117, 114], [105, 108], [97, 112], [121, 274], [41, 32], [105, 99], [259, 32], [98, 257], [34, 10], [114, 111], [46, 32], [260, 116], [300, 258], [105, 107], [41, 10], [84, 73], [32, 80], [100, 105], [100, 261], [107, 101], [32, 112], [97, 99], [117, 107], [32, 116], [32, 67], [97, 103], [115, 116], [117, 108], [82, 69], [279, 103], [105, 100], [265, 34], [84, 69], [267, 269], [121, 258], [32, 69], [259, 105], [101, 100], [261, 109], [108, 111], [32, 38], [270, 115], [256, 100], [108, 273], [84, 65], [111, 109], [115, 101], [65, 76], [339, 269], [272, 264], [83, 73], [79, 78], [105, 109], [121, 32], [76, 76], [306, 32], [44, 10], [105, 104], [32, 68], [65, 82], [32, 107], [107, 261], [58, 10], [259, 258], [266, 264], [117, 100], [32, 77], [65, 121], [370, 309], [69, 82], [287, 285], [97, 98], [112, 297], [114, 105], [115, 281], [109, 315], [99, 104], [65, 73], [266, 258], [269, 116], [271, 321], [267, 100], [32, 78], [117, 112], [112, 283], [121, 97], [311, 371], [117, 109], [32, 99], [259, 101], [32, 98], [65, 78], [32, 75], [384, 278], [103, 262], [97, 280], [69, 83], [276, 98], [268, 274], [32, 73], [115, 105], [101, 112], [282, 264], [111, 108], [97, 119], [260, 267], [265, 39], [377, 102], [69, 78], [32, 265], [295, 116], [105, 114], [116, 104], [101, 102], [288, 103], [116, 114], [276, 112], [267, 118], [304, 108], [421, 266], [262, 116], [117, 259], [101, 120], [110, 397], [105, 103], [278, 261], [101, 98], [100, 264], [76, 69], [273, 32], [342, 103], [111, 301], [279, 116], [111, 100], [267, 99], [107, 259], [268, 32], [291, 261], [103, 117], [73, 78], [294, 67], [270, 32], [40, 77], [322, 387], [285, 107], [84, 296], [272, 280], [68, 261], [46, 112], [32, 66], [267, 102], [256, 99], [277, 115], [99, 105], [263, 272], [268, 275], [256, 116], [112, 257], [262, 100], [97, 262], [103, 296], [389, 49], [305, 258], [300, 97], [408, 116], [301, 257], [288, 32], [114, 270], [32, 70], [327, 79], [108, 346], [45, 78], [272, 109], [332, 463], [46, 39], [114, 101], [65, 83], [435, 325], [262, 302], [121, 110], [82, 73], [97, 121], [109, 314], [265, 77], [32, 84], [32, 10], [276, 268], [278, 341], [79, 84], [296, 475], [492, 283], [110, 337], [447, 398], [75, 352], [91, 380], [77, 80], [468, 95], [84, 85], [85, 357], [272, 258], [116, 345], [266, 100], [283, 258], [396, 32], [267, 107], [99, 344], [111, 103], [110, 105], [112, 269]]}}

requirements.txt

@@ -1,3 +1,2 @@
 torch>=2.0
 safetensors>=0.4
-tensorflow>=2.16

train.py

@@ -1,13 +1,19 @@
-import argparse
-import math
 import os
-import random
 import time
-
+import math
 import torch
+import random
+import argparse
+
 from safetensors.torch import save_file
 
-from indigo.common import collect_text_files, load_meta, read_clean, save_meta
+from indigo.common import (
+    build_tokenizer,
+    collect_text_files,
+    load_meta,
+    read_clean,
+    save_meta,
+)
 from indigo.model import GPT, GPTConfig
 from indigo.tokenizer import CharTokenizer
 
@@ -25,9 +31,15 @@ def load_init(path):
                 opt = torch.load(opt_path, map_location="cpu", weights_only=True)
             except Exception as e:
                 print(f"optimizer state dilewati: {e}")
-        return state, meta["config"], meta.get("step", 0), opt
+        return state, meta, opt
     ckpt = torch.load(path, map_location="cpu", weights_only=True)
-    return ckpt["model"], ckpt["config"], ckpt.get("step", 0), ckpt.get("optimizer")
+    meta = {
+        "config": ckpt["config"],
+        "vocab": ckpt["vocab"],
+        "step": ckpt.get("step", 0),
+        "tokenizer": ckpt.get("tokenizer"),
+    }
+    return ckpt["model"], meta, ckpt.get("optimizer")
 
 
 def get_batch(data, block_size, batch_size, device):
@@ -68,6 +80,8 @@ def main():
     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")
     args = parser.parse_args()
 
@@ -90,36 +104,56 @@ def main():
 
     train_text = "".join(read_clean(p) for p in files[n_val:])
     val_text = "".join(read_clean(p) for p in files[:n_val])
-
-    tokenizer = CharTokenizer.from_text(train_text + val_text)
-    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:
-        raise SystemExit(f"data latih terlalu pendek ({len(train_data)} token), minimal {args.block_size * 2}")
-    print(
-        f"tokens latih={len(train_data):,} | tokens validasi={len(val_data):,} | vocab={tokenizer.vocab_size}"
-    )
+    all_text = train_text + val_text
 
     init_state = None
     init_opt = None
     start_step = 0
+    init_meta = None
+    config = None
     if args.init_from:
-        init_state, init_config, start_step, init_opt = load_init(args.init_from)
-        config = GPTConfig(**init_config)
+        init_state, init_meta, init_opt = load_init(args.init_from)
+        config = GPTConfig(**init_meta["config"])
+        start_step = init_meta.get("step", 0)
         print(f"melanjutkan dari {args.init_from} (step {start_step})")
+        tokenizer = build_tokenizer(init_meta.get("tokenizer") or {"type": "char"}, init_meta["vocab"])
+        tinfo = init_meta.get("tokenizer") or {"type": "char"}
     else:
-        config = GPTConfig(
-            vocab_size=tokenizer.vocab_size,
-            block_size=args.block_size,
-            n_layer=args.n_layer,
-            n_head=args.n_head,
-            n_embd=args.n_embd,
-            dropout=args.dropout,
-        )
+        if args.tokenizer == "bpe":
+            from indigo.bpe import BPETokenizer
+
+            tokenizer = BPETokenizer.train(all_text, args.vocab_size)
+            tinfo = tokenizer.state()
+            n_chars = len(all_text.encode("utf-8"))
+            print(
+                f"tokenizer=bpe | vocab={tokenizer.vocab_size} | "
+                f"kompresi {n_chars:,} karakter -> rasio {n_chars / max(1, len(tokenizer.encode(all_text))):.2f}x"
+            )
+        else:
+            tokenizer = CharTokenizer.from_text(all_text)
+            tinfo = {"type": "char"}
+        if config is None:
+            config = GPTConfig(
+                vocab_size=tokenizer.vocab_size,
+                block_size=args.block_size,
+                n_layer=args.n_layer,
+                n_head=args.n_head,
+                n_embd=args.n_embd,
+                dropout=args.dropout,
+            )
     if config.vocab_size != tokenizer.vocab_size:
         raise SystemExit(
-            f"vocab tidak cocok: checkpoint={config.vocab_size}, data={tokenizer.vocab_size}"
+            f"vocab tidak cocok: checkpoint={config.vocab_size}, tokenizer={tokenizer.vocab_size}"
         )
+
+    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:
+        raise SystemExit(f"data latih terlalu pendek ({len(train_data)} token), minimal {args.block_size * 2}")
+    print(
+        f"tokens latih={len(train_data):,} | tokens validasi={len(val_data):,}"
+    )
+
     model = GPT(config)
     if init_state is not None:
         missing, unexpected = model.load_state_dict(init_state, strict=False)
@@ -145,7 +179,15 @@ def main():
     def save_model(base_path, val_loss):
         tensors = {k: v.detach().clone().contiguous() for k, v in model.state_dict().items()}
         save_file(tensors, base_path)
-        save_meta(base_path, config.__dict__, tokenizer.itos, total_steps, val_loss, backend="pytorch")
+        save_meta(
+            base_path,
+            config.__dict__,
+            tokenizer.itos if hasattr(tokenizer, "itos") else None,
+            total_steps,
+            val_loss,
+            backend="pytorch",
+            tokenizer=tinfo,
+        )
 
     def lr_at(step):
         if step < args.warmup: