add save dataset

app.py

@@ -1,4 +1,5 @@
 import io
+import os
 import re
 import time
 from itertools import islice
@@ -10,16 +11,18 @@ from typing import Callable, Iterable, Iterator, Optional, TypeVar
 import gradio as gr
 import pandas as pd
 import requests.exceptions
-from huggingface_hub import InferenceClient
+from huggingface_hub import InferenceClient, create_repo, whoami, DatasetCard
 
 
 model_id = "microsoft/Phi-3-mini-4k-instruct"
 client = InferenceClient(model_id)
+save_dataset_hf_token = os.environ.get("SAVE_DATASET_HF_TOKEN")
 
 MAX_TOTAL_NB_ITEMS = 100  # almost infinite, don't judge me (actually it's because gradio needs a fixed number of components)
 MAX_NB_ITEMS_PER_GENERATION_CALL = 10
 NUM_ROWS = 100
 NUM_VARIANTS = 10
+NAMESPACE = "infinite-dataset-hub"
 URL = "https://huggingface.co/spaces/infinite-dataset-hub/infinite-dataset-hub"
 
 GENERATE_DATASET_NAMES_FOR_SEARCH_QUERY = (
@@ -54,8 +57,6 @@ LONG_RARITIES = [
     "very nice but still plausible",
 ]
 
-landing_page_query = "various datasets on many different subjects and topics, from classification to language modeling, from science to sport to finance to news"
-
 landing_page_datasets_generated_text = """
 1. NewsEventsPredict (classification, media, trend)
 2. FinancialForecast (economy, stocks, regression)
@@ -71,6 +72,29 @@ landing_page_datasets_generated_text = """
 default_output = landing_page_datasets_generated_text.strip().split("\n")
 assert len(default_output) == MAX_NB_ITEMS_PER_GENERATION_CALL
 
+DATASET_CARD_CONTENT = """
+---
+license: mit
+tags:
+- infinite-dataset-hub
+- synthetic
+---
+
+{title}
+
+_Note: This is an AI-generated dataset so its content may be inaccurate or false_
+
+{content}
+
+**Source of the data:**
+
+The dataset was generated using the [Infinite Dataset Hub]({url}) and {model_id} using the query '{search_query}':
+
+- **Dataset Generation Page**: {dataset_url}
+- **Model**: https://huggingface.co/{model_id}
+- **More Datasets**: https://huggingface.co/datasets?other=infinite-dataset-hub
+"""
+
 css = """
 a {
     color: var(--body-text-color);
@@ -145,35 +169,27 @@ a {
     color: transparent;
     background-clip: text;
 }
+.settings {
+    background: transparent;
+}
+.settings button span {
+    color: var(--body-text-color-subdued);
+}
 """
 
 
 with gr.Blocks(css=css) as demo:
     generated_texts_state = gr.State((landing_page_datasets_generated_text,))
-    with gr.Row():
-        with gr.Column(scale=4, min_width=0):
-            pass
-        with gr.Column(scale=10):
-            gr.Markdown(
-                "# 🤗 Infinite Dataset Hub ♾️\n\n"
-                "An endless catalog of datasets, created just for you.\n\n"
-            )
-        with gr.Column(scale=4, min_width=0):
-            pass
     with gr.Column() as search_page:
         with gr.Row():
-            with gr.Column(scale=4, min_width=0):
-                pass
             with gr.Column(scale=10):
+                gr.Markdown(
+                    "# 🤗 Infinite Dataset Hub ♾️\n\n"
+                    "An endless catalog of datasets, created just for you.\n\n"
+                )
                 with gr.Row():
                     search_bar = gr.Textbox(max_lines=1, placeholder="Search datasets, get infinite results", show_label=False, container=False, scale=9)
                     search_button = gr.Button("🔍", variant="primary", scale=1)
-            with gr.Column(scale=4, min_width=0):
-                pass
-        with gr.Row():
-            with gr.Column(scale=4, min_width=0):
-                pass
-            with gr.Column(scale=10):
                 button_groups: list[gr.Group] = []
                 buttons: list[gr.Button] = []
                 for i in range(MAX_TOTAL_NB_ITEMS):
@@ -195,20 +211,28 @@ with gr.Blocks(css=css) as demo:
 
                 load_more_datasets = gr.Button("Load more datasets")  # TODO: dosable when reaching end of page
                 gr.Markdown(f"_powered by [{model_id}](https://huggingface.co/{model_id})_")
-            with gr.Column(scale=4, min_width=0):
-                pass
+            with gr.Column(scale=4, min_width="200px"):
+                with gr.Accordion("Settings", open=False, elem_classes="settings"):
+                    gr.Markdown("Save datasets to your account")
+                    gr.LoginButton()
+                    select_namespace_dropdown = gr.Dropdown(choices=[NAMESPACE], value=NAMESPACE, label="Select user or organization", visible=False)
+                    gr.Markdown("Save datasets as public or private datasets")
+                    visibility_radio = gr.Radio(["public", "private"], value="public", container=False, interactive=False)
     with gr.Column(visible=False) as dataset_page:
+        gr.Markdown(
+            "# 🤗 Infinite Dataset Hub ♾️\n\n"
+            "An endless catalog of datasets, created just for you.\n\n"
+        )
         dataset_title = gr.Markdown()
         gr.Markdown("_Note: This is an AI-generated dataset so its content may be inaccurate or false_")
         dataset_content = gr.Markdown()
         generate_full_dataset_button = gr.Button("Generate Full Dataset", variant="primary")
         dataset_dataframe = gr.DataFrame(visible=False, interactive=False, wrap=True)
         save_dataset_button = gr.Button("💾 Save Dataset", variant="primary", visible=False)
+        open_dataset_message = gr.Markdown("", visible=False)
         dataset_share_button = gr.Button("Share Dataset URL")
         dataset_share_textbox = gr.Textbox(visible=False, show_copy_button=True, label="Copy this URL:", interactive=False, show_label=True)
         back_button = gr.Button("< Back", size="sm")
-    
-    app_state = gr.State({})
 
     ###################################
     #
@@ -254,7 +278,7 @@ with gr.Blocks(css=css) as demo:
 
     def gen_datasets_line_by_line(search_query: str, generated_texts: tuple[str] = ()) -> Iterator[str]:
         search_query = search_query or ""
-        search_query = search_query[:1000] if search_query.strip() else landing_page_query
+        search_query = search_query[:1000] if search_query.strip() else ""
         generated_text = ""
         current_line = ""
         for token in stream_reponse(
@@ -273,7 +297,7 @@ with gr.Blocks(css=css) as demo:
 
     def gen_dataset_content(search_query: str, dataset_name: str, tags: str) -> Iterator[str]:
         search_query = search_query or ""
-        search_query = search_query[:1000] if search_query.strip() else landing_page_query
+        search_query = search_query[:1000] if search_query.strip() else ""
         generated_text = ""
         for token in stream_reponse(GENERATE_DATASET_CONTENT_FOR_SEARCH_QUERY_AND_NAME_AND_TAGS.format(
             search_query=search_query,
@@ -418,7 +442,7 @@ with gr.Blocks(css=css) as demo:
 
 
     def _search_datasets(search_query):
-        yield {generated_texts_state: [], app_state: {"search_query": search_query}}
+        yield {generated_texts_state: []}
         yield {
             button_group: gr.Group(elem_classes="buttonsGroup insivibleButtonGroup")
             for button_group in button_groups[MAX_NB_ITEMS_PER_GENERATION_CALL:]
@@ -453,12 +477,12 @@ with gr.Blocks(css=css) as demo:
                 current_item_idx += 1
 
 
-    @search_button.click(inputs=search_bar, outputs=button_groups + buttons + [generated_texts_state, app_state])
+    @search_button.click(inputs=search_bar, outputs=button_groups + buttons + [generated_texts_state])
     def search_dataset_from_search_button(search_query):
         yield from _search_datasets(search_query)
 
 
-    @search_bar.submit(inputs=search_bar, outputs=button_groups + buttons + [generated_texts_state, app_state])
+    @search_bar.submit(inputs=search_bar, outputs=button_groups + buttons + [generated_texts_state])
     def search_dataset_from_search_bar(search_query):
         yield from _search_datasets(search_query)
 
@@ -497,20 +521,16 @@ with gr.Blocks(css=css) as demo:
             dataset_title: f"# {dataset_name}\n\n tags: {tags}",
             dataset_share_textbox: gr.Textbox(visible=False),
             dataset_dataframe: gr.DataFrame(visible=False),
-            generate_full_dataset_button: gr.Button(visible=True),
+            generate_full_dataset_button: gr.Button(interactive=True),
             save_dataset_button: gr.Button(visible=False),
-            app_state: {
-                "search_query": search_query,
-                "dataset_name": dataset_name,
-                "tags": tags
-            }
+            open_dataset_message: gr.Markdown(visible=False)
         }
         for generated_text in gen_dataset_content(search_query=search_query, dataset_name=dataset_name, tags=tags):
             yield {dataset_content: generated_text}
 
 
     show_dataset_inputs = [search_bar, *buttons]
-    show_dataset_outputs = [app_state, search_page, dataset_page, dataset_title, dataset_content, generate_full_dataset_button, dataset_dataframe, save_dataset_button, dataset_share_textbox]
+    show_dataset_outputs = [search_page, dataset_page, dataset_title, dataset_content, generate_full_dataset_button, dataset_dataframe, save_dataset_button, open_dataset_message, dataset_share_textbox]
     scroll_to_top_js = """
     function (...args) {
         console.log(args);
@@ -537,8 +557,8 @@ with gr.Blocks(css=css) as demo:
         return gr.Column(visible=True), gr.Column(visible=False)
 
 
-    @generate_full_dataset_button.click(inputs=[dataset_title, dataset_content, search_bar], outputs=[dataset_dataframe, generate_full_dataset_button, save_dataset_button])
-    def generate_full_dataset(title, content, search_query):
+    @generate_full_dataset_button.click(inputs=[dataset_title, dataset_content, search_bar, select_namespace_dropdown, visibility_radio], outputs=[dataset_dataframe, generate_full_dataset_button, save_dataset_button])
+    def generate_full_dataset(title, content, search_query, namespace, visability):
         dataset_name, tags = title.strip("# ").split("\ntags:", 1)
         dataset_name, tags = dataset_name.strip(), tags.strip()
         csv_header, preview_df = parse_preview_df(content)
@@ -556,8 +576,8 @@ with gr.Blocks(css=css) as demo:
         output[:len(preview_df)] = [{"idx": i, **x} for i, x in enumerate(preview_df.to_dict(orient="records"))]
         yield {
             dataset_dataframe: gr.DataFrame(pd.DataFrame([{"idx": i, **x} for i, x in enumerate(output) if x]), visible=True),
-            generate_full_dataset_button: gr.Button(visible=False),
-            save_dataset_button: gr.Button(visible=True, interactive=False)
+            generate_full_dataset_button: gr.Button(interactive=False),
+            save_dataset_button: gr.Button(f"💾 Save Dataset {namespace}/{dataset_name}" + (" (private)" if visability != "public" else ""), visible=True, interactive=False)
         }
         kwargs_iterable = [
             {
@@ -573,24 +593,48 @@ with gr.Blocks(css=css) as demo:
         ]
         for _ in iflatmap_unordered(generate_partial_dataset, kwargs_iterable=kwargs_iterable):
             yield {dataset_dataframe: pd.DataFrame([{"idx": i, **{column_name: x.get(column_name) for column_name in columns}} for i, x in enumerate(output) if x])}
-        yield {save_dataset_button: gr.Button(visible=True, interactive=True)}
+        yield {save_dataset_button: gr.Button(interactive=True)}
         print(f"Generated {dataset_name}!")
 
 
-    @save_dataset_button.click(inputs=[dataset_title, dataset_content, search_bar, dataset_dataframe])
-    def save_dataset(title, content, search_query, df):
-        raise gr.Error("Not implemented yet sorry ! Request your dataset to be saved in the Discussion tab (provide the dataset URL)")
-
-
-    @dataset_share_button.click(inputs=[app_state], outputs=[dataset_share_textbox])
-    def show_dataset_url(state):
+    @save_dataset_button.click(inputs=[dataset_title, dataset_content, search_bar, dataset_dataframe, select_namespace_dropdown, visibility_radio], outputs=[save_dataset_button, open_dataset_message])
+    def save_dataset(title: str, content: str, search_query: str, df: pd.DataFrame, namespace: str, visability: str, oauth_token: Optional[gr.OAuthToken]):
+        dataset_name, tags = title.strip("# ").split("\ntags:", 1)
+        dataset_name, tags = dataset_name.strip(), tags.strip()
+        token = oauth_token.token if oauth_token else save_dataset_hf_token
+        repo_id = f"{namespace}/{dataset_name}"
+        dataset_url = f"{URL}?q={search_query.replace(' ', '+')}&dataset={dataset_name.replace(' ', '+')}&tags={tags.replace(' ', '+')}"
+        gr.Info("Saving dataset...")
+        yield {save_dataset_button: gr.Button(interactive=False)}
+        create_repo(repo_id=repo_id, repo_type="dataset", private=visability!="public", exist_ok=True, token=token)
+        df.to_csv(f"hf://datasets/{repo_id}/data.csv", storage_options={"token": token}, index=False)
+        DatasetCard(DATASET_CARD_CONTENT.format(title=title, content=content, url=URL, dataset_url=dataset_url, model_id=model_id, search_query=search_query)).push_to_hub(repo_id=repo_id, repo_type="dataset", token=token)
+        gr.Info(f"✅ Dataset saved at {repo_id}")
+        additional_message = "PS: You can also save datasets under your account in the Settings ;)"
+        yield {open_dataset_message: gr.Markdown(f"# 🎉 Yay ! Your dataset has been saved to [{repo_id}](https://huggingface.co/datasets/{repo_id}) !\n\nDataset link: [https://huggingface.co/datasets/{repo_id}](https://huggingface.co/datasets/{repo_id})\n\n{additional_message}", visible=True)}
+
+
+    @dataset_share_button.click(inputs=[dataset_title, search_bar], outputs=[dataset_share_textbox])
+    def show_dataset_url(title, search_query):
+        dataset_name, tags = title.strip("# ").split("\ntags:", 1)
+        dataset_name, tags = dataset_name.strip(), tags.strip()
         return gr.Textbox(
-            f"{URL}?q={state['search_query'].replace(' ', '+')}&dataset={state['dataset_name'].replace(' ', '+')}&tags={state['tags'].replace(' ', '+')}",
+            f"{URL}?q={search_query.replace(' ', '+')}&dataset={dataset_name.replace(' ', '+')}&tags={tags.replace(' ', '+')}",
             visible=True,
         )
 
-    @demo.load(outputs=[app_state, search_page, search_bar, dataset_page, dataset_title, dataset_content, dataset_share_textbox] + button_groups + buttons + [generated_texts_state])
-    def load_app(request: gr.Request):
+    @demo.load(outputs=show_dataset_outputs + button_groups + buttons + [generated_texts_state] + [select_namespace_dropdown, visibility_radio])
+    def load_app(request: gr.Request, oauth_token: Optional[gr.OAuthToken]):
+        if oauth_token:
+            user_info = whoami(oauth_token.token)
+            yield {
+                select_namespace_dropdown: gr.Dropdown(
+                        choices=[user_info["name"]] + [org_info["name"] for org_info in user_info["orgs"]],
+                        value=user_info["name"],
+                        visible=True,
+                    ),
+                visibility_radio: gr.Radio(interactive=True),
+                }
         query_params = dict(request.query_params)
         if "dataset" in query_params:
             yield from _show_dataset(