Instructions to use adyoi/indigo.tf with libraries, inference providers, notebooks, and local apps. Follow these links to get started.
- Libraries
- Keras
How to use adyoi/indigo.tf with Keras:
# Available backend options are: "jax", "torch", "tensorflow". import os os.environ["KERAS_BACKEND"] = "jax" import keras model = keras.saving.load_model("hf://adyoi/indigo.tf") - Notebooks
- Google Colab
- Kaggle
Upload folder using huggingface_hub
Browse files- .gitignore +5 -0
- README.md +28 -0
- data/sample.txt +50 -0
- generate.py +67 -0
- indigotf/__init__.py +5 -0
- indigotf/bpe.py +75 -0
- indigotf/common.py +56 -0
- indigotf/model.py +61 -0
- indigotf/tokenizer.py +33 -0
- requirements.txt +2 -0
- train.py +224 -0
.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()
|