app.py

"""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"))