pipeline.py

"""
Pipeline end-to-end untuk training Indigo.

Alur kerja:
1. Kumpulkan data: copy file lokal + tarik dataset dari Hugging Face
2. Konversi JSON Alpaca ke .txt (opsional, --format-qa)
3. Jalankan training (train.py) dengan semua hyperparameter
4. Evaluasi guard kamus (opsional, --guard)
5. Simpan manifest.json (statistik + metadata run)
6. Finalize: promosikan checkpoint terbaik ke out/ (opsional, --finalize)
7. Push ke Hugging Face Hub (opsional, --push)

Cara pakai:
    python pipeline.py --tag run01 --data data/sample.txt --steps 2000
    python pipeline.py --tag run02 --hf-dataset adyoi/indigo --steps 900 --device cpu
    python pipeline.py --tag run03 --data data/ --steps 2000 --finalize --push
    python pipeline.py --tag qa01 --hf-dataset rohanrdy/CS-Theory-QA-Dataset --format-qa --steps 3000
"""

import argparse
import datetime
import json
import os
import shutil
import subprocess
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parent
sys.path.insert(0, str(ROOT))

import train


def preprocess_alpaca_json(data_dir, patterns=("*.json",)):
    """Konversi file JSON Alpaca format ke .txt flat.

    Format Alpaca: [{"instruction": "...", "input": "...", "output": "..."}, ...]
    Format output:
        ### instruction
        (input jika ada)

        output

    Fungsi ini auto-detect: hanya mengkonversi file JSON yang punya field "instruction" dan "output".
    File yang sudah ada .txt-nya tidak ditimpa.

    Args:
        data_dir:  Direktori yang berisi file JSON (akan dicari rekursif).
        patterns:  Pola file JSON (default: "*.json").

    Returns:
        Int — jumlah file JSON yang berhasil dikonversi.
    """
    converted = 0
    for pat in patterns:
        for jf in Path(data_dir).rglob(pat):
            if ".cache" in jf.parts:
                continue
            try:
                with open(jf, encoding="utf-8") as f:
                    data = json.load(f)
            except (json.JSONDecodeError, UnicodeDecodeError):
                continue
            if not isinstance(data, list) or not data:
                continue
            # Cek apakah ini format Alpaca (punya field instruction + output)
            keys = set(data[0].keys()) if isinstance(data[0], dict) else set()
            if not {"instruction", "output"}.issubset(keys):
                continue
            lines = []
            for row in data:
                inst = str(row.get("instruction", "")).strip()
                inp = str(row.get("input", "")).strip()
                out = str(row.get("output", "")).strip()
                if not inst or not out:
                    continue
                q = inst if not inp else f"{inst}\n{inp}"
                lines.append(f"### {q}\n\n{out}\n\n")
            if not lines:
                continue
            # Simpan .txt di sebelah file JSON asli
            txt_path = jf.with_suffix(".txt")
            txt_path.write_text("".join(lines), encoding="utf-8")
            print(f"[pipeline] alpaca: {jf.name} -> {txt_path.name} ({len(lines)} QA)")
            converted += 1
    return converted


def git_short():
    """Ambil hash commit git pendek (7 karakter) atau "nogit" jika bukan repo git."""
    try:
        return subprocess.check_output(
            ["git", "rev-parse", "--short", "HEAD"], cwd=str(ROOT), text=True
        ).strip()
    except Exception:
        return "nogit"


def pull_hf(repo_id, dest, patterns):
    """Tarik dataset dari Hugging Face Hub ke direktori lokal.

    Menggunakan snapshot_download dengan allow_patterns untuk mengunduh
    hanya file yang sesuai pola (mis. *.txt, *.md, *.json).

    Args:
        repo_id:  ID repo HF (mis. "adyoi/indigo" atau "rohanrdy/CS-Theory-QA-Dataset").
        dest:     Direktori tujuan download.
        patterns: Pola file yang akan diunduh.

    Returns:
        Str — path ke direktori lokal yang berisi file yang diunduh.
    """
    from huggingface_hub import snapshot_download

    path = snapshot_download(
        repo_id=repo_id,
        repo_type="dataset",
        allow_patterns=patterns,
        local_dir=str(dest),
    )
    n = sum(
        1 for f in Path(path).rglob("*")
        if f.is_file() and ".cache" not in f.parts and f.suffix.lstrip(".") in
        {p.lstrip("*.") for p in patterns}
    )
    print(f"[pipeline] dataset '{repo_id}' ditarik -> {dest} ({n} file)")
    return str(dest)


def main():
    ap = argparse.ArgumentParser(
        description="Pipeline latih Indigo: kumpul data -> tokenizer -> training -> manifest -> finalize/push"
    )

    # --- Identitas Run ---
    ap.add_argument("--tag", required=True,
                    help="nama run; semua artefak disimpan di runs/<tag>/")

    # --- Sumber Data ---
    ap.add_argument("--data", nargs="*", default=[],
                    help="file/folder teks lokal tambahan (banyak, spasi-separated)")
    ap.add_argument("--hf-dataset", default=None,
                    help="repo dataset HF untuk menarik file teks (mis. adyoi/indigo)")
    ap.add_argument("--hf-patterns", nargs="*", default=["*.txt", "*.md"],
                    help="pola file yang diambil dari HF (default: *.txt *.md)")
    ap.add_argument("--format-qa", action="store_true",
                    help="auto-convert Alpaca JSON (instruction/output) ke .txt sebelum training")

    # --- Output & Finalize ---
    ap.add_argument("--runs", default="runs",
                    help="folder root untuk semua run (default: runs)")
    ap.add_argument("--device", default=None,
                    help="device untuk training (diteruskan ke train.py, default: auto)")
    ap.add_argument("--finalize", action="store_true",
                    help="promosikan checkpoint terbaik run ini ke folder out/ kanonik")
    ap.add_argument("--push", action="store_true",
                    help="upload checkpoint terbaik ke repo HF")
    ap.add_argument("--repo", default="adyoi/indigo",
                    help="repo HF tujuan upload (default: adyoi/indigo)")

    # --- Guard Kamus ---
    ap.add_argument("--guard", default=None,
                    help="file kamus kata (satu kata per baris); tambahkan metrik rasio ejaan ke manifest")
    ap.add_argument("--guard-max-new", type=int, default=120,
                    help="jumlah token generate untuk evaluasi guard (default: 120)")
    ap.add_argument("--guard-prefiks", default=None,
                    help="file prefiks Indonesia (default: data/prefiks.txt bila ada)")
    ap.add_argument("--guard-sufiks", default=None,
                    help="file sufiks Indonesia (default: data/sufiks.txt bila ada)")

    # --- Hyperparameter Training (diteruskan ke train.py) ---
    t = ap.add_argument_group("hyperparameter training (diteruskan ke train.py)")
    t.add_argument("--steps", type=int, default=2000,
                   help="jumlah langkah training (default: 2000)")
    t.add_argument("--batch-size", type=int, default=32,
                   help="batch size (default: 32)")
    t.add_argument("--block-size", type=int, default=128,
                   help="panjang konteks (default: 128)")
    t.add_argument("--n-layer", type=int, default=4,
                   help="jumlah layer transformer (default: 4)")
    t.add_argument("--n-head", type=int, default=4,
                   help="jumlah head attention (default: 4)")
    t.add_argument("--n-embd", type=int, default=128,
                   help="dimensi embedding (default: 128)")
    t.add_argument("--dropout", type=float, default=0.1,
                   help="dropout rate (default: 0.1)")
    t.add_argument("--lr", type=float, default=3e-4,
                   help="learning rate (default: 3e-4)")
    t.add_argument("--warmup", type=int, default=100,
                   help="langkah warmup (default: 100)")
    t.add_argument("--weight-decay", type=float, default=0.1,
                   help="weight decay (default: 0.1)")
    t.add_argument("--eval-interval", type=int, default=200,
                   help="evaluasi setiap N langkah (default: 200)")
    t.add_argument("--eval-iters", type=int, default=20,
                   help="jumlah batch evaluasi (default: 20)")
    t.add_argument("--seed", type=int, default=1337,
                   help="seed random (default: 1337)")
    t.add_argument("--init-from", default=None,
                   help="checkpoint untuk resume training")
    t.add_argument("--tokenizer", default="char", choices=["char", "bpe"],
                   help="jenis tokenizer (default: char)")
    t.add_argument("--vocab-size", type=int, default=512,
                   help="vocab size untuk BPE (default: 512)")
    t.add_argument("--val-fraction", type=float, default=0.1,
                   help="proporsi file validasi (default: 0.1)")
    args = ap.parse_args()

    # --- Siapkan direktori run ---
    run_dir = ROOT / args.runs / args.tag
    data_dir = run_dir / "data"
    data_dir.mkdir(parents=True, exist_ok=True)

    # --- Kumpulkan data lokal ---
    sources = []
    for p in args.data:
        src = Path(p)
        if not src.exists():
            raise SystemExit(f"sumber data tidak ada: {src}")
        if src.is_dir():
            shutil.copytree(src, data_dir / src.name, dirs_exist_ok=True)
        else:
            shutil.copy2(src, data_dir / src.name)
        sources.append(str(src))

    # --- Tarik dataset dari HF ---
    if args.hf_dataset:
        patterns = list(args.hf_patterns)
        if args.format_qa and "*.json" not in patterns:
            patterns.append("*.json")  # butuh JSON untuk Alpaca conversion
        pull_hf(args.hf_dataset, data_dir / "_hf", patterns)
        sources.append(f"hf:{args.hf_dataset}")

    # --- Konversi JSON Alpaca → .txt ---
    if args.format_qa:
        n = preprocess_alpaca_json(data_dir)
        if n:
            print(f"[pipeline] {n} file JSON dikonversi ke .txt")
        else:
            print("[pipeline] peringatan: tidak ada file Alpaca JSON ditemukan")

    # --- Jalankan training ---
    out_dir = run_dir / "ckpt"
    argv = [
        "--data", str(data_dir),
        "--out", str(out_dir),
        "--steps", str(args.steps),
        "--batch-size", str(args.batch_size),
        "--block-size", str(args.block_size),
        "--n-layer", str(args.n_layer),
        "--n-head", str(args.n_head),
        "--n-embd", str(args.n_embd),
        "--dropout", str(args.dropout),
        "--lr", str(args.lr),
        "--warmup", str(args.warmup),
        "--weight-decay", str(args.weight_decay),
        "--eval-interval", str(args.eval_interval),
        "--eval-iters", str(args.eval_iters),
        "--seed", str(args.seed),
        "--tokenizer", args.tokenizer,
        "--vocab-size", str(args.vocab_size),
        "--val-fraction", str(args.val_fraction),
    ]
    if args.init_from:
        argv += ["--init-from", args.init_from]
    if args.device:
        argv += ["--device", args.device]

    print(f"[pipeline] run='{args.tag}' | sumber={sources or ['<kosong>' ]}")
    stats = train.main(argv) or {}

    # --- Evaluasi guard kamus (opsional) ---
    if args.guard:
        import torch

        from indigo.common import load_wordlist, load_meta, build_tokenizer, word_known_ratio
        from indigo.model import GPT, GPTConfig
        from safetensors.torch import load_file

        # Cari checkpoint: prioritaskan best, fallback ke final
        probe_path = out_dir / "indigo_best.safetensors"
        if not probe_path.exists():
            probe_path = out_dir / "indigo.safetensors"
        wordset = load_wordlist(args.guard)
        p_def, s_def = ROOT / "data" / "prefiks.txt", ROOT / "data" / "sufiks.txt"

        def muat_afiks(flag, default_path):
            """Muat file afiks: prioritaskan flag CLI → default path."""
            if flag:
                return load_wordlist(flag) if Path(flag).exists() else None
            return load_wordlist(str(default_path)) if default_path.exists() else None

        prefiks = muat_afiks(args.guard_prefiks, p_def)
        sufiks = muat_afiks(args.guard_sufiks, s_def)
        meta = load_meta(str(probe_path))
        tok = build_tokenizer(meta.get("tokenizer") or {"type": "char"}, meta.get("vocab"))
        mdl = GPT(GPTConfig(**meta["config"]))
        mdl.load_state_dict(load_file(str(probe_path)), strict=False)
        mdl.eval()
        # Generate teks dari prompt newline (netral)
        ids = tok.encode("\n") or [0]
        idx = torch.tensor([ids], dtype=torch.long)
        with torch.no_grad():
            out = mdl.generate(idx, args.guard_max_new, temperature=0.8, top_k=40)
        text = tok.decode(out[0].tolist())
        stats["kamus_ratio"] = round(word_known_ratio(text, wordset, prefiks, sufiks), 4)
        print(f"[pipeline] kamus_ratio={stats['kamus_ratio']:.2%}")

    # --- Simpan manifest.json ---
    stats["pipeline"] = {
        "tag": args.tag,
        "timestamp": datetime.datetime.now().isoformat(timespec="seconds"),
        "git": git_short(),
        "sources": sources,
        "hf_dataset": args.hf_dataset,
        "run_dir": str(run_dir),
    }

    manifest = run_dir / "manifest.json"
    with open(manifest, "w", encoding="utf-8") as f:
        json.dump(stats, f, ensure_ascii=False, indent=2)
    print(f"[pipeline] manifest -> {manifest}")

    # --- Finalize: promosikan checkpoint terbaik ke out/ ---
    best = out_dir / "indigo_best.safetensors"
    if args.finalize:
        has_best = (out_dir / "indigo_best.safetensors").exists()
        has_final = (out_dir / "indigo.safetensors").exists()
        if not has_best and not has_final:
            raise SystemExit("finalize dibatalkan: tidak ada checkpoint")
        if not has_best:
            print("[pipeline] peringatan: tanpa data validasi; memakai checkpoint akhir sebagai best")
        dst = ROOT / "out"
        dst.mkdir(exist_ok=True)
        copied = []
        for name in ("indigo_best.safetensors", "indigo_best_meta.json",
                     "indigo.safetensors", "indigo_meta.json", "indigo_optimizer.pt"):
            srcp = out_dir / name
            if srcp.exists():
                shutil.copy2(srcp, dst / name)
                copied.append(name)
        print(f"[pipeline] finalisasi ke out/: {', '.join(copied)}")
        info = {
            "tag": args.tag,
            "timestamp": datetime.datetime.now().isoformat(timespec="seconds"),
            "git": git_short(),
            "best_val": stats.get("best_val"),
            "nats_per_char_best": stats.get("nats_per_char_best"),
            "kamus_ratio": stats.get("kamus_ratio"),
            "sources": sources,
        }
        (dst / "run_info.json").write_text(
            json.dumps(info, ensure_ascii=False, indent=2), encoding="utf-8"
        )
        print("[pipeline] silsilah -> out/run_info.json")

    # --- Push ke Hugging Face Hub ---
    if args.push:
        from huggingface_hub import HfApi

        api = HfApi()
        api.upload_file(
            path_or_fileobj=str(best),
            path_in_repo="out/indigo_best.safetensors",
            repo_id=args.repo,
            repo_type="model",
        )
        meta_f = out_dir / "indigo_best_meta.json"
        if meta_f.exists():
            api.upload_file(
                path_or_fileobj=str(meta_f),
                path_in_repo="out/indigo_best_meta.json",
                repo_id=args.repo,
                repo_type="model",
            )
        print(f"[pipeline] checkpoint best ter-upload ke {args.repo}")

    print("[pipeline] selesai.")


if __name__ == "__main__":
    main()