"""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", "") # e.g. "your-org/transformers-snippets" UV_IMAGE = "ghcr.io/astral-sh/uv:python3.12-bookworm" api = HfApi() # ---------------------------------------------------------------- index ---- 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: # fall through to a local build 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"]} # ------------------------------------------------------------- hardware ---- 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 # weights + KV/activation headroom 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 # ------------------------------------------------------------- snippets ---- 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) # ------------------------------------------------------------ callbacks ---- 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 # let the job surface any other access problem 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}`." # ------------------------------------------------------------------- UI ---- 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"))