| """Build the snippet index from transformers' docs/source/en/model_doc. |
| |
| Run locally: python indexer.py --out snippets.json |
| Push to the Hub: python indexer.py --push your-org/transformers-snippets |
| (Schedule it as a daily HF Job to keep the index fresh.) |
| """ |
| import argparse |
| import json |
| import os |
| import re |
| import shutil |
| import subprocess |
| import time |
|
|
| REPO = "huggingface/transformers" |
| DOC_DIR = "docs/source/en/model_doc" |
|
|
| FENCE = re.compile(r"```(py|python|bash|sh)\s*\n(.*?)```", re.S) |
| HFOPTIONS = re.compile(r'<hfoptions id="([^"]+)">(.*?)</hfoptions>', re.S) |
| HFOPTION = re.compile(r'<hfoption id="([^"]+)">(.*?)</hfoption>', re.S) |
| HEADING = re.compile(r"^(#{2,3})\s+(.+)$", re.M) |
| CHECKPOINT = re.compile( |
| r'(?:from_pretrained\(\s*|model\s*=\s*|--model\s+)["\']([\w.\-]+/[\w.\-]+)["\']?' |
| ) |
|
|
| |
| EXCLUDE = re.compile( |
| r"\bTrainer\b|\.train\(|TrainingArguments|SFTTrainer|trl\b|loss\.backward|optimizer\." |
| r"|BitsAndBytesConfig|TorchAoConfig|quantization_config|GPTQConfig|AwqConfig|HqqConfig" |
| ) |
| INFERENCE_HINT = re.compile(r"pipeline\(|\.generate\(|from_pretrained|transformers (chat|run|serve)") |
|
|
|
|
| def fetch_docs(ref="main", workdir="/tmp/transformers-docs"): |
| """Shallow, sparse clone of just model_doc. No GitHub API calls, so no rate limits.""" |
| shutil.rmtree(workdir, ignore_errors=True) |
| run = lambda *a, **kw: subprocess.run(a, check=True, capture_output=True, text=True, **kw) |
| run("git", "clone", "--quiet", "--depth", "1", "--filter=blob:none", "--sparse", |
| "--branch", ref, f"https://github.com/{REPO}", workdir) |
| run("git", "sparse-checkout", "set", DOC_DIR, cwd=workdir) |
| sha = run("git", "rev-parse", "HEAD", cwd=workdir).stdout.strip() |
| doc_root = os.path.join(workdir, DOC_DIR) |
| files = {} |
| for name in sorted(os.listdir(doc_root)): |
| if name.endswith(".md"): |
| with open(os.path.join(doc_root, name), encoding="utf-8") as f: |
| files[f"{DOC_DIR}/{name}"] = f.read() |
| return sha, files |
|
|
|
|
| def _section_for(text, pos): |
| """Nearest heading above a position, used as context for loose code blocks.""" |
| title = "" |
| for m in HEADING.finditer(text): |
| if m.start() > pos: |
| break |
| title = m.group(2).strip() |
| return title |
|
|
|
|
| def _clean(code): |
| code = code.strip("\n") |
| return re.sub(r"^\s*<!--.*?-->\s*$", "", code, flags=re.M).strip("\n") |
|
|
|
|
| def _snippet(code, lang, label, section): |
| code = _clean(code) |
| if not code or EXCLUDE.search(code): |
| return None |
| if not INFERENCE_HINT.search(code): |
| return None |
| lang = "python" if lang in ("py", "python") else "bash" |
| ckpt = CHECKPOINT.search(code) |
| return { |
| "label": label, |
| "section": section, |
| "language": lang, |
| "code": code, |
| "checkpoint": ckpt.group(1) if ckpt else None, |
| "needs_gpu": bool(re.search(r'cuda|device_map|attn_implementation="flash', code)), |
| } |
|
|
|
|
| def parse_doc(text): |
| snippets, covered = [], [] |
| |
| for group in HFOPTIONS.finditer(text): |
| covered.append(group.span()) |
| section = _section_for(text, group.start()) |
| for opt in HFOPTION.finditer(group.group(2)): |
| for fence in FENCE.finditer(opt.group(2)): |
| s = _snippet(fence.group(2), fence.group(1), opt.group(1), section) |
| if s: |
| snippets.append(s) |
| |
| for fence in FENCE.finditer(text): |
| if any(a <= fence.start() < b for a, b in covered): |
| continue |
| section = _section_for(text, fence.start()) |
| if not re.search(r"usage|example|inference|generat|quick|tips|notes", section, re.I): |
| continue |
| s = _snippet(fence.group(2), fence.group(1), section or "Example", section) |
| if s: |
| snippets.append(s) |
| |
| seen, out = set(), [] |
| for s in snippets: |
| if s["code"] not in seen: |
| seen.add(s["code"]) |
| out.append(s) |
| return out |
|
|
|
|
| def build_index(ref="main"): |
| sha, files = fetch_docs(ref) |
| docs = {} |
| for path, text in files.items(): |
| arch = os.path.basename(path)[:-3] |
| title = re.search(r"^#\s+(.+)$", text, re.M) |
| docs[arch] = { |
| "title": title.group(1).strip() if title else arch, |
| "doc_url": f"https://huggingface.co/docs/transformers/model_doc/{arch}", |
| "source_url": f"https://github.com/{REPO}/blob/{sha}/{path}", |
| "snippets": parse_doc(text), |
| } |
| return {"ref": ref, "sha": sha, "built_at": int(time.time()), "docs": docs} |
|
|
|
|
| if __name__ == "__main__": |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--ref", default="main") |
| ap.add_argument("--out", default="snippets.json") |
| ap.add_argument("--push", help="dataset repo id to upload snippets.json to") |
| args = ap.parse_args() |
|
|
| index = build_index(args.ref) |
| with open(args.out, "w") as f: |
| json.dump(index, f) |
| n = sum(len(d["snippets"]) for d in index["docs"].values()) |
| print(f"{len(index['docs'])} docs, {n} snippets @ {index['sha'][:8]}") |
|
|
| if args.push: |
| from huggingface_hub import HfApi |
| api = HfApi() |
| api.create_repo(args.push, repo_type="dataset", exist_ok=True) |
| api.upload_file(path_or_fileobj=args.out, path_in_repo="snippets.json", |
| repo_id=args.push, repo_type="dataset", |
| commit_message=f"Index transformers@{index['sha'][:8]}") |
| print(f"Pushed to https://huggingface.co/datasets/{args.push}") |
|
|
