Upload scripts/multi_quant_build_upload.py with huggingface_hub

scripts/multi_quant_build_upload.py

@@ -0,0 +1,480 @@
+#!/usr/bin/env python3
+import argparse
+import importlib
+import json
+import os
+import shlex
+import signal
+import subprocess
+import sys
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Optional
+
+# Must match your converter's accepted quant modes
+SUPPORTED_BY_CONVERTER = {"none", "dynamic_int8", "float16", "int8", "int4"}
+
+
+def run(
+    cmd,
+    cwd=None,
+    check=True,
+    env=None,
+    log_file: Optional[Path] = None,
+    timeout_sec: int = 0,
+):
+    print(f"[CMD] {' '.join(shlex.quote(c) for c in cmd)}", flush=True)
+
+    p = subprocess.Popen(
+        cmd,
+        cwd=cwd,
+        env=env,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.STDOUT,
+        text=True,
+        bufsize=1,
+        preexec_fn=os.setsid,  # new process group so we can kill children too
+    )
+
+    lines = []
+    timed_out = False
+
+    try:
+        assert p.stdout is not None
+        for line in p.stdout:
+            print(line, end="", flush=True)
+            lines.append(line)
+
+        if timeout_sec > 0:
+            rc = p.wait(timeout=timeout_sec)
+        else:
+            rc = p.wait()
+
+    except subprocess.TimeoutExpired:
+        timed_out = True
+        print(f"\n[!] Timeout after {timeout_sec}s. Terminating process group...", flush=True)
+        try:
+            os.killpg(os.getpgid(p.pid), signal.SIGTERM)
+        except Exception:
+            pass
+        try:
+            p.wait(timeout=8)
+        except Exception:
+            try:
+                os.killpg(os.getpgid(p.pid), signal.SIGKILL)
+            except Exception:
+                pass
+        rc = 124
+
+    except KeyboardInterrupt:
+        print("\n[!] KeyboardInterrupt received. Terminating process group...", flush=True)
+        try:
+            os.killpg(os.getpgid(p.pid), signal.SIGTERM)
+        except Exception:
+            pass
+        try:
+            p.wait(timeout=8)
+        except Exception:
+            try:
+                os.killpg(os.getpgid(p.pid), signal.SIGKILL)
+            except Exception:
+                pass
+        raise
+
+    out = "".join(lines)
+    if log_file:
+        log_file.parent.mkdir(parents=True, exist_ok=True)
+        log_file.write_text(out, encoding="utf-8")
+
+    if check and rc != 0:
+        msg = f"Command failed ({rc})"
+        if timed_out:
+            msg += " [timeout]"
+        msg += f": {' '.join(cmd)}"
+        raise RuntimeError(msg)
+
+    return rc, out, timed_out
+
+
+def ensure_hf_repo(repo_id: str, private: bool):
+    cmd = ["hf", "repo", "create", repo_id, "--type", "model"]
+    if private:
+        cmd.append("--private")
+    else:
+        cmd.append("--public")
+
+    rc, out, _ = run(cmd, check=False)
+    low = out.lower()
+    if rc != 0 and "already exists" not in low:
+        raise RuntimeError(f"Failed creating repo {repo_id}")
+
+
+def hf_upload(repo_id: str, local_path: Path, path_in_repo: str):
+    cmd = [
+        "hf",
+        "upload",
+        repo_id,
+        str(local_path),
+        path_in_repo,
+        "--repo-type",
+        "model",
+    ]
+    run(cmd, check=True)
+
+
+def normalize_quant(q: str) -> str:
+    q = q.strip().lower()
+    aliases = {
+        "fp16": "float16",
+        "f16": "float16",
+        "i8": "int8",
+        "i4": "int4",
+        "q8": "int8",
+        "q4": "int4",
+        "fp32": "none",
+        "off": "none",
+        "no": "none",
+    }
+    return aliases.get(q, q)
+
+
+def detect_native_quant_possible(model_dir: Path):
+    """
+    Heuristic for your environment:
+    - For gemma3, native quant is likely usable only if build_model_4b exists.
+    """
+    cfg = model_dir / "config.json"
+    if not cfg.exists():
+        return False, "config.json missing"
+
+    try:
+        j = json.loads(cfg.read_text())
+        model_type = (j.get("model_type") or "").lower()
+    except Exception as e:
+        return False, f"config parse failed: {e}"
+
+    if model_type == "gemma3":
+        try:
+            mod = importlib.import_module("litert_torch.generative.examples.gemma3.gemma3")
+            available = {n for n in dir(mod) if n.startswith("build_model")}
+            if "build_model_4b" in available:
+                return True, "gemma3 build_model_4b present"
+            return False, f"gemma3 4b builder missing; available={sorted(available)}"
+        except Exception as e:
+            return False, f"cannot import gemma3 converter module: {e}"
+
+    return True, f"model_type={model_type or 'unknown'}"
+
+
+def plan_quants(requested_quants, native_ok: bool):
+    """
+    Build plan for requested quantization modes.
+
+    Strategy 3 (post-TFLite quantization) gives int4/int8/dynamic_int8 genuinely
+    different output sizes, so they are each built independently.
+
+    The only deduplication is: if 'none' is requested alongside int4/int8/dynamic_int8,
+    the float32 TFLite produced by those builds also satisfies 'none', so we skip a
+    redundant bare 'none' build.
+    """
+    requested = []
+    unsupported = []
+
+    for q in requested_quants:
+        if q not in SUPPORTED_BY_CONVERTER:
+            unsupported.append(q)
+        else:
+            requested.append(q)
+
+    # All supported modes each get their own build — Strategy 3 makes them genuinely distinct
+    build_plan = requested[:]
+    alias_map = {q: q for q in requested}
+    mode = "native_quant_available" if native_ok else "strategy3_post_tflite"
+
+    return build_plan, alias_map, unsupported, mode
+
+
+def main():
+    ap = argparse.ArgumentParser(
+        description="Run multi-quant conversion+bundle and upload successful artifacts to HF."
+    )
+    ap.add_argument("--converter-script", default="/home/ubuntu/convert_translategemma_android.py")
+    ap.add_argument("--model-id", default="google/translategemma-4b-it")
+    ap.add_argument("--model-dir", default="/home/ubuntu/translategemma-4b-it")
+    ap.add_argument("--tflite-root", default="/home/ubuntu/tflite_output")
+    ap.add_argument("--output-dir", default="/home/ubuntu/output")
+    ap.add_argument("--log-dir", default="/home/ubuntu/logs")
+    ap.add_argument("--prefill", type=int, default=1024)
+    ap.add_argument("--kvcache", type=int, default=1024)
+    ap.add_argument("--timeout-sec", type=int, default=0, help="Per-quant timeout. 0 = no timeout.")
+
+    ap.add_argument(
+        "--quants",
+        default="int4,int8,fp8,fp16,dynamic_int8",
+        help="Comma-separated requested modes",
+    )
+    ap.add_argument(
+        "--repo-id",
+        default="barakplasma/translategemma-4b-it-android-task-quantized",
+    )
+
+    vis = ap.add_mutually_exclusive_group()
+    vis.add_argument("--private", action="store_true", default=True)
+    vis.add_argument("--public", action="store_true")
+
+    ap.add_argument("--no-upload", action="store_true")
+    args = ap.parse_args()
+
+    converter_script = Path(args.converter_script)
+    if not converter_script.exists():
+        print(f"[x] converter script not found: {converter_script}", file=sys.stderr)
+        sys.exit(1)
+
+    model_dir = Path(args.model_dir)
+    tflite_root = Path(args.tflite_root)
+    output_dir = Path(args.output_dir)
+    log_dir = Path(args.log_dir)
+
+    tflite_root.mkdir(parents=True, exist_ok=True)
+    output_dir.mkdir(parents=True, exist_ok=True)
+    log_dir.mkdir(parents=True, exist_ok=True)
+
+    # Parse and dedup quants
+    quant_list = [normalize_quant(x) for x in args.quants.split(",") if x.strip()]
+    seen = set()
+    quant_list = [q for q in quant_list if not (q in seen or seen.add(q))]
+
+    native_ok, native_reason = detect_native_quant_possible(model_dir)
+    print(f"[+] native quant capability: {native_ok} ({native_reason})")
+
+    build_plan, alias_map, unsupported, plan_mode = plan_quants(quant_list, native_ok=native_ok)
+    print(f"[+] plan mode: {plan_mode}")
+    print(f"[+] requested quants: {quant_list}")
+    print(f"[+] build plan: {build_plan}")
+    if unsupported:
+        print(f"[!] unsupported (skipped): {unsupported}")
+
+    built = {}
+    results = []
+
+    try:
+        for q in build_plan:
+            print(f"\n=== BUILD QUANT: {q} ===", flush=True)
+            q_tflite_dir = tflite_root / q
+            q_tflite_dir.mkdir(parents=True, exist_ok=True)
+
+            task_file = output_dir / f"translategemma-4b-it-{q}.task"
+            log_file = log_dir / f"convert_{q}.log"
+
+            # Reuse existing float32 TFLite if available — Strategy 3 will quantize it
+            base_tflite = tflite_root / "none" / "translategemma-4b-it-generic-none.tflite"
+            if q not in ("none", "float16") and base_tflite.exists():
+                print(f"[+] Reusing existing float32 TFLite for {q} (Strategy 3 will quantize)", flush=True)
+                cmd = [
+                    sys.executable,
+                    str(converter_script),
+                    "--bundle-only",
+                    "--existing-tflite",
+                    str(base_tflite),
+                    "--tflite-dir",
+                    str(q_tflite_dir),
+                    "--output-dir",
+                    str(output_dir),
+                    "--task-file",
+                    str(task_file),
+                    "--quantize",
+                    q,
+                    "--model-dir",
+                    str(model_dir),
+                ]
+            else:
+                cmd = [
+                    sys.executable,
+                    str(converter_script),
+                    "--model-id",
+                    args.model_id,
+                    "--model-dir",
+                    str(model_dir),
+                    "--tflite-dir",
+                    str(q_tflite_dir),
+                    "--output-dir",
+                    str(output_dir),
+                    "--task-file",
+                    str(task_file),
+                    "--quantize",
+                    q,
+                    "--prefill",
+                    str(args.prefill),
+                    "--kvcache",
+                    str(args.kvcache),
+                    "--allow-no-token",
+                ]
+
+            rc, _, timed_out = run(
+                cmd,
+                check=False,
+                log_file=log_file,
+                timeout_sec=args.timeout_sec,
+            )
+
+            tflites = sorted(q_tflite_dir.glob("*.tflite"))
+            tflite_file = tflites[-1] if tflites else None
+            ok = (rc == 0) and task_file.exists()
+
+            built[q] = {
+                "quant": q,
+                "ok": ok,
+                "rc": rc,
+                "timed_out": timed_out,
+                "task": str(task_file) if task_file.exists() else "",
+                "task_size": task_file.stat().st_size if task_file.exists() else 0,
+                "tflite": str(tflite_file) if tflite_file and tflite_file.exists() else "",
+                "log": str(log_file),
+            }
+
+    except KeyboardInterrupt:
+        print("\n[!] Stopped by user (Ctrl+C). Partial results will be saved.", flush=True)
+
+    # Expand to requested list
+    for q in quant_list:
+        if q in unsupported:
+            results.append(
+                {
+                    "quant": q,
+                    "ok": False,
+                    "rc": 2,
+                    "timed_out": False,
+                    "task": "",
+                    "task_size": 0,
+                    "tflite": "",
+                    "log": "",
+                    "status": "unsupported",
+                    "alias_of": "",
+                }
+            )
+            continue
+
+        bq = alias_map.get(q, q)
+        b = built.get(bq)
+        if not b:
+            results.append(
+                {
+                    "quant": q,
+                    "ok": False,
+                    "rc": 130,
+                    "timed_out": False,
+                    "task": "",
+                    "task_size": 0,
+                    "tflite": "",
+                    "log": "",
+                    "status": "not_built",
+                    "alias_of": bq,
+                }
+            )
+            continue
+
+        status = "built" if q == bq else f"aliased_to_{bq}"
+        results.append(
+            {
+                **b,
+                "quant": q,
+                "status": status,
+                "alias_of": bq if q != bq else "",
+            }
+        )
+
+    # Write summary JSON
+    summary = {
+        "timestamp_utc": datetime.now(timezone.utc).isoformat(),
+        "native_quant_capability": native_ok,
+        "native_quant_reason": native_reason,
+        "plan_mode": plan_mode,
+        "requested_quants": quant_list,
+        "build_plan": build_plan,
+        "unsupported_quants": unsupported,
+        "results": results,
+    }
+
+    summary_json = output_dir / "quantization_summary.json"
+    summary_json.write_text(json.dumps(summary, indent=2), encoding="utf-8")
+    print(f"\n[+] Wrote summary: {summary_json}")
+
+    # Write README
+    readme = output_dir / "README.md"
+    lines = []
+    lines.append("---")
+    lines.append("license: other")
+    lines.append("library_name: mediapipe")
+    lines.append("pipeline_tag: text-generation")
+    lines.append("---\n")
+    lines.append("# TranslateGemma 4B IT - Quantized Android Task Bundles\n")
+    lines.append(f"Generated: `{datetime.now(timezone.utc).isoformat()}`\n")
+    lines.append(f"- Native quant capability: `{native_ok}`")
+    lines.append(f"- Reason: `{native_reason}`")
+    lines.append(f"- Plan mode: `{plan_mode}`\n")
+    lines.append("| Requested quant | Status | Built from | Task file | Size (bytes) |")
+    lines.append("|---|---|---|---|---|")
+    for r in results:
+        if r.get("status") == "unsupported":
+            status = "⏭️ unsupported by converter"
+            built_from = "-"
+        elif str(r.get("status", "")).startswith("aliased_to_"):
+            status = "↪️ aliased"
+            built_from = f"`{r.get('alias_of','-')}`"
+        else:
+            if r.get("timed_out"):
+                status = "⏱️ timeout"
+            else:
+                status = "✅ success" if r.get("ok") else f"❌ failed (rc={r.get('rc')})"
+            built_from = "`self`"
+
+        task_name = Path(r["task"]).name if r.get("task") else "-"
+        lines.append(
+            f"| `{r['quant']}` | {status} | {built_from} | `{task_name}` | `{r.get('task_size',0)}` |"
+        )
+
+    lines.append("\n## Notes")
+    lines.append("- Aliased entries are not rebuilt; they point to an equivalent built variant.")
+    lines.append("- `fp8` is often unsupported in current converter/runtime stacks.")
+    lines.append("- Verify on-device compatibility before public release.")
+
+    readme.write_text("\n".join(lines), encoding="utf-8")
+    print(f"[+] Wrote README: {readme}")
+
+    if args.no_upload:
+        print("[!] --no-upload set. Done.")
+        return
+
+    private = False if args.public else True
+    ensure_hf_repo(args.repo_id, private=private)
+
+    hf_upload(args.repo_id, readme, "README.md")
+    hf_upload(args.repo_id, summary_json, "quantization_summary.json")
+
+    uploaded = set()
+
+    # Upload only unique built outputs
+    for q, b in built.items():
+        if b.get("log"):
+            lp = Path(b["log"])
+            if lp.exists() and str(lp) not in uploaded:
+                hf_upload(args.repo_id, lp, f"logs/{lp.name}")
+                uploaded.add(str(lp))
+
+        if b.get("tflite"):
+            tp = Path(b["tflite"])
+            if tp.exists() and str(tp) not in uploaded:
+                hf_upload(args.repo_id, tp, f"artifacts/{q}/{tp.name}")
+                uploaded.add(str(tp))
+
+        if b.get("task"):
+            tk = Path(b["task"])
+            if tk.exists() and str(tk) not in uploaded:
+                hf_upload(args.repo_id, tk, f"artifacts/{q}/{tk.name}")
+                uploaded.add(str(tk))
+
+    print(f"\n[+] Upload complete: https://huggingface.co/{args.repo_id}")
+
+
+if __name__ == "__main__":
+    main()