Spaces:
Running
Running
| """Build Your Own SafeTulu: pick a removal threshold and sources, see what gets removed, export.""" | |
| from __future__ import annotations | |
| import os | |
| import tempfile | |
| import time | |
| from pathlib import Path | |
| from urllib.parse import urlencode | |
| import gradio as gr | |
| import numpy as np | |
| import plotly.graph_objects as go | |
| import labs_theme as lt | |
| import safetulu_data as sd | |
| # ZeroGPU Spaces refuse to start unless at least one function is decorated with | |
| # @spaces.GPU. This app needs no GPU, so register a no-op that is never called; | |
| # everything else runs on the CPU side as usual. Elsewhere this is skipped. | |
| try: | |
| import spaces | |
| def _zerogpu_placeholder(): | |
| return None | |
| except ImportError: | |
| pass | |
| DF = sd.load_scores() | |
| SOURCES = DF["source"].astype(str).value_counts().index.tolist() # largest first | |
| SOURCE_N = DF["source"].astype(str).value_counts().to_dict() | |
| def short(name, n=40): | |
| return name if len(name) <= n else name[: n // 2 - 1] + "…" + name[-(n // 2 - 1):] | |
| SOURCE_CHOICES = [(f"{short(s)} ({SOURCE_N[s]:,})", s) for s in SOURCES] | |
| SCORE_MIN = float(np.floor(DF[sd.SCORE_COL].min() * 1000) / 1000) | |
| SCORE_MAX = float(np.ceil(DF[sd.SCORE_COL].max() * 1000) / 1000) | |
| OFFICIAL_PCT = 100.0 * DF["official_removed"].mean() | |
| SPACE_ID = os.environ.get("SPACE_ID") # set automatically on HF Spaces | |
| ON_SPACES = bool(SPACE_ID) or os.environ.get("SAFETULU_ENABLE_PUSH") == "1" | |
| # Colors (validated categorical slots 1 and 2); text stays neutral. | |
| KEPT, REMOVED = lt.KEPT, lt.REMOVED | |
| INK, GRID = lt.INK, lt.GRID | |
| # --------------------------------------------------------------------------- | |
| # Figures | |
| # --------------------------------------------------------------------------- | |
| def _layout(fig, height, **kw): | |
| fig.update_layout( | |
| height=height, | |
| margin=dict(l=8, r=8, t=8, b=8), | |
| paper_bgcolor="rgba(0,0,0,0)", | |
| plot_bgcolor="rgba(0,0,0,0)", | |
| font=dict(color=INK, size=12, family=f"{lt.FONT}, ui-sans-serif, sans-serif"), | |
| hoverlabel=dict(font_size=12, bgcolor=lt.NAVY_2, bordercolor=lt.ELECTRIC, | |
| font=dict(color=lt.TEXT, family=f"{lt.FONT}, sans-serif")), | |
| legend=dict(traceorder="normal", orientation="h", yanchor="bottom", y=1.02, xanchor="right", x=1), | |
| **kw, | |
| ) | |
| fig.update_xaxes(gridcolor=GRID, zeroline=False, linecolor=GRID) | |
| fig.update_yaxes(gridcolor=GRID, zeroline=False, linecolor=GRID) | |
| return fig | |
| def hist_figure(in_pool, removed, cutoff, log_y): | |
| edges, kept_c, rem_c = sd.histogram(DF, in_pool, removed) | |
| mids = (edges[:-1] + edges[1:]) / 2 | |
| width = edges[1] - edges[0] | |
| ranges = [f"{a:.3f} – {b:.3f}" for a, b in zip(edges[:-1], edges[1:])] | |
| fig = go.Figure() | |
| for name, counts, color in (("Kept", kept_c, KEPT), ("Removed", rem_c, REMOVED)): | |
| fig.add_bar( | |
| x=mids, y=counts, width=width * 0.92, name=name, marker_color=color, | |
| customdata=ranges, | |
| hovertemplate=f"<b>{name}</b><br>score %{{customdata}}<br>%{{y:,}} examples<extra></extra>", | |
| ) | |
| if cutoff is not None: | |
| fig.add_vline(x=cutoff, line=dict(color=lt.CUTOFF_LINE, width=1, dash="dot")) | |
| fig.add_annotation(x=cutoff, y=1, yref="paper", text=f"cutoff {cutoff:.3f}", showarrow=False, | |
| xanchor="left", yanchor="top", xshift=4, font=dict(color=lt.TEXT, size=11)) | |
| fig = _layout(fig, 340, barmode="stack", bargap=0) | |
| fig.update_layout(margin=dict(t=36)) | |
| fig.update_xaxes(title_text="score (higher = closer to SafeRLHF harmful data)") | |
| fig.update_yaxes(title_text="examples", type="log" if log_y else "linear") | |
| return fig | |
| def source_figure(table): | |
| t = table.sort_values("pct_removed", ascending=True) | |
| fig = go.Figure(go.Bar( | |
| y=[short(x) for x in t.source], x=t.pct_removed, orientation="h", marker_color=REMOVED, | |
| customdata=np.stack([t.removed, t.examples, t.mean_score], axis=1), | |
| hovertemplate="<b>%{y}</b><br>%{x:.1f}% removed<br>%{customdata[0]:,} of %{customdata[1]:,}" | |
| "<br>mean score %{customdata[2]:.3f}<extra></extra>", | |
| text=[f"{v:.1f}%" for v in t.pct_removed], textposition="outside", cliponaxis=False, | |
| textfont=dict(color=INK, size=11), | |
| )) | |
| fig = _layout(fig, max(220, 26 * len(t) + 60), showlegend=False) | |
| fig.update_xaxes(title_text="% of source removed", ticksuffix="%", range=[0, max(1.0, float(t.pct_removed.max())) * 1.15]) | |
| fig.update_yaxes(automargin=True) | |
| return fig | |
| # --------------------------------------------------------------------------- | |
| # Main update | |
| # --------------------------------------------------------------------------- | |
| def settings_from(mode_label, pct, cutoff, include, exempt) -> sd.Settings: | |
| mode = next(k for k, v in sd.MODES.items() if v == mode_label) | |
| inc = None if set(include or []) == set(SOURCES) else list(include or []) | |
| return sd.Settings(mode=mode, pct=float(pct), cutoff=float(cutoff), include=inc, exempt=list(exempt or [])) | |
| def stat_tiles(st, s): | |
| def tile(label, value, sub="", cls=""): | |
| return (f'<div class="tile {cls}"><div class="tile-label">{label}</div>' | |
| f'<div class="tile-value">{value}</div><div class="tile-sub">{sub}</div></div>') | |
| cut = "n/a" if st["score_cutoff"] is None else f"{st['score_cutoff']:.4f}" | |
| jac = st["jaccard_with_official"] | |
| agree = "n/a" if jac is None else f"{100 * jac:.0f}%" | |
| return '<div class="tiles">' + "".join([ | |
| tile("Kept", f"{st['kept']:,}", f"of {st['pool']:,} in the selected sources"), | |
| tile("Removed", f"{st['removed']:,}", f"{st['removed_pct_of_pool']:.1f}% of the selection", "accent"), | |
| tile("Lowest removed score", cut, "examples at or above this are removed" if s.mode != "per_source" | |
| else "varies by source in this mode"), | |
| tile("Agreement with released SafeTulu", agree, | |
| f"Jaccard overlap of removed sets ({st['overlap_with_official']:,} shared)"), | |
| ]) + "</div>" | |
| def update(mode_label, pct, cutoff, include, exempt, log_y): | |
| s = settings_from(mode_label, pct, cutoff, include, exempt) | |
| if not include: | |
| empty = '<div class="tiles"><div class="tile"><div class="tile-label">No sources selected</div></div></div>' | |
| return empty, go.Figure(), go.Figure(), None, share_link(s), sd.snippet(s) | |
| in_pool, removed = sd.select(DF, s) | |
| st = sd.summarize(DF, in_pool, removed) | |
| table = sd.per_source(DF, in_pool, removed) | |
| shown = table.rename(columns={"pct_removed": "% removed", "mean_score": "mean score"}).round( | |
| {"% removed": 1, "mean score": 4}) | |
| return (stat_tiles(st, s), hist_figure(in_pool, removed, st["score_cutoff"] if s.mode != "per_source" else None, log_y), | |
| source_figure(table), shown, share_link(s), sd.snippet(s)) | |
| def share_link(s: sd.Settings) -> str: | |
| qs = urlencode(s.to_query()) | |
| base = f"https://huggingface.co/spaces/{SPACE_ID}" if SPACE_ID else "" | |
| return f"{base}?{qs}" | |
| def mode_visibility(mode_label): | |
| is_cut = mode_label == sd.MODES["cutoff"] | |
| return gr.update(visible=not is_cut), gr.update(visible=is_cut) | |
| def on_load(request: gr.Request): | |
| q = dict(request.query_params) if request else {} | |
| s = sd.Settings.from_query(q, SOURCES) | |
| include = SOURCES if s.include is None else s.include | |
| is_cut = s.mode == "cutoff" | |
| controls_v = (sd.MODES[s.mode], s.pct, s.cutoff, include, s.exempt) | |
| return (controls_v[0], gr.update(value=s.pct, visible=not is_cut), gr.update(value=s.cutoff, visible=is_cut), | |
| include, s.exempt, gr.update(open=bool(s.exempt)), *update(*controls_v, False)) | |
| # --------------------------------------------------------------------------- | |
| # Export | |
| # --------------------------------------------------------------------------- | |
| def export_ids(mode_label, pct, cutoff, include, exempt): | |
| s = settings_from(mode_label, pct, cutoff, include, exempt) | |
| in_pool, removed = sd.select(DF, s) | |
| out = Path(tempfile.mkdtemp(prefix="safetulu_ids_")) | |
| return [str(p) for p in sd.write_id_files(DF, in_pool, removed, s, out)] | |
| def export_hub(mode_label, pct, cutoff, include, exempt, repo_name, private, with_scores, with_emb, | |
| oauth_token: gr.OAuthToken | None, profile: gr.OAuthProfile | None, | |
| progress=gr.Progress()): | |
| if oauth_token is None or profile is None: | |
| raise gr.Error("Sign in with Hugging Face first.") | |
| repo_name = (repo_name or "").strip().strip("/") | |
| if not repo_name: | |
| raise gr.Error("Enter a dataset name.") | |
| repo_id = repo_name if "/" in repo_name else f"{profile.username}/{repo_name}" | |
| s = settings_from(mode_label, pct, cutoff, include, exempt) | |
| in_pool, removed = sd.select(DF, s) | |
| st = sd.summarize(DF, in_pool, removed) | |
| keep_ids = DF[sd.ID_COL].to_numpy()[in_pool & ~removed] | |
| t0 = time.time() | |
| work = Path(tempfile.mkdtemp(prefix="safetulu_push_")) | |
| files = sd.materialize(keep_ids, work / "data", include_scores=with_scores, | |
| include_embeddings=with_emb, progress=progress) | |
| progress(0.95, desc="Uploading to the Hub") | |
| url = sd.push_to_hub(files, sd.dataset_card(s, st), repo_id, oauth_token.token, private=private) | |
| return f"Pushed {st['kept']:,} examples to [{repo_id}]({url}) in {time.time() - t0:.0f}s." | |
| # --------------------------------------------------------------------------- | |
| # UI | |
| # --------------------------------------------------------------------------- | |
| with gr.Blocks(title="Build Your Own SafeTulu") as demo: | |
| gr.HTML(lt.hero_html(sd.REPO_ID, len(DF), OFFICIAL_PCT)) | |
| with gr.Row(equal_height=False): | |
| with gr.Column(scale=1, min_width=300): | |
| mode = gr.Radio(list(sd.MODES.values()), value=sd.MODES["global"], label="How to remove") | |
| pct = gr.Slider(0, 50, value=20, step=0.5, label="Remove (%)", | |
| info="Share of the filterable examples to remove") | |
| cutoff = gr.Slider(SCORE_MIN, SCORE_MAX, value=round((SCORE_MIN + SCORE_MAX) / 2, 4), step=0.0005, | |
| label="Score cutoff", info="Remove examples with score ≥ this", visible=False) | |
| include = gr.CheckboxGroup(SOURCE_CHOICES, value=SOURCES, label="Sources to include in the output") | |
| with gr.Row(): | |
| all_btn = gr.Button("Select all", size="sm") | |
| none_btn = gr.Button("Clear", size="sm") | |
| with gr.Accordion("Never filter these sources", open=False) as exempt_acc: | |
| exempt = gr.CheckboxGroup(SOURCE_CHOICES, value=[], show_label=False, | |
| info="Kept in full; the removal % applies to the rest") | |
| log_y = gr.Checkbox(False, label="Log scale for the histogram") | |
| with gr.Column(scale=2): | |
| tiles = gr.HTML() | |
| gr.Markdown("Score distribution", elem_classes="section-title") | |
| hist = gr.Plot(show_label=False) | |
| gr.Markdown("Share removed, by source", elem_classes="section-title") | |
| src_plot = gr.Plot(show_label=False) | |
| with gr.Accordion("Per-source table", open=False): | |
| table = gr.Dataframe(interactive=False, wrap=True) | |
| with gr.Tab("Share"): | |
| link = gr.Textbox(label="Link to this exact selection", buttons=["copy"], interactive=False) | |
| with gr.Tab("Download IDs"): | |
| gr.Markdown("Parquet files of kept and removed `document_id`s (with source and score), plus the settings as JSON.") | |
| ids_btn = gr.Button("Generate files", variant="primary") | |
| ids_files = gr.File(label="Files", file_count="multiple") | |
| with gr.Tab("Python script"): | |
| gr.Markdown("Reproduces this selection from the Hub without the Space and builds a `datasets.Dataset` of the kept examples.") | |
| code = gr.Code(language="python", label="Reproduce", max_lines=80, wrap_lines=True) | |
| with gr.Tab("Push to Hub"): | |
| gr.Markdown("Writes the kept examples (`document_id`, `content_hash`, `messages`) to a dataset repo in your account, with a dataset card recording these settings. Takes a few minutes.") | |
| if ON_SPACES: | |
| gr.LoginButton() | |
| repo_name = gr.Textbox(label="Dataset name", placeholder="my-safetulu-10pct") | |
| with gr.Row(): | |
| private = gr.Checkbox(True, label="Private") | |
| with_scores = gr.Checkbox(False, label="Include score columns") | |
| with_emb = gr.Checkbox(False, label="Include gradient embeddings (adds ~7 GB)") | |
| push_btn = gr.Button("Push to Hub", variant="primary") | |
| push_status = gr.Markdown() | |
| else: | |
| gr.Markdown("*Sign-in only works on the hosted Space. Running locally, use the Python script tab and `ds.push_to_hub(...)`.*") | |
| controls = [mode, pct, cutoff, include, exempt] | |
| outputs = [tiles, hist, src_plot, table, link, code] | |
| for c in controls + [log_y]: | |
| c.input(update, controls + [log_y], outputs, trigger_mode="always_last", show_progress="hidden") | |
| pct.release(update, controls + [log_y], outputs, trigger_mode="always_last", show_progress="hidden") | |
| cutoff.release(update, controls + [log_y], outputs, trigger_mode="always_last", show_progress="hidden") | |
| mode.input(mode_visibility, mode, [pct, cutoff]) | |
| all_btn.click(lambda: SOURCES, None, include).then(update, controls + [log_y], outputs, show_progress="hidden") | |
| none_btn.click(lambda: [], None, include).then(update, controls + [log_y], outputs, show_progress="hidden") | |
| ids_btn.click(export_ids, controls, ids_files) | |
| if ON_SPACES: | |
| push_btn.click(export_hub, controls + [repo_name, private, with_scores, with_emb], push_status) | |
| demo.load(on_load, None, controls + [exempt_acc] + outputs) | |
| if __name__ == "__main__": | |
| demo.launch(css=lt.CSS, theme=lt.build_theme(), js=lt.THEME_JS) | |