app.py

import os
import gradio as gr
from huggingface_hub import HfApi
from helpers import clean_model_id, HF_TOKEN, resolve_token

# gr.LoginButton triggers Gradio's OAuth wiring, which requires OAUTH_CLIENT_ID
# (set automatically by Spaces with `hf_oauth: true`). Skip it elsewhere.
OAUTH_AVAILABLE = bool(os.environ.get("OAUTH_CLIENT_ID"))
from checks_ci import check_ci
from checks_attention import check_attention_support
from checks_attention_jobs import check_attention_jobs, DEFAULT_ATTN_HARDWARE, HARDWARE_CHOICES
from checks_serve_jobs import check_serve_jobs, DEFAULT_SERVE_HARDWARE
from checks_reach import gather_reach, DERIVATIVE_KINDS, fmt_date


def _fmt(n) -> str:
    return f"{n:,}" if isinstance(n, int) else "—"


def _fmt_bytes(n: int | None) -> str:
    if not n:
        return "—"
    units = ["B", "KB", "MB", "GB", "TB"]
    for u in units:
        if n < 1024:
            return f"{n:.1f} {u}"
        n /= 1024
    return f"{n:.1f} PB"


def _general_table(model_id: str, token: str | None = None) -> str:
    """Static, technical metadata about the model.

    Includes architecture, model_type, pipeline tag, license, weight-file size,
    and last-modified date. Size comes from summing `.safetensors` file sizes
    on the repo.
    """
    api = HfApi(token=resolve_token(token))
    try:
        info = api.model_info(model_id, files_metadata=True)
    except Exception as e:
        return f"### 📋 General info\n❌ Could not fetch model info: `{e}`"

    config = getattr(info, "config", None) or {}
    card_data = getattr(info, "card_data", None)
    if hasattr(card_data, "to_dict"):
        card_data = card_data.to_dict()
    elif card_data is None:
        card_data = {}

    arch = (config.get("architectures") or ["—"])[0]
    model_type = config.get("model_type") or "—"
    pipeline_tag = getattr(info, "pipeline_tag", None) or "—"
    license_ = card_data.get("license") or "—"

    siblings = getattr(info, "siblings", None) or []
    safetensors_bytes = sum(
        (getattr(f, "size", 0) or 0) for f in siblings
        if (getattr(f, "rfilename", "") or "").endswith(".safetensors")
    )

    rows = [
        ("🏷️ Architecture", f"`{arch}`"),
        ("🔖 Model type", f"`{model_type}`"),
        ("🎯 Pipeline", f"`{pipeline_tag}`" if pipeline_tag != "—" else "—"),
        ("📜 License", f"`{license_}`" if license_ != "—" else "—"),
        ("💾 Safetensors size", _fmt_bytes(safetensors_bytes)),
        ("📅 Last modified", fmt_date(getattr(info, "last_modified", None))),
    ]
    lines = ["### 📋 General info", "", "| Property | Value |", "|---|---|"]
    lines += [f"| {label} | {val} |" for label, val in rows]
    return "\n".join(lines)


def _reach_table(data: dict) -> str:
    """Compact community-signals table — no accordion, no clutter."""
    deriv = data["derivatives"]
    spaces = data["spaces"]
    spaces_n = spaces.get("count")
    open_hf_i = data["issues_by_status"].get("open", 0)
    open_hf_p = data["prs_by_status"].get("open", 0)

    def _d(kind):
        d = deriv.get(kind, {})
        c = d.get("count")
        if c is None:
            return "—"
        return f"{c:,}" + ("+" if d.get("capped") else "")

    rows = [
        ("📥 Downloads (30d)", _fmt(data["downloads_30d"])),
        ("❤️ Likes", _fmt(data["likes"])),
        ("🚀 Spaces using model", _fmt(spaces_n) if spaces_n is not None else "—"),
        ("🧬 Finetunes", _d("finetune")),
        ("🧩 Adapters", _d("adapter")),
        ("📦 Quantized", _d("quantized")),
        ("🔀 Merges", _d("merge")),
        ("💬 Open HF issues / PRs", f"{open_hf_i} / {open_hf_p}"),
    ]
    lines = ["### 🌐 Community signals", "", "| Signal | Count |", "|---|---:|"]
    lines += [f"| {label} | {val} |" for label, val in rows]
    return "\n".join(lines)


def _hf_discussion_url(d) -> str | None:
    url = getattr(d, "url", None)
    if not url:
        return None
    return url if url.startswith("http") else f"https://huggingface.co{url}"


def _discussions_md(data: dict) -> str:
    """List of open issues + PRs on the HF Hub for the model."""
    issues = data.get("issues_open") or []
    prs = data.get("prs_open") or []
    lines: list[str] = []

    if issues:
        lines.append(f"**Open issues — {len(issues)} shown**")
        for d in issues:
            url = _hf_discussion_url(d)
            num = f"[#{d.num}]({url})" if url else f"#{d.num}"
            title = (d.title or "")[:100]
            lines.append(f"- {num} — `@{d.author or '?'}` — {title}")
    else:
        lines.append("_No open issues._")

    lines.append("")

    if prs:
        lines.append(f"**Open PRs — {len(prs)} shown**")
        for d in prs:
            url = _hf_discussion_url(d)
            num = f"[#{d.num}]({url})" if url else f"#{d.num}"
            title = (d.title or "")[:100]
            lines.append(f"- {num} — `@{d.author or '?'}` — {title}")
    else:
        lines.append("_No open PRs._")

    return "\n".join(lines)


def _token(oauth_token: gr.OAuthToken | None) -> str | None:
    return oauth_token.token if oauth_token else HF_TOKEN


# ── orchestrators ─────────────────────────────────────────────────────────────

def run_checks(model_id_raw: str, oauth_token: gr.OAuthToken | None = None):
    if not model_id_raw.strip():
        msg = "Please enter a model ID or URL."
        return msg, msg, msg, msg, ""
    model_id = clean_model_id(model_id_raw)
    tok = _token(oauth_token)
    general = _general_table(model_id, token=tok)
    attn = check_attention_support(model_id, token=tok)
    ci = check_ci(model_id, token=tok)
    try:
        reach_data = gather_reach(model_id)
        reach = _reach_table(reach_data)
        discussions = _discussions_md(reach_data)
    except Exception as e:
        reach = f"### 🌐 Community signals\n❌ Could not fetch: `{e}`"
        discussions = "_Could not fetch open issues / PRs._"
    return general, attn, ci, reach, discussions


def run_job_checks(
    model_id_raw: str,
    hw_eager: str,
    hw_sdpa: str,
    hw_flash: str,
    hw_flex: str,
    oauth_token: gr.OAuthToken | None = None,
):
    if not model_id_raw.strip():
        yield "Please enter a model ID or URL.", ""
        return
    hardware = {
        "eager": hw_eager,
        "sdpa": hw_sdpa,
        "flash_attention_2": hw_flash,
        "flex_attention": hw_flex,
    }
    yield from check_attention_jobs(
        clean_model_id(model_id_raw),
        token=_token(oauth_token),
        hardware=hardware,
    )


def run_serve_jobs(
    model_id_raw: str,
    serve_hw_eager: str,
    serve_hw_sdpa: str,
    serve_hw_flash: str,
    serve_hw_flex: str,
    oauth_token: gr.OAuthToken | None = None,
):
    """Spin one HF Job per (attn × continuous_batching) cell and stream results."""
    if not model_id_raw.strip():
        yield "Please enter a model ID or URL.", ""
        return
    hardware = {
        "eager": serve_hw_eager,
        "sdpa": serve_hw_sdpa,
        "flash_attention_2": serve_hw_flash,
        "flex_attention": serve_hw_flex,
    }
    yield from check_serve_jobs(
        clean_model_id(model_id_raw),
        token=_token(oauth_token),
        hardware=hardware,
    )


# ── UI ────────────────────────────────────────────────────────────────────────

EXAMPLES = [
    ["Qwen/Qwen2.5-7B-Instruct"],
    ["google/gemma-2-2b-it"],
    ["meta-llama/Llama-3.2-1B"],
    ["mistralai/Ministral-8B-Instruct-2410"],
]

with gr.Blocks(title="vibecheck") as demo:
    gr.Markdown(
        "# ✅ vibecheck\n"
        "Paste a Hugging Face checkpoint to check config/tokenizer/generation, "
        "CI test status, and attention implementation support."
    )

    with gr.Row():
        gr.LoginButton()

    with gr.Row():
        model_input = gr.Textbox(
            label="Model ID or URL",
            placeholder="Qwen/Qwen2.5-7B-Instruct  or  https://huggingface.co/...",
            scale=5,
        )
        run_btn = gr.Button("Run Checks", variant="primary", scale=1, min_width=130)

    # ── transformers Support — main section, all subsections stacked ────────
    gr.Markdown("## 🤗 transformers Support")
    general_md = gr.Markdown()
    attn_summary = gr.Markdown()
    ci_out = gr.Markdown()
    reach_table = gr.Markdown()

    with gr.Accordion("📬 Open issues / PRs on the HF Hub", open=False):
        discussions_md = gr.Markdown()

    run_btn.click(
        fn=run_checks,
        inputs=[model_input],
        outputs=[
            general_md,
            attn_summary,
            ci_out,
            reach_table,
            discussions_md,
        ],
    )

    gr.Markdown("---")

    with gr.Accordion("Job hardware", open=False):
        with gr.Row():
            hw_eager = gr.Dropdown(
                choices=HARDWARE_CHOICES,
                value=DEFAULT_ATTN_HARDWARE["eager"].value,
                label="eager",
            )
            hw_sdpa = gr.Dropdown(
                choices=HARDWARE_CHOICES,
                value=DEFAULT_ATTN_HARDWARE["sdpa"].value,
                label="sdpa",
            )
            hw_flash = gr.Dropdown(
                choices=HARDWARE_CHOICES,
                value=DEFAULT_ATTN_HARDWARE["flash_attention_2"].value,
                label="flash_attention_2",
            )
            hw_flex = gr.Dropdown(
                choices=HARDWARE_CHOICES,
                value=DEFAULT_ATTN_HARDWARE["flex_attention"].value,
                label="flex_attention",
            )

    with gr.Row():
        jobs_btn = gr.Button(
            "☁️ Test model generation (HF Jobs)",
            variant="secondary",
        )

    with gr.Row():
        with gr.Column():
            jobs_summary = gr.Markdown(label="HF Jobs Results")
            with gr.Accordion("Details", open=False):
                jobs_details = gr.Markdown()

    jobs_btn.click(
        fn=run_job_checks,
        inputs=[model_input, hw_eager, hw_sdpa, hw_flash, hw_flex],
        outputs=[jobs_summary, jobs_details],
    )

    gr.Markdown("---")

    # ── transformers serve (attn × continuous batching, via HF Jobs) ────────
    with gr.Accordion("Job hardware for `transformers serve`", open=False):
        with gr.Row():
            serve_hw_eager = gr.Dropdown(
                choices=HARDWARE_CHOICES,
                value=DEFAULT_SERVE_HARDWARE["eager"].value,
                label="eager",
            )
            serve_hw_sdpa = gr.Dropdown(
                choices=HARDWARE_CHOICES,
                value=DEFAULT_SERVE_HARDWARE["sdpa"].value,
                label="sdpa",
            )
            serve_hw_flash = gr.Dropdown(
                choices=HARDWARE_CHOICES,
                value=DEFAULT_SERVE_HARDWARE["flash_attention_2"].value,
                label="flash_attention_2",
            )
            serve_hw_flex = gr.Dropdown(
                choices=HARDWARE_CHOICES,
                value=DEFAULT_SERVE_HARDWARE["flex_attention"].value,
                label="flex_attention",
            )

    with gr.Row():
        serve_btn = gr.Button(
            "🚀 Test transformers serve (HF Jobs)",
            variant="secondary",
        )

    with gr.Row():
        with gr.Column():
            serve_summary = gr.Markdown(label="Serve Jobs Results")
            with gr.Accordion("Details", open=False):
                serve_details = gr.Markdown()

    serve_btn.click(
        fn=run_serve_jobs,
        inputs=[model_input, serve_hw_eager, serve_hw_sdpa, serve_hw_flash, serve_hw_flex],
        outputs=[serve_summary, serve_details],
    )

    gr.Examples(examples=EXAMPLES, inputs=[model_input])


if __name__ == "__main__":
    demo.launch(theme=gr.themes.Soft())