| """Transformers Playground: find a model, grab its doc snippets, run them on HF Jobs.""" |
| import json |
| import os |
| import re |
| import time |
|
|
| import gradio as gr |
| from huggingface_hub import ( |
| HfApi, cancel_job, fetch_job_logs, hf_hub_download, inspect_job, run_job, |
| ) |
| from huggingface_hub.utils import GatedRepoError, HfHubHTTPError, RepositoryNotFoundError |
|
|
| from indexer import build_index |
|
|
| INDEX_REPO = os.getenv("INDEX_REPO", "") |
| UV_IMAGE = "ghcr.io/astral-sh/uv:python3.12-bookworm" |
| api = HfApi() |
|
|
| |
|
|
| def load_index(): |
| if INDEX_REPO: |
| try: |
| path = hf_hub_download(INDEX_REPO, "snippets.json", repo_type="dataset") |
| with open(path) as f: |
| return json.load(f) |
| except Exception as e: |
| print(f"Could not load {INDEX_REPO}: {e}") |
| if os.path.exists("snippets.json"): |
| with open("snippets.json") as f: |
| return json.load(f) |
| print("Building snippet index from transformers@main ...") |
| index = build_index("main") |
| with open("snippets.json", "w") as f: |
| json.dump(index, f) |
| return index |
|
|
|
|
| INDEX = load_index() |
| _norm = lambda s: re.sub(r"[^a-z0-9]", "", (s or "").lower()) |
| DOC_BY_NORM = {_norm(k): k for k in INDEX["docs"]} |
|
|
| |
|
|
| def _gb(s): |
| m = re.search(r"([\d.]+)", s or "") |
| return float(m.group(1)) if m else 0.0 |
|
|
|
|
| def load_hardware(): |
| try: |
| hw = api.list_jobs_hardware() |
| return [{ |
| "name": h.name, |
| "ram": _gb(h.ram), |
| "vram": _gb(h.accelerator.vram) if h.accelerator else 0.0, |
| "gpus": int(h.accelerator.quantity) if h.accelerator else 0, |
| "gpu": h.accelerator.model if h.accelerator else None, |
| "usd_hr": h.unit_cost_usd * 60, |
| } for h in hw] |
| except Exception as e: |
| print(f"Hardware list unavailable ({e}); using a minimal fallback.") |
| return [ |
| {"name": "cpu-basic", "ram": 16, "vram": 0, "gpus": 0, "gpu": None, "usd_hr": 0.01}, |
| {"name": "t4-small", "ram": 15, "vram": 16, "gpus": 1, "gpu": "T4", "usd_hr": 0.40}, |
| {"name": "l4x1", "ram": 30, "vram": 24, "gpus": 1, "gpu": "L4", "usd_hr": 0.80}, |
| {"name": "a100-large", "ram": 142, "vram": 80, "gpus": 1, "gpu": "A100", "usd_hr": 2.50}, |
| ] |
|
|
|
|
| HARDWARE = sorted(load_hardware(), key=lambda h: h["usd_hr"]) |
| HW_BY_NAME = {h["name"]: h for h in HARDWARE} |
| BYTES = {"F64": 8, "F32": 4, "BF16": 2, "F16": 2, "F8_E4M3": 1, "F8_E5M2": 1, "I8": 1, "U8": 1, "I32": 4} |
|
|
|
|
| def estimate_gb(meta): |
| if not meta.get("params"): |
| return None |
| return meta["params"] * meta["bytes"] / 1e9 * 1.2 + 1.0 |
|
|
|
|
| def recommend(meta, needs_gpu, multi_gpu_ok): |
| need = estimate_gb(meta) |
| if need is None: |
| return "t4-small", "Model size unknown (no safetensors metadata), starting with a T4." |
| if not needs_gpu and meta["params"] <= 1.5e9: |
| for h in HARDWARE: |
| if h["gpus"] == 0 and h["ram"] >= need * 1.5: |
| return h["name"], f"~{need:.1f} GB needed; small enough for CPU." |
| for h in HARDWARE: |
| if h["gpus"] and h["vram"] >= need and h["ram"] >= need and (h["gpus"] == 1 or multi_gpu_ok): |
| return h["name"], f"~{need:.1f} GB needed; cheapest GPU that fits." |
| return HARDWARE[-1]["name"], f"~{need:.1f} GB needed; larger than any single flavor, expect OOM." |
|
|
|
|
| def hw_choices(meta): |
| need = estimate_gb(meta) or 0 |
| out = [] |
| for h in HARDWARE: |
| spec = f"{h['gpus']}x {h['gpu']}, {h['vram']:.0f} GB" if h["gpus"] else f"CPU, {h['ram']:.0f} GB RAM" |
| small = " (too small)" if (h["gpus"] and h["vram"] < need) else "" |
| out.append((f"{h['name']} [{spec}] ${h['usd_hr']:.2f}/h{small}", h["name"])) |
| return out |
|
|
| |
|
|
| SAMPLE_IMG = "http://images.cocodataset.org/val2017/000000039769.jpg" |
| SAMPLE_AUDIO = "https://huggingface.co/datasets/Narsil/asr_dummy/resolve/main/mlk.flac" |
| GENERIC_INPUTS = { |
| "text-generation": '"Plants create energy through a process known as", max_new_tokens=50', |
| "text-classification": '"I love how easy this library makes things!"', |
| "token-classification": '"My name is Sarah and I live in London."', |
| "fill-mask": 'f"Paris is the {pipe.tokenizer.mask_token} of France."', |
| "question-answering": 'question="Where do I live?", context="My name is Sarah and I live in London."', |
| "zero-shot-classification": '"I just bought a new laptop", candidate_labels=["tech", "cooking", "sports"]', |
| "summarization": '"The tower is 324 metres tall, about the same height as an 81-storey building, and the tallest structure in Paris."', |
| "translation": '"How old are you?"', |
| "feature-extraction": '"Transformers is great."', |
| "sentence-similarity": '"Transformers is great."', |
| "image-classification": f'"{SAMPLE_IMG}"', |
| "object-detection": f'"{SAMPLE_IMG}"', |
| "image-segmentation": f'"{SAMPLE_IMG}"', |
| "depth-estimation": f'"{SAMPLE_IMG}"', |
| "image-to-text": f'"{SAMPLE_IMG}"', |
| "zero-shot-image-classification": f'"{SAMPLE_IMG}", candidate_labels=["cats", "dogs", "birds"]', |
| "image-feature-extraction": f'"{SAMPLE_IMG}"', |
| "automatic-speech-recognition": f'"{SAMPLE_AUDIO}"', |
| "audio-classification": f'"{SAMPLE_AUDIO}"', |
| "image-text-to-text": ( |
| f'text=[{{"role": "user", "content": [{{"type": "image", "url": "{SAMPLE_IMG}"}}, ' |
| f'{{"type": "text", "text": "Describe this image."}}]}}], max_new_tokens=50' |
| ), |
| } |
| TASK_ALIAS = {"sentence-similarity": "feature-extraction"} |
|
|
|
|
| def generic_snippet(task): |
| if task not in GENERIC_INPUTS: |
| return None |
| code = ( |
| "from transformers import pipeline\n\n" |
| f'pipe = pipeline(task="{TASK_ALIAS.get(task, task)}", model="MODEL_ID", device_map="auto")\n' |
| f"print(pipe({GENERIC_INPUTS[task]}))\n" |
| ) |
| return {"label": "Generic pipeline (auto-generated)", "section": "", "language": "python", |
| "code": code, "checkpoint": "MODEL_ID", "needs_gpu": False} |
|
|
|
|
| def adapt(snippet, model_id): |
| code = snippet["code"] |
| if snippet.get("checkpoint"): |
| code = code.replace(snippet["checkpoint"], model_id) |
| return code |
|
|
|
|
| def snippet_task(code): |
| m = re.search(r'pipeline\(\s*(?:task\s*=\s*)?["\']([\w-]+)["\']', code) |
| return m.group(1) if m else None |
|
|
|
|
| def infer_deps(code, task): |
| deps = ["transformers", "torch", "accelerate"] |
| checks = { |
| "pillow": r"\bPIL\b|Image\.open|image", |
| "torchvision": r"torchvision|image|Image|vision|video", |
| "requests": r"\brequests\b", |
| "datasets": r"load_dataset|from datasets|import datasets", |
| "soundfile": r"audio|soundfile|\.flac|\.wav", |
| "librosa": r"librosa|audio|\.flac|\.wav", |
| "sentencepiece": r"sentencepiece|T5|Marian", |
| "num2words": r"SmolVLM|Idefics3", |
| } |
| blob = code + " " + (task or "") |
| deps += [d for d, pat in checks.items() if re.search(pat, blob)] |
| return ", ".join(deps) |
|
|
| |
|
|
| def search_models(query, task): |
| models = api.list_models(search=query or None, filter="transformers", |
| pipeline_tag=task or None, sort="downloads", limit=30) |
| ids = [m.id for m in models] |
| return gr.update(choices=ids, value=ids[0] if ids else None) |
|
|
|
|
| def _token(oauth_token): |
| return oauth_token.token if oauth_token else None |
|
|
|
|
| def select_model(model_id, oauth_token: gr.OAuthToken | None): |
| empty = ("", gr.update(choices=[], value=None), [], {}, "", "", gr.update(), "", "") |
| if not model_id: |
| return empty |
| try: |
| info = api.model_info(model_id, token=_token(oauth_token)) |
| except RepositoryNotFoundError: |
| return ("Model not found, or it is private and you are not signed in.",) + empty[1:] |
|
|
| cfg = info.config or {} |
| model_type = cfg.get("model_type") |
| arch = DOC_BY_NORM.get(_norm(model_type)) |
| params, dtype_bytes = None, 2 |
| if info.safetensors: |
| params = info.safetensors.total |
| main_dtype = max(info.safetensors.parameters, key=info.safetensors.parameters.get) |
| dtype_bytes = BYTES.get(main_dtype, 2) |
| meta = {"model_id": model_id, "params": params, "bytes": dtype_bytes, |
| "task": info.pipeline_tag, "gated": bool(info.gated)} |
|
|
| snippets = list(INDEX["docs"][arch]["snippets"]) if arch else [] |
| generic = generic_snippet(info.pipeline_tag) |
| if generic: |
| snippets.append(generic) |
| snippets = [dict(s, adapted=adapt(s, model_id)) for s in snippets] |
|
|
| size = f"{params / 1e9:.2f}B params" if params else "size unknown" |
| lines = [f"### [{model_id}](https://huggingface.co/{model_id})", |
| f"`{model_type or 'unknown arch'}` | {info.pipeline_tag or 'no task tag'} | {size}"] |
| if arch: |
| d = INDEX["docs"][arch] |
| lines.append(f"Snippets from [{d['title']} docs]({d['doc_url']}) " |
| f"at transformers [`{INDEX['sha'][:8]}`]({d['source_url']})") |
| else: |
| lines.append("No model_doc page matched this architecture.") |
| if info.gated: |
| lines.append("This model is gated. Accept its license on the model page before running.") |
| if not snippets: |
| lines.append("No runnable snippets found for this model.") |
|
|
| choices = [(f"{s['label']} ({s['language']})", i) for i, s in enumerate(snippets)] |
| first = 0 if snippets else None |
| return ("\n\n".join(lines), gr.update(choices=choices, value=first), snippets, meta, |
| *snippet_outputs(first, snippets, meta)) |
|
|
|
|
| def snippet_outputs(idx, snippets, meta): |
| if idx is None or not snippets: |
| return "", "", gr.update(choices=hw_choices(meta), value="cpu-basic"), "", "" |
| s = snippets[idx] |
| code = s["adapted"] |
| flavor, why = recommend(meta, s["needs_gpu"], "device_map" in code) |
| notes = [f"Recommended hardware: **{flavor}**. {why}"] |
| st = snippet_task(code) |
| if st and meta.get("task") and TASK_ALIAS.get(meta["task"], meta["task"]) != st: |
| notes.append(f"Heads up: this snippet uses the `{st}` task but the model is tagged " |
| f"`{meta['task']}`. It may need edits.") |
| if s["language"] != "python": |
| notes.append("Shell snippets are copy-only; run Python snippets on Jobs.") |
| return (code, s["code"], gr.update(choices=hw_choices(meta), value=flavor), |
| "\n\n".join(notes), infer_deps(code, meta.get("task"))) |
|
|
|
|
| def cost_note(flavor, timeout): |
| h = HW_BY_NAME.get(flavor) |
| if not h: |
| return "" |
| return f"{flavor}: ${h['usd_hr']:.2f}/h. Worst case at a {int(timeout)} min timeout: ${h['usd_hr'] * timeout / 60:.2f}" |
|
|
|
|
| def build_script(code, deps, tf_source): |
| deps = [d.strip() for d in deps.split(",") if d.strip()] |
| if tf_source == "main": |
| deps = [d for d in deps if d != "transformers"] |
| deps.append("transformers @ https://github.com/huggingface/transformers/archive/refs/heads/main.zip") |
| header = "# /// script\n# requires-python = \">=3.10\"\n# dependencies = [\n" |
| header += "".join(f'# "{d}",\n' for d in deps) + "# ]\n# ///\n\n" |
| return header + code |
|
|
|
|
| def run_snippet(code, idx, snippets, meta, flavor, deps, tf_source, timeout, |
| oauth_token: gr.OAuthToken | None): |
| token = _token(oauth_token) |
| if not token: |
| yield "Sign in with Hugging Face first. Jobs run on your account.", "", None |
| return |
| if idx is not None and snippets and snippets[idx]["language"] != "python": |
| yield "This is a shell snippet. Copy it to run locally.", "", None |
| return |
| try: |
| api.auth_check(meta["model_id"], token=token) |
| except GatedRepoError: |
| yield f"You don't have access to {meta['model_id']} yet. Accept its license on the model page.", "", None |
| return |
| except Exception: |
| pass |
|
|
| script = build_script(code, deps, tf_source) |
| try: |
| job = run_job( |
| image=UV_IMAGE, |
| command=["bash", "-c", 'printf "%s" "$PLAYGROUND_SCRIPT" > /tmp/playground.py && uv run /tmp/playground.py'], |
| env={"PLAYGROUND_SCRIPT": script, "PYTHONUNBUFFERED": "1"}, |
| secrets={"HF_TOKEN": token}, |
| flavor=flavor, |
| timeout=f"{int(timeout)}m", |
| labels={"app": "transformers-playground"}, |
| token=token, |
| ) |
| except HfHubHTTPError as e: |
| hint = " Jobs need a PRO, Team or Enterprise account with credits." if e.response is not None and e.response.status_code in (402, 403) else "" |
| yield f"Could not start the job: {e}.{hint}", "", None |
| return |
|
|
| status = f"Job started on **{flavor}**. [Open job page]({job.url})" |
| logs, last = "", 0.0 |
| yield status + "\n\nWaiting for logs...", logs, job.id |
| try: |
| for line in fetch_job_logs(job_id=job.id, follow=True, token=token): |
| logs += line.rstrip("\n") + "\n" |
| if time.time() - last > 0.5: |
| last = time.time() |
| yield status + "\n\nRunning...", logs, job.id |
| except Exception as e: |
| logs += f"\n[log stream interrupted: {e}]\n" |
| stage = inspect_job(job_id=job.id, token=token).status.stage |
| yield f"{status}\n\nFinished with status **{stage}**.", logs, job.id |
|
|
|
|
| def stop_job(job_id, oauth_token: gr.OAuthToken | None): |
| if not job_id or not oauth_token: |
| return "No running job to cancel." |
| cancel_job(job_id=job_id, token=oauth_token.token) |
| return f"Cancel requested for job `{job_id}`." |
|
|
| |
|
|
| TASKS = ["", "text-generation", "image-text-to-text", "text-classification", "token-classification", |
| "fill-mask", "question-answering", "summarization", "translation", "feature-extraction", |
| "automatic-speech-recognition", "audio-classification", "image-classification", |
| "object-detection", "image-segmentation", "depth-estimation", "zero-shot-classification"] |
|
|
| with gr.Blocks(title="Transformers Playground") as demo: |
| snippets_state, meta_state, job_state = gr.State([]), gr.State({}), gr.State(None) |
|
|
| with gr.Row(equal_height=True): |
| gr.Markdown("# Transformers Playground\nPick a model, copy the official usage snippet, " |
| "or run it on Hugging Face Jobs with your own account.") |
| gr.LoginButton(scale=0, min_width=240) |
|
|
| with gr.Row(): |
| query = gr.Textbox(label="Search models", placeholder="llama, whisper, Qwen/Qwen2.5-0.5B-Instruct ...", scale=3) |
| task = gr.Dropdown(TASKS, value="", label="Task", scale=2) |
| search_btn = gr.Button("Search", scale=1) |
| model = gr.Dropdown(label="Model", choices=[], allow_custom_value=True, |
| info="Pick from results, or paste any model ID.") |
| model_md = gr.Markdown() |
|
|
| with gr.Row(): |
| with gr.Column(scale=3): |
| snippet_pick = gr.Radio(label="Snippet", choices=[]) |
| code = gr.Code(language="python", interactive=True, lines=18, |
| label="Snippet, adapted to your model (editable)", buttons=["copy", "download"]) |
| with gr.Accordion("Original snippet from the docs", open=False): |
| original = gr.Code(language="python", interactive=False, buttons=["copy"]) |
| with gr.Column(scale=2): |
| notes = gr.Markdown() |
| flavor = gr.Dropdown(label="Hardware", choices=hw_choices({}), value="cpu-basic") |
| timeout = gr.Slider(5, 120, value=30, step=5, label="Timeout (minutes)") |
| cost = gr.Markdown(cost_note("cpu-basic", 30)) |
| deps = gr.Textbox(label="Dependencies", info="Comma-separated pip specs added to the job.") |
| tf_source = gr.Radio(["latest release", "main"], value="latest release", |
| label="transformers version", |
| info="Snippets come from the main docs; pick main if a newer API fails.") |
| with gr.Row(): |
| run_btn = gr.Button("Run on HF Jobs", variant="primary") |
| stop_btn = gr.Button("Cancel job", variant="stop") |
| status = gr.Markdown() |
| logs = gr.Code(language="shell", label="Job logs", lines=14, interactive=False) |
|
|
| search_btn.click(search_models, [query, task], model) |
| query.submit(search_models, [query, task], model) |
| model.change(select_model, [model], |
| [model_md, snippet_pick, snippets_state, meta_state, code, original, flavor, notes, deps]) |
| snippet_pick.input(snippet_outputs, [snippet_pick, snippets_state, meta_state], |
| [code, original, flavor, notes, deps]) |
| flavor.change(cost_note, [flavor, timeout], cost) |
| timeout.change(cost_note, [flavor, timeout], cost) |
| run_btn.click(run_snippet, [code, snippet_pick, snippets_state, meta_state, flavor, deps, tf_source, timeout], |
| [status, logs, job_state]) |
| stop_btn.click(stop_job, [job_state], status) |
|
|
| if __name__ == "__main__": |
| demo.launch(theme=gr.themes.Soft(primary_hue="amber", secondary_hue="slate")) |
|
|
