File size: 2,387 Bytes
47dfda4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 | import sys
import random
import argparse
import numpy as np
import tensorflow as tf
from indigotf.common import build_tokenizer, load_meta
from indigotf.model import build_gpt, generate
CONFIG_KEYS = ("vocab_size", "block_size", "n_layer", "n_head", "n_embd", "dropout")
def load_model(path):
from safetensors.numpy import load_file
meta = load_meta(path)
if meta.get("backend") != "tensorflow":
raise SystemExit(
f"{path} berasal dari backend {meta.get('backend')}, gunakan repo indigo (PyTorch)"
)
state = load_file(path)
by_path = {k.replace("/", "_"): v for k, v in state.items()}
model = build_gpt(**{k: meta["config"][k] for k in CONFIG_KEYS})
missing = [v.path for v in model.weights if v.path.replace("/", "_") not in by_path]
if missing:
raise SystemExit(f"bobot tidak cocok dengan checkpoint: {missing[:5]}")
model.set_weights([by_path[v.path.replace("/", "_")] for v in model.weights])
tokenizer = build_tokenizer(meta.get("tokenizer") or {"type": "char"}, meta.get("vocab"))
return model, meta, tokenizer
def main():
sys.stdout.reconfigure(encoding="utf-8", errors="replace")
parser = argparse.ArgumentParser(description="Generate teks dari checkpoint Indigo-TF")
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("--seed", type=int, default=None)
parser.add_argument("--device", default="auto", choices=["auto", "cpu", "gpu"])
args = parser.parse_args()
if args.device == "cpu":
tf.config.set_visible_devices([], "GPU")
if args.seed is not None:
random.seed(args.seed)
np.random.seed(args.seed)
tf.random.set_seed(args.seed)
model, meta, tokenizer = load_model(args.ckpt)
ids = tokenizer.encode(args.prompt) or [0]
idx = tf.constant([ids], dtype=tf.int64)
out = generate(
model,
idx,
args.max_new,
block_size=meta["config"]["block_size"],
temperature=args.temperature,
top_k=args.top_k,
)
print(tokenizer.decode(out.numpy()[0].tolist()))
if __name__ == "__main__":
main()
|