adyoi
/

indigo

@@ -1,197 +0,0 @@
-import argparse
-import math
-import os
-import random
-import time
-import numpy as np
-import tensorflow as tf
-from safetensors.numpy import load_file, save_file
-from indigo.common import collect_text_files, load_meta, read_clean, save_meta
-from indigo.model_keras import build_gpt
-from indigo.tokenizer import CharTokenizer
-CONFIG_KEYS = ("vocab_size", "block_size", "n_layer", "n_head", "n_embd", "dropout")
-def get_batch(np_data, block_size, batch_size):
-    ix = np.random.randint(0, len(np_data) - block_size - 1, size=batch_size)
-    x = np.stack([np_data[i : i + block_size] for i in ix])
-    y = np.stack([np_data[i + 1 : i + block_size + 1] for i in ix])
-    return tf.constant(x), tf.constant(y)
-@tf.function(reduce_retracing=True)
-def train_step(model, optimizer, loss_fn, x, y):
-    with tf.GradientTape() as tape:
-        logits = model(x, training=True)
-        loss = loss_fn(y, logits)
-    grads = tape.gradient(loss, model.trainable_variables)
-    optimizer.apply_gradients(zip(grads, model.trainable_variables))
-    return loss
-@tf.function(reduce_retracing=True)
-def eval_step(model, loss_fn, x, y):
-    logits = model(x, training=False)
-    return loss_fn(y, logits)
-@tf.function(reduce_retracing=True)
-def forward_last(model, x):
-    return model(x, training=False)[:, -1, :]
-def save_weights_tf(model, base_path, config, vocab, step, val_loss):
-    tensors = {v.path: np.asarray(v) for v in model.weights}
-    save_file(tensors, base_path)
-    save_meta(base_path, config, vocab, step, val_loss, backend="tensorflow")
-def main():
-    parser = argparse.ArgumentParser(description="Latih model Indigo (backend TensorFlow/Keras)")
-    parser.add_argument("--data", nargs="+", default=["data/sample.txt"])
-    parser.add_argument("--out", default="out_tf")
-    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 safetensors dari backend TF")
-    parser.add_argument("--device", default="auto", choices=["auto", "cpu", "gpu"])
-    args = parser.parse_args()
-    if args.device == "cpu":
-        tf.config.set_visible_devices([], "GPU")
-    gpus = tf.config.list_physical_devices("GPU")
-    device_label = f"gpu({len(gpus)})" if gpus and args.device != "cpu" else "cpu"
-    random.seed(args.seed)
-    np.random.seed(args.seed)
-    tf.random.set_seed(args.seed)
-    os.makedirs(args.out, exist_ok=True)
-    files = sorted(collect_text_files(args.data))
-    if not files:
-        raise SystemExit("tidak ada file teks ditemukan")
-    rng = random.Random(args.seed)
-    rng.shuffle(files)
-    n_val = max(1, round(len(files) * 0.1)) if len(files) > 1 else 0
-    print(f"file latih={len(files) - n_val} | file validasi={n_val}")
-    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_np = np.array(tokenizer.encode(train_text), dtype=np.int64)
-    val_np = np.array(tokenizer.encode(val_text), dtype=np.int64)
-    if len(train_np) < args.block_size * 2:
-        raise SystemExit(f"data latih terlalu pendek ({len(train_np)} token)")
-    print(
-        f"tokens latih={len(train_np):,} | tokens validasi={len(val_np):,} | vocab={tokenizer.vocab_size}"
-    )
-    start_step = 0
-    init_state = None
-    if args.init_from:
-        meta = load_meta(args.init_from)
-        if meta.get("backend") != "tensorflow":
-            raise SystemExit(f"{args.init_from} bukan checkpoint backend TensorFlow")
-        config_d = meta["config"]
-        start_step = meta.get("step", 0)
-        init_state = {k.replace("/", "_"): v for k, v in load_file(args.init_from).items()}
-        print(f"melanjutkan dari {args.init_from} (step {start_step})")
-    else:
-        config_d = {
-            "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,
-            "bias": False,
-        }
-        if config_d["vocab_size"] != tokenizer.vocab_size:
-            raise SystemExit("vocab tidak cocok")
-    total_steps = start_step + args.steps
-    model = build_gpt(
-        vocab_size=config_d["vocab_size"],
-        block_size=config_d["block_size"],
-        n_layer=config_d["n_layer"],
-        n_head=config_d["n_head"],
-        n_embd=config_d["n_embd"],
-        dropout=config_d["dropout"],
-    )
-    if init_state is not None:
-        by_path = {v.path.replace("/", "_"): v for v in model.weights}
-        missing = [p for p in by_path if p not in init_state]
-        if missing:
-            raise SystemExit(f"bobot tidak cocok dengan checkpoint: {missing[:5]}")
-        model.set_weights([init_state[v.path.replace("/", "_")] for v in model.weights])
-    n_params = int(sum(int(np.prod(v.shape)) for v in model.weights))
-    print(f"device={device_label} | params={n_params / 1e6:.2f}M | vocab={config_d['vocab_size']} | total_steps={total_steps}")
-    optimizer = tf.keras.optimizers.AdamW(
-        learning_rate=args.lr, beta_1=0.9, beta_2=0.95, weight_decay=args.weight_decay, clipnorm=1.0
-    )
-    loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
-    def lr_at(step):
-        if step < args.warmup:
-            return args.lr * (step + 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))
-    best_val = float("inf")
-    last_val = None
-    t0 = time.time()
-    for step in range(start_step, total_steps):
-        optimizer.learning_rate.assign(lr_at(step))
-        x, y = get_batch(train_np, config_d["block_size"], args.batch_size)
-        loss = train_step(model, optimizer, loss_fn, x, y)
-        if step % args.eval_interval == 0 or step == total_steps - 1:
-            if len(val_np) > config_d["block_size"] + 1:
-                losses = []
-                for _ in range(args.eval_iters):
-                    vx, vy = get_batch(val_np, config_d["block_size"], args.batch_size)
-                    losses.append(float(eval_step(model, loss_fn, vx, vy)))
-                val_loss = sum(losses) / len(losses)
-                marker = ""
-                if val_loss < best_val:
-                    best_val = val_loss
-                    save_weights_tf(
-                        model,
-                        os.path.join(args.out, "indigo_best.safetensors"),
-                        config_d,
-                        tokenizer.itos,
-                        total_steps,
-                        val_loss,
-                    )
-                    marker = " <- best"
-                last_val = val_loss
-                val_str = f"{val_loss:.4f}{marker}"
-            else:
-                val_str = "n/a"
-            print(
-                f"step {step + 1:5d}/{total_steps} | "
-                f"loss {float(loss):.4f} | val {val_str} | {time.time() - t0:.1f}s"
-            )
-    final_path = os.path.join(args.out, "indigo.safetensors")
-    save_weights_tf(model, final_path, config_d, tokenizer.itos, total_steps, last_val)
-    print(f"model tersimpan di {final_path} (+_meta.json)")
-if __name__ == "__main__":
-    main()