adyoi commited on
Commit
47dfda4
·
verified ·
1 Parent(s): f2fbf5f

Upload folder using huggingface_hub

Browse files
.gitignore ADDED
@@ -0,0 +1,5 @@
 
 
 
 
 
 
1
+ __pycache__/
2
+ *.pyc
3
+ out/
4
+ *.pt
5
+ .git/
README.md CHANGED
@@ -1,3 +1,31 @@
1
  ---
 
 
2
  license: mit
 
 
 
 
 
 
 
3
  ---
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
  ---
2
+ language:
3
+ - id
4
  license: mit
5
+ tags:
6
+ - text-generation
7
+ - from-scratch
8
+ - gpt
9
+ - transformer
10
+ - tensorflow
11
+ - keras
12
  ---
13
+
14
+ # Indigo-TF
15
+
16
+ Versi **TensorFlow/Keras** dari proyek [Indigo](https://huggingface.co/adyoi/indigo) — model bahasa GPT kecil dari nol. Repo utama (PyTorch): `D:\Documents\GitHub\Indigo`.
17
+
18
+ Format checkpoint identik (`.safetensors` + `_meta.json`, penanda `backend: tensorflow`), sehingga tidak bisa tertukar dengan checkpoint PyTorch.
19
+
20
+ ## Pakai
21
+
22
+ ```bash
23
+ pip install -r requirements.txt
24
+
25
+ python train.py --data data/sample.txt --steps 2000
26
+ python generate.py --prompt "Indigo"
27
+ ```
28
+
29
+ Fitur setara repo utama: pembersihan teks, split validasi per-file, best-checkpoint, resume (`--init-from`), tokenizer `char`/`bpe`, cosine LR + warmup, pilihan `--device`.
30
+
31
+ Catatan: GPU TensorFlow di Windows native tidak didukung sejak TF 2.11 — gunakan WSL2/Colab/Linux.
data/sample.txt ADDED
@@ -0,0 +1,50 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ Indigo adalah model bahasa kecil yang dibangun dari nol.
2
+ Model ini dilatih dengan PyTorch murni tanpa pustaka tambahan.
3
+ Arsitektur Indigo berupa transformer decoder sederhana.
4
+ Setiap blok transformer memuat perhatian kausal dan jaringan saraf.
5
+ Perhatian kausal membuat model hanya melihat token masa lalu.
6
+ Token masa lalu dipakai untuk memprediksi token berikutnya.
7
+ Prediksi token berikutnya adalah inti dari pemodelan bahasa.
8
+ Pemodelan bahasa bisa dipakai untuk membuat teks baru.
9
+ Teks baru dihasilkan token demi token secara berurutan.
10
+ Urutan token ditentukan oleh distribusi probabilitas model.
11
+ Probabilitas model dipelajari dari data latih.
12
+ Data latih berupa kumpulan teks biasa dalam berkas txt.
13
+ Berkas teks dibaca lalu diubah menjadi deret angka.
14
+ Deret angka tersebut menjadi masukan bagi jaringan saraf.
15
+ Jaringan saraf belajar dengan menurunkan fungsi kerugian.
16
+ Fungsi kerugian dihitung menggunakan cross entropy.
17
+ Cross entropy mengukur selisih prediksi dengan jawaban sebenarnya.
18
+ Jawaban sebenarnya adalah token yang muncul pada data asli.
19
+ Data asli dibagi menjadi data latih dan data validasi.
20
+ Data validasi dipakai untuk memantau kemajuan pelatihan.
21
+ Pelatihan berjalan selama ribuan langkah kecil.
22
+ Setiap langkah memperbarui bobot model sedikit demi sedikit.
23
+ Bobot model diperbarui oleh optimizer AdamW.
24
+ Optimizer menurunkan kerugian dengan gradien turun.
25
+ Gradien dihitung lewat propagasi balik otomatis dari PyTorch.
26
+ Propagasi balik meneruskan kesalahan ke setiap lapisan.
27
+ Lapisan pertama adalah lapisan penyematan token.
28
+ Penyematan token mengubah angka menjadi vektor padat.
29
+ Vektor padat digabung dengan penyematan posisi.
30
+ Penyematan posisi memberi tahu urutan kata dalam kalimat.
31
+ Kalimat panjang dipotong menjadi jendela tetap.
32
+ Jendela tetap disebut block size atau konteks.
33
+ Konteks dibatasi agar perhatian tetap efisien.
34
+ Perhatian menghitung hubungan antar token dalam konteks.
35
+ Hubungan antar token membentuk pemahaman sederhana.
36
+ Pemahaman ini tumbuh seiring bertambahnya langkah latih.
37
+ Langkah latih dicatat beserta nilai kerugiannya.
38
+ Nilai kerugian yang menurun menandakan model belajar.
39
+ Model yang sudah belajar bisa menyimpan checkpoint.
40
+ Checkpoint berisi bobot, konfigurasi, dan kosakata.
41
+ Kosakata Indigo berupa karakter unik dari data.
42
+ Tokenizer karakter mengubah huruf menjadi angka sederhana.
43
+ Angka sederhana mudah dipahami tanpa pustaka eksternal.
44
+ Tanpa pustaka eksternal proyek tetap ringan dan jelas.
45
+ Proyek ringan cocok untuk belajar dari dasar.
46
+ Dasar yang kuat memudahkan eksperimen lanjutan.
47
+ Eksperimen bisa berupa penambahan lapisan atau kepala perhatian.
48
+ Kepala perhatian tambahan menangkap pola yang lebih beragam.
49
+ Pola beragam menghasilkan teks yang lebih koheren.
50
+ Teks koheren adalah tujuan akhir dari pelatihan ini.
generate.py ADDED
@@ -0,0 +1,67 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import sys
2
+ import random
3
+ import argparse
4
+ import numpy as np
5
+ import tensorflow as tf
6
+
7
+ from indigotf.common import build_tokenizer, load_meta
8
+ from indigotf.model import build_gpt, generate
9
+
10
+ CONFIG_KEYS = ("vocab_size", "block_size", "n_layer", "n_head", "n_embd", "dropout")
11
+
12
+
13
+ def load_model(path):
14
+ from safetensors.numpy import load_file
15
+
16
+ meta = load_meta(path)
17
+ if meta.get("backend") != "tensorflow":
18
+ raise SystemExit(
19
+ f"{path} berasal dari backend {meta.get('backend')}, gunakan repo indigo (PyTorch)"
20
+ )
21
+ state = load_file(path)
22
+ by_path = {k.replace("/", "_"): v for k, v in state.items()}
23
+ model = build_gpt(**{k: meta["config"][k] for k in CONFIG_KEYS})
24
+ missing = [v.path for v in model.weights if v.path.replace("/", "_") not in by_path]
25
+ if missing:
26
+ raise SystemExit(f"bobot tidak cocok dengan checkpoint: {missing[:5]}")
27
+ model.set_weights([by_path[v.path.replace("/", "_")] for v in model.weights])
28
+ tokenizer = build_tokenizer(meta.get("tokenizer") or {"type": "char"}, meta.get("vocab"))
29
+ return model, meta, tokenizer
30
+
31
+
32
+ def main():
33
+ sys.stdout.reconfigure(encoding="utf-8", errors="replace")
34
+ parser = argparse.ArgumentParser(description="Generate teks dari checkpoint Indigo-TF")
35
+ parser.add_argument("--ckpt", default="out/indigo_best.safetensors")
36
+ parser.add_argument("--prompt", default="")
37
+ parser.add_argument("--max-new", type=int, default=300)
38
+ parser.add_argument("--temperature", type=float, default=0.8)
39
+ parser.add_argument("--top-k", type=int, default=40)
40
+ parser.add_argument("--seed", type=int, default=None)
41
+ parser.add_argument("--device", default="auto", choices=["auto", "cpu", "gpu"])
42
+ args = parser.parse_args()
43
+
44
+ if args.device == "cpu":
45
+ tf.config.set_visible_devices([], "GPU")
46
+ if args.seed is not None:
47
+ random.seed(args.seed)
48
+ np.random.seed(args.seed)
49
+ tf.random.set_seed(args.seed)
50
+
51
+ model, meta, tokenizer = load_model(args.ckpt)
52
+
53
+ ids = tokenizer.encode(args.prompt) or [0]
54
+ idx = tf.constant([ids], dtype=tf.int64)
55
+ out = generate(
56
+ model,
57
+ idx,
58
+ args.max_new,
59
+ block_size=meta["config"]["block_size"],
60
+ temperature=args.temperature,
61
+ top_k=args.top_k,
62
+ )
63
+ print(tokenizer.decode(out.numpy()[0].tolist()))
64
+
65
+
66
+ if __name__ == "__main__":
67
+ main()
indigotf/__init__.py ADDED
@@ -0,0 +1,5 @@
 
 
 
 
 
 
1
+ from .bpe import BPETokenizer
2
+ from .tokenizer import CharTokenizer
3
+ from .model import build_gpt, generate
4
+
5
+ __all__ = ["build_gpt", "generate", "BPETokenizer", "CharTokenizer"]
indigotf/bpe.py ADDED
@@ -0,0 +1,75 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ def _get_stats(ids):
2
+ stats = {}
3
+ for pair in zip(ids, ids[1:]):
4
+ stats[pair] = stats.get(pair, 0) + 1
5
+ return stats
6
+
7
+
8
+ def _merge(ids, pair, idx):
9
+ out = []
10
+ i = 0
11
+ while i < len(ids):
12
+ if i < len(ids) - 1 and ids[i] == pair[0] and ids[i + 1] == pair[1]:
13
+ out.append(idx)
14
+ i += 2
15
+ else:
16
+ out.append(ids[i])
17
+ i += 1
18
+ return out
19
+
20
+
21
+ class BPETokenizer:
22
+ def __init__(self, merges=None):
23
+ self.merges = [tuple(p) for p in (merges or [])]
24
+ self.ranks = {pair: i for i, pair in enumerate(self.merges)}
25
+ self.vocab = [bytes([i]) for i in range(256)]
26
+ for a, b in self.merges:
27
+ self.vocab.append(self.vocab[a] + self.vocab[b])
28
+
29
+ @classmethod
30
+ def train(cls, text, vocab_size):
31
+ tok = cls()
32
+ ids = list(text.encode("utf-8"))
33
+ next_id = 256
34
+ while next_id < vocab_size and len(ids) > 1:
35
+ stats = _get_stats(ids)
36
+ pair = max(stats, key=stats.get)
37
+ if stats[pair] < 2:
38
+ break
39
+ ids = _merge(ids, pair, next_id)
40
+ tok.ranks[pair] = len(tok.merges)
41
+ tok.merges.append(pair)
42
+ tok.vocab.append(tok.vocab[pair[0]] + tok.vocab[pair[1]])
43
+ next_id += 1
44
+ return tok
45
+
46
+ @property
47
+ def vocab_size(self):
48
+ return 256 + len(self.merges)
49
+
50
+ def _encode_chunk(self, ids):
51
+ while len(ids) >= 2:
52
+ best = None
53
+ best_rank = None
54
+ for pair in zip(ids, ids[1:]):
55
+ rank = self.ranks.get(pair)
56
+ if rank is not None and (best_rank is None or rank < best_rank):
57
+ best = pair
58
+ best_rank = rank
59
+ if best is None:
60
+ break
61
+ ids = _merge(ids, best, 256 + best_rank)
62
+ return ids
63
+
64
+ def encode(self, text):
65
+ return self._encode_chunk(list(text.encode("utf-8")))
66
+
67
+ def decode(self, ids):
68
+ return b"".join(self.vocab[i] for i in ids).decode("utf-8", errors="replace")
69
+
70
+ def state(self):
71
+ return {"type": "bpe", "merges": [list(p) for p in self.merges]}
72
+
73
+ @classmethod
74
+ def from_state(cls, state):
75
+ return cls(state["merges"])
indigotf/common.py ADDED
@@ -0,0 +1,56 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import re
3
+ import json
4
+
5
+ DECOR_LINE = re.compile(r"^[\s=\-_~*#.]{4,}$")
6
+
7
+
8
+ def clean_text(text):
9
+ lines = [ln for ln in text.splitlines() if not DECOR_LINE.match(ln)]
10
+ text = "\n".join(lines)
11
+ text = re.sub(r"\n{3,}", "\n\n", text)
12
+ return text.strip() + "\n"
13
+
14
+
15
+ def collect_text_files(paths):
16
+ files = []
17
+ for p in paths:
18
+ if os.path.isdir(p):
19
+ for root, _, names in os.walk(p):
20
+ files.extend(os.path.join(root, n) for n in sorted(names) if n.lower().endswith(".txt"))
21
+ else:
22
+ files.append(p)
23
+ return sorted(files)
24
+
25
+
26
+ def read_clean(path):
27
+ with open(path, encoding="utf-8") as f:
28
+ return clean_text(f.read())
29
+
30
+
31
+ def save_meta(base_path, config, vocab, step, val_loss, backend, tokenizer=None):
32
+ meta = {
33
+ "config": config,
34
+ "vocab": vocab,
35
+ "step": step,
36
+ "val_loss": val_loss,
37
+ "backend": backend,
38
+ "tokenizer": tokenizer or {"type": "char"},
39
+ }
40
+ with open(os.path.splitext(base_path)[0] + "_meta.json", "w", encoding="utf-8") as f:
41
+ json.dump(meta, f, ensure_ascii=False)
42
+
43
+
44
+ def load_meta(path):
45
+ with open(os.path.splitext(path)[0] + "_meta.json", encoding="utf-8") as f:
46
+ return json.load(f)
47
+
48
+
49
+ def build_tokenizer(tokenizer_info, vocab):
50
+ if tokenizer_info.get("type") == "bpe":
51
+ from indigotf.bpe import BPETokenizer
52
+
53
+ return BPETokenizer.from_state(tokenizer_info)
54
+ from indigotf.tokenizer import CharTokenizer
55
+
56
+ return CharTokenizer(vocab)
indigotf/model.py ADDED
@@ -0,0 +1,61 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import tensorflow as tf
2
+ from keras import layers
3
+
4
+
5
+ class PositionEmbedding(layers.Layer):
6
+ def __init__(self, block_size, **kwargs):
7
+ super().__init__(**kwargs)
8
+ self.block_size = block_size
9
+
10
+ def build(self, input_shape):
11
+ self.pos_emb = self.add_weight(
12
+ name="pos_emb", shape=(self.block_size, input_shape[-1]), initializer="random_normal"
13
+ )
14
+
15
+ def call(self, x):
16
+ T = tf.shape(x)[1]
17
+ return x + self.pos_emb[tf.newaxis, :T, :]
18
+
19
+ def get_config(self):
20
+ config = super().get_config()
21
+ config["block_size"] = self.block_size
22
+ return config
23
+
24
+
25
+ def build_gpt(vocab_size, block_size, n_layer=4, n_head=4, n_embd=128, dropout=0.1, name="indigo"):
26
+ tokens = tf.keras.Input(shape=(None,), dtype="int64", name="tokens")
27
+ x = layers.Embedding(vocab_size, n_embd, name="tok_emb")(tokens)
28
+ x = PositionEmbedding(block_size, name="pos_emb")(x)
29
+ x = layers.Dropout(dropout)(x)
30
+ for i in range(n_layer):
31
+ xn = layers.LayerNormalization(epsilon=1e-5, name=f"ln1_{i}")(x)
32
+ attn = layers.MultiHeadAttention(
33
+ num_heads=n_head, key_dim=n_embd // n_head, dropout=dropout, name=f"attn_{i}"
34
+ )
35
+ x = x + attn(xn, xn, use_causal_mask=True)
36
+ xn = layers.LayerNormalization(epsilon=1e-5, name=f"ln2_{i}")(x)
37
+ h = layers.Dense(4 * n_embd, activation="gelu", name=f"fc_{i}")(xn)
38
+ h = layers.Dense(n_embd, name=f"proj_{i}")(h)
39
+ h = layers.Dropout(dropout)(h)
40
+ x = x + h
41
+ x = layers.LayerNormalization(epsilon=1e-5, name="ln_f")(x)
42
+ logits = layers.Dense(vocab_size, use_bias=False, name="head")(x)
43
+ return tf.keras.Model(tokens, logits, name=name)
44
+
45
+
46
+ def generate(model, idx, max_new_tokens, block_size, temperature=1.0, top_k=None):
47
+ for _ in range(max_new_tokens):
48
+ idx_cond = idx[:, -block_size:]
49
+ logits = model(idx_cond, training=False)[:, -1, :]
50
+ logits = logits / max(temperature, 1e-8)
51
+ if top_k is not None:
52
+ k = min(top_k, int(logits.shape[-1]))
53
+ vals, _ = tf.math.top_k(logits, k=k)
54
+ logits = tf.where(
55
+ logits < vals[:, -1:],
56
+ tf.fill(tf.shape(logits), tf.float32.min),
57
+ logits,
58
+ )
59
+ next_id = tf.random.categorical(logits, num_samples=1, dtype=tf.int64)
60
+ idx = tf.concat([idx, next_id], axis=1)
61
+ return idx
indigotf/tokenizer.py ADDED
@@ -0,0 +1,33 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import json
2
+
3
+
4
+ class CharTokenizer:
5
+ def __init__(self, vocab=None):
6
+ if vocab is None:
7
+ self.itos = []
8
+ else:
9
+ self.itos = list(vocab)
10
+ self.stoi = {ch: i for i, ch in enumerate(self.itos)}
11
+
12
+ @classmethod
13
+ def from_text(cls, text):
14
+ return cls(sorted(set(text)))
15
+
16
+ @property
17
+ def vocab_size(self):
18
+ return len(self.itos)
19
+
20
+ def encode(self, text):
21
+ return [self.stoi[ch] for ch in text if ch in self.stoi]
22
+
23
+ def decode(self, ids):
24
+ return "".join(self.itos[i] for i in ids)
25
+
26
+ def save(self, path):
27
+ with open(path, "w", encoding="utf-8") as f:
28
+ json.dump(self.itos, f, ensure_ascii=False)
29
+
30
+ @classmethod
31
+ def load(cls, path):
32
+ with open(path, encoding="utf-8") as f:
33
+ return cls(json.load(f))
requirements.txt ADDED
@@ -0,0 +1,2 @@
 
 
 
1
+ tensorflow>=2.16
2
+ safetensors>=0.4
train.py ADDED
@@ -0,0 +1,224 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import time
3
+ import math
4
+ import random
5
+ import argparse
6
+ import numpy as np
7
+ import tensorflow as tf
8
+
9
+ from safetensors.numpy import load_file, save_file
10
+
11
+ from indigotf.bpe import BPETokenizer
12
+ from indigotf.common import (
13
+ build_tokenizer,
14
+ collect_text_files,
15
+ load_meta,
16
+ read_clean,
17
+ save_meta,
18
+ )
19
+ from indigotf.model import build_gpt
20
+ from indigotf.tokenizer import CharTokenizer
21
+
22
+ CONFIG_KEYS = ("vocab_size", "block_size", "n_layer", "n_head", "n_embd", "dropout")
23
+
24
+
25
+ def get_batch(np_data, block_size, batch_size):
26
+ ix = np.random.randint(0, len(np_data) - block_size - 1, size=batch_size)
27
+ x = np.stack([np_data[i : i + block_size] for i in ix])
28
+ y = np.stack([np_data[i + 1 : i + block_size + 1] for i in ix])
29
+ return tf.constant(x), tf.constant(y)
30
+
31
+
32
+ @tf.function(reduce_retracing=True)
33
+ def train_step(model, optimizer, loss_fn, x, y):
34
+ with tf.GradientTape() as tape:
35
+ logits = model(x, training=True)
36
+ loss = loss_fn(y, logits)
37
+ grads = tape.gradient(loss, model.trainable_variables)
38
+ optimizer.apply_gradients(zip(grads, model.trainable_variables))
39
+ return loss
40
+
41
+
42
+ @tf.function(reduce_retracing=True)
43
+ def eval_step(model, loss_fn, x, y):
44
+ logits = model(x, training=False)
45
+ return loss_fn(y, logits)
46
+
47
+
48
+ def save_weights_tf(model, base_path, config, vocab, step, val_loss, tinfo):
49
+ tensors = {v.path: np.asarray(v) for v in model.weights}
50
+ save_file(tensors, base_path)
51
+ save_meta(base_path, config, vocab, step, val_loss, backend="tensorflow", tokenizer=tinfo)
52
+
53
+
54
+ def main():
55
+ parser = argparse.ArgumentParser(description="Latih model Indigo-TF (backend TensorFlow/Keras)")
56
+ parser.add_argument("--data", nargs="+", default=["data/sample.txt"])
57
+ parser.add_argument("--out", default="out")
58
+ parser.add_argument("--steps", type=int, default=2000)
59
+ parser.add_argument("--batch-size", type=int, default=32)
60
+ parser.add_argument("--block-size", type=int, default=128)
61
+ parser.add_argument("--n-layer", type=int, default=4)
62
+ parser.add_argument("--n-head", type=int, default=4)
63
+ parser.add_argument("--n-embd", type=int, default=128)
64
+ parser.add_argument("--dropout", type=float, default=0.1)
65
+ parser.add_argument("--lr", type=float, default=3e-4)
66
+ parser.add_argument("--warmup", type=int, default=100)
67
+ parser.add_argument("--weight-decay", type=float, default=0.1)
68
+ parser.add_argument("--eval-interval", type=int, default=200)
69
+ parser.add_argument("--eval-iters", type=int, default=20)
70
+ parser.add_argument("--seed", type=int, default=1337)
71
+ parser.add_argument("--init-from", default=None, help="checkpoint safetensors sebelumnya")
72
+ parser.add_argument("--tokenizer", default="char", choices=["char", "bpe"])
73
+ parser.add_argument("--vocab-size", type=int, default=512)
74
+ parser.add_argument("--device", default="auto", choices=["auto", "cpu", "gpu"])
75
+ args = parser.parse_args()
76
+
77
+ if args.device == "cpu":
78
+ tf.config.set_visible_devices([], "GPU")
79
+ gpus = tf.config.list_physical_devices("GPU")
80
+ device_label = f"gpu({len(gpus)})" if gpus and args.device != "cpu" else "cpu"
81
+
82
+ random.seed(args.seed)
83
+ np.random.seed(args.seed)
84
+ tf.random.set_seed(args.seed)
85
+ os.makedirs(args.out, exist_ok=True)
86
+
87
+ files = sorted(collect_text_files(args.data))
88
+ if not files:
89
+ raise SystemExit("tidak ada file teks ditemukan")
90
+ rng = random.Random(args.seed)
91
+ rng.shuffle(files)
92
+ n_val = max(1, round(len(files) * 0.1)) if len(files) > 1 else 0
93
+ print(f"file latih={len(files) - n_val} | file validasi={n_val}")
94
+
95
+ train_text = "".join(read_clean(p) for p in files[n_val:])
96
+ val_text = "".join(read_clean(p) for p in files[:n_val])
97
+ all_text = train_text + val_text
98
+
99
+ start_step = 0
100
+ init_state = None
101
+ if args.init_from:
102
+ meta = load_meta(args.init_from)
103
+ if meta.get("backend") != "tensorflow":
104
+ raise SystemExit(f"{args.init_from} bukan checkpoint backend TensorFlow")
105
+ config_d = meta["config"]
106
+ start_step = meta.get("step", 0)
107
+ init_state = {k.replace("/", "_"): v for k, v in load_file(args.init_from).items()}
108
+ print(f"melanjutkan dari {args.init_from} (step {start_step})")
109
+ tokenizer = build_tokenizer(meta.get("tokenizer") or {"type": "char"}, meta.get("vocab"))
110
+ tinfo = meta.get("tokenizer") or {"type": "char"}
111
+ else:
112
+ if args.tokenizer == "bpe":
113
+ tokenizer = BPETokenizer.train(all_text, args.vocab_size)
114
+ tinfo = tokenizer.state()
115
+ else:
116
+ tokenizer = CharTokenizer.from_text(all_text)
117
+ tinfo = {"type": "char"}
118
+ config_d = {
119
+ "vocab_size": tokenizer.vocab_size,
120
+ "block_size": args.block_size,
121
+ "n_layer": args.n_layer,
122
+ "n_head": args.n_head,
123
+ "n_embd": args.n_embd,
124
+ "dropout": args.dropout,
125
+ "bias": False,
126
+ }
127
+ if config_d["vocab_size"] != tokenizer.vocab_size:
128
+ raise SystemExit(
129
+ f"vocab tidak cocok: checkpoint={config_d['vocab_size']}, tokenizer={tokenizer.vocab_size}"
130
+ )
131
+
132
+ train_np = np.array(tokenizer.encode(train_text), dtype=np.int64)
133
+ val_np = np.array(tokenizer.encode(val_text), dtype=np.int64)
134
+ if len(train_np) < config_d["block_size"] * 2:
135
+ raise SystemExit(f"data latih terlalu pendek ({len(train_np)} token)")
136
+ print(
137
+ f"tokenizer={tinfo['type']} | tokens latih={len(train_np):,} | "
138
+ f"tokens validasi={len(val_np):,} | vocab={tokenizer.vocab_size}"
139
+ )
140
+ total_steps = start_step + args.steps
141
+
142
+ model = build_gpt(
143
+ vocab_size=config_d["vocab_size"],
144
+ block_size=config_d["block_size"],
145
+ n_layer=config_d["n_layer"],
146
+ n_head=config_d["n_head"],
147
+ n_embd=config_d["n_embd"],
148
+ dropout=config_d["dropout"],
149
+ )
150
+ if init_state is not None:
151
+ by_path = {v.path.replace("/", "_"): v for v in model.weights}
152
+ missing = [p for p in by_path if p not in init_state]
153
+ if missing:
154
+ raise SystemExit(f"bobot tidak cocok dengan checkpoint: {missing[:5]}")
155
+ model.set_weights([init_state[v.path.replace("/", "_")] for v in model.weights])
156
+
157
+ n_params = int(sum(int(np.prod(v.shape)) for v in model.weights))
158
+ print(
159
+ f"device={device_label} | params={n_params / 1e6:.2f}M | "
160
+ f"vocab={config_d['vocab_size']} | total_steps={total_steps}"
161
+ )
162
+
163
+ optimizer = tf.keras.optimizers.AdamW(
164
+ learning_rate=args.lr, beta_1=0.9, beta_2=0.95, weight_decay=args.weight_decay, clipnorm=1.0
165
+ )
166
+ loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
167
+
168
+ def lr_at(step):
169
+ if step < args.warmup:
170
+ return args.lr * (step + 1) / args.warmup
171
+ progress = (step - args.warmup) / max(1, total_steps - args.warmup)
172
+ return 0.1 * args.lr + 0.45 * args.lr * (1 + math.cos(math.pi * progress))
173
+
174
+ best_val = float("inf")
175
+ last_val = None
176
+ t0 = time.time()
177
+ for step in range(start_step, total_steps):
178
+ optimizer.learning_rate.assign(lr_at(step))
179
+ x, y = get_batch(train_np, config_d["block_size"], args.batch_size)
180
+ loss = train_step(model, optimizer, loss_fn, x, y)
181
+ if step % args.eval_interval == 0 or step == total_steps - 1:
182
+ if len(val_np) > config_d["block_size"] + 1:
183
+ losses = []
184
+ for _ in range(args.eval_iters):
185
+ vx, vy = get_batch(val_np, config_d["block_size"], args.batch_size)
186
+ losses.append(float(eval_step(model, loss_fn, vx, vy)))
187
+ val_loss = sum(losses) / len(losses)
188
+ marker = ""
189
+ if val_loss < best_val:
190
+ best_val = val_loss
191
+ save_weights_tf(
192
+ model,
193
+ os.path.join(args.out, "indigo_best.safetensors"),
194
+ config_d,
195
+ tokenizer.itos if hasattr(tokenizer, "itos") else None,
196
+ total_steps,
197
+ val_loss,
198
+ tinfo,
199
+ )
200
+ marker = " <- best"
201
+ last_val = val_loss
202
+ val_str = f"{val_loss:.4f}{marker}"
203
+ else:
204
+ val_str = "n/a"
205
+ print(
206
+ f"step {step + 1:5d}/{total_steps} | "
207
+ f"loss {float(loss):.4f} | val {val_str} | {time.time() - t0:.1f}s"
208
+ )
209
+
210
+ final_path = os.path.join(args.out, "indigo.safetensors")
211
+ save_weights_tf(
212
+ model,
213
+ final_path,
214
+ config_d,
215
+ tokenizer.itos if hasattr(tokenizer, "itos") else None,
216
+ total_steps,
217
+ last_val,
218
+ tinfo,
219
+ )
220
+ print(f"model tersimpan di {final_path} (+_meta.json)")
221
+
222
+
223
+ if __name__ == "__main__":
224
+ main()