File size: 2,779 Bytes
a9cfb03 | 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 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 | """`python -m arch_svg ...`
python -m arch_svg one --model llama --mode full --out out/
python -m arch_svg one --model gemma --mode diff --out out/
python -m arch_svg all --mode diff --out out/ --jobs 8
python -m arch_svg all --mode diff --out out/ --limit 20 # quick subset
"""
from __future__ import annotations
import argparse
import os
import sys
def _cmd_one(args) -> int:
from .discover import discover_models, model_type_for
from .gallery import render_one
model_type = None
for e in discover_models():
if e.name == args.model:
model_type = model_type_for(e)
break
svg, record = render_one(args.model, model_type or args.model, args.mode)
os.makedirs(args.out, exist_ok=True)
path = os.path.join(args.out, f"{args.model}.svg")
with open(path, "w", encoding="utf-8") as f:
f.write(svg)
print(f"wrote {path} (status={record['status']})")
if record.get("build_error"):
print(f" note: {record['build_error']}")
if args.mode == "diff" and record.get("is_modular") is not None:
print(
f" modular={record.get('is_modular')} parent={record.get('parent_model')} totals={record.get('diff_totals')}"
)
return 0 if record["status"] != "failed" else 1
def _cmd_all(args) -> int:
from .gallery import run_all
res = run_all(out=args.out, mode=args.mode, jobs=args.jobs, limit=args.limit)
s = res["summary"]
print(f"\ndone → {os.path.join(args.out, 'index.html')}")
print(f" {s}")
return 0
def _cmd_list(args) -> int:
from .discover import discover_models
for e in discover_models():
tag = "modular" if e.has_modular else "standalone"
print(f"{e.name:32s} {tag}")
return 0
def main(argv=None) -> int:
p = argparse.ArgumentParser(prog="arch_svg", description="transformers architecture SVG generator")
sub = p.add_subparsers(dest="cmd", required=True)
one = sub.add_parser("one", help="render a single model")
one.add_argument("--model", required=True)
one.add_argument("--mode", choices=["full", "diff"], default="full")
one.add_argument("--out", default="out")
one.set_defaults(func=_cmd_one)
alle = sub.add_parser("all", help="render the whole model zoo + index.html")
alle.add_argument("--mode", choices=["full", "diff", "both"], default="both")
alle.add_argument("--out", default="out")
alle.add_argument("--jobs", type=int, default=4)
alle.add_argument("--limit", type=int, default=None)
alle.set_defaults(func=_cmd_all)
lst = sub.add_parser("list", help="list discovered models")
lst.set_defaults(func=_cmd_list)
args = p.parse_args(argv)
return args.func(args)
if __name__ == "__main__":
sys.exit(main())
|