better buttons

app.py

@@ -29,8 +29,10 @@ NAMESPACE = "dataset-rewriter"
 URL = "https://huggingface.co/spaces/dataset-rewriter/dataset-rewriter"
 
 NUM_ROWS_PREVIEW = 3
+PARTIAL_SUFFIX = {10: "-10", 100: "-100", 1000: "-1k", 10_000: "-10k", 100_000: "-100k", 1000_000: "-1M"}
 MAX_NUM_ROWS_TO_REWRITE = int(os.environ.get("MAX_NUM_ROWS_TO_REWRITE") or 1000)
-PARTIAL_SUFFIX = "-1k"
+assert MAX_NUM_ROWS_TO_REWRITE in PARTIAL_SUFFIX, "allowed max num rows are 100, 1000, 10000, 100000 and 1000000"
+
 NUM_PARALLEL_CALLS = 10
 NUM_ROWS_PER_CALL = 5
 MAX_PROGRESS_UPDATES_PER_SECOND = 4
@@ -88,6 +90,12 @@ a {
 }
 """
 
+examples = [
+    ["fka/awesome-chatgpt-prompts", "make the prompt 6 words long maximum"],
+    ["lhoestq/CudyPokemonAdventures", "Pikachu as main character"],
+    ["infinite-dataset-hub/SmallTalkDialogues", "translate to proper French"],
+]
+
 with gr.Blocks(css=css) as demo:
     dataset_info_json = gr.JSON(visible=False)
     with gr.Row():
@@ -111,7 +119,7 @@ with gr.Blocks(css=css) as demo:
 
             gr.Markdown("### ReWrite")
             with gr.Group():
-                input_prompt = gr.Textbox(label="Enter the adjustment or transformation to apply to the dataset:")
+                input_prompt = gr.Textbox(label="Adjustment or transformation to apply to the dataset")
                 with gr.Accordion("(Advanced) Edit columns", open=False):
                     output_format_dataframe = gr.DataFrame(col_count=(2, "fixed"), headers=["column", "type"])
             rewrite_preview_button = gr.Button("Preview Results", variant="primary")
@@ -119,8 +127,9 @@ with gr.Blocks(css=css) as demo:
             gr.Markdown("#### Output")
             full_dataset_generation_label = gr.Label(visible=False, show_label=False)
             pretty_output_preview = gr.DataFrame(interactive=False)
-            full_dataset_generation_success_markdown = gr.Markdown("")
             pretty_full_dataset_generation_output = gr.DataFrame(interactive=False, visible=False)
+            full_dataset_generation_success_html = gr.HTML()
+            gr.Examples(examples, inputs=[dataset_search, input_prompt])
             gr.Markdown(f"_powered by [{model_id}](https://huggingface.co/{model_id})_")
         with gr.Column(scale=4, min_width="200px"):
             with gr.Accordion("Settings", open=False, elem_classes="settings"):
@@ -130,8 +139,9 @@ with gr.Blocks(css=css) as demo:
                 gr.Markdown("Save datasets as public or private datasets")
                 visibility_radio = gr.Radio(["public", "private"], value="public", container=False, interactive=False)
                 gr.Markdown("Maximum number of rows to ReWrite")
-                gr.Dropdown(choices=[MAX_NUM_ROWS_TO_REWRITE], value=MAX_NUM_ROWS_TO_REWRITE, interactive=False, container=False)
-                gr.Markdown(f"_[duplicate]({URL}?duplicate=true) this space to rewrite bigger datasets_")
+                max_num_rows_dropdown = gr.Dropdown(choices=[num_rows for num_rows in PARTIAL_SUFFIX if num_rows <= MAX_NUM_ROWS_TO_REWRITE], value=MAX_NUM_ROWS_TO_REWRITE, container=False)
+                gr.Markdown("Duplicate this space to ReWrite more rows")
+                gr.HTML(f'<a href="{URL}?duplicate=true" target="_blank"><img src="https://huggingface.co/datasets/huggingface/badges/resolve/main/duplicate-this-space-xl.svg" alt="Duplicate this Space"></a>')
 
 
     ############
@@ -187,10 +197,10 @@ with gr.Blocks(css=css) as demo:
         prompt = prompt[:1000] if prompt.strip() else ""
         messages = [{"role": "user", "content": REWRITE_DATASET_PREVIEW.format(
             dataset=dataset,
-            rows=json.dumps({"data": rows}),
+            rows=json.dumps({"data": rows}, ensure_ascii=False),
             prompt=prompt,
         )}]
-        response_format = {"type": "json", "value": {"properties": {"data": {"type": "array", "items": format}}, "required": ["data"]}}
+        response_format = {"type": "json", "value": {"properties": {"data": {"type": "array", "items": format, "minItems": len(rows), "maxItems": len(rows)}}, "required": ["data"]}}
         yield from ijson.items(StringIteratorIO(stream_reponse(messages, response_format=response_format)), "data.item", buf_size=4)
 
 
@@ -198,13 +208,17 @@ with gr.Blocks(css=css) as demo:
         prompt = prompt[:1000] if prompt.strip() else ""
         messages = [{"role": "user", "content": REWRITE_DATASET.format(
             dataset=dataset,
-            rows=json.dumps({"data": rows}),
+            rows=json.dumps({"data": rows}, ensure_ascii=False),
             prompt=prompt,
-            input_preview_rows=json.dumps({"data": input_preview_rows}),
-            output_preview_rows=json.dumps({"data": output_preview_rows}),
+            input_preview_rows=json.dumps({"data": input_preview_rows}, ensure_ascii=False),
+            output_preview_rows=json.dumps({"data": output_preview_rows}, ensure_ascii=False),
         )}]
-        response_format = {"type": "json", "value": {"properties": {"data": {"type": "array", "items": format}}, "required": ["data"]}}
-        yield from ijson.items(StringIteratorIO(stream_reponse(messages, response_format=response_format)), "data.item", buf_size=4)
+        response_format = {"type": "json", "value": {"properties": {"data": {"type": "array", "items": format, "minItems": len(rows), "maxItems": len(rows)}}, "required": ["data"]}}
+        try:
+            yield from ijson.items(StringIteratorIO(stream_reponse(messages, response_format=response_format)), "data.item", buf_size=4)
+        except ijson.IncompleteJSONError as e:
+            print(f"{type(e).__name__}: {e}")
+            print("Warning: Some rows were missing during ReWriting.")
 
 
     def find_new_name(dataset: str, prompt: str) -> str:
@@ -311,11 +325,11 @@ with gr.Blocks(css=css) as demo:
     def show_input_from_dataset_search(dataset: str) -> dict:
         return _show_input_preview(dataset, default_subset="default", default_split="train")
 
-    @subset_dropdown.change(inputs=[dataset_search, subset_dropdown], outputs=[pretty_input_preview, subset_dropdown, split_dropdown, output_format_dataframe, dataset_info_json])
+    @subset_dropdown.select(inputs=[dataset_search, subset_dropdown], outputs=[pretty_input_preview, subset_dropdown, split_dropdown, output_format_dataframe, dataset_info_json])
     def show_input_from_subset_dropdown(dataset: str, subset: str) -> dict:
         return _show_input_preview(dataset, default_subset=subset, default_split="train")
 
-    @split_dropdown.change(inputs=[dataset_search, subset_dropdown, split_dropdown], outputs=[pretty_input_preview, subset_dropdown, split_dropdown, output_format_dataframe, dataset_info_json])
+    @split_dropdown.select(inputs=[dataset_search, subset_dropdown, split_dropdown], outputs=[pretty_input_preview, subset_dropdown, split_dropdown, output_format_dataframe, dataset_info_json])
     def show_input_from_split_dropdown(dataset: str, subset: str, split: str) -> dict:
         return _show_input_preview(dataset, default_subset=subset, default_split=split)
     
@@ -344,14 +358,14 @@ with gr.Blocks(css=css) as demo:
         print(f"(preview) Done ReWriting {dataset} with instruction '{prompt}'")
 
 
-    @rewrite_full_dataset_button.click(inputs=[dataset_search, subset_dropdown, split_dropdown, pretty_input_preview, pretty_output_preview, input_prompt, output_format_dataframe, dataset_info_json, select_namespace_dropdown], outputs=[full_dataset_generation_label, full_dataset_generation_success_markdown, pretty_output_preview, pretty_full_dataset_generation_output])
-    def rewrite_full_dataset(dataset: str, subset: str, split: str, pretty_input_preview_df: pd.DataFrame, pretty_output_preview_df: pd.DataFrame, prompt: str, output_format_df: pd.DataFrame, dataset_info: dict[str, Any], namespace: str, oauth_token: Optional[gr.OAuthToken]) -> Iterator[pd.DataFrame]:
+    @rewrite_full_dataset_button.click(inputs=[dataset_search, subset_dropdown, split_dropdown, pretty_input_preview, pretty_output_preview, input_prompt, output_format_dataframe, dataset_info_json, select_namespace_dropdown, max_num_rows_dropdown], outputs=[full_dataset_generation_label, full_dataset_generation_success_html, pretty_output_preview, pretty_full_dataset_generation_output])
+    def rewrite_full_dataset(dataset: str, subset: str, split: str, pretty_input_preview_df: pd.DataFrame, pretty_output_preview_df: pd.DataFrame, prompt: str, output_format_df: pd.DataFrame, dataset_info: dict[str, Any], namespace: str, max_num_rows: int, oauth_token: Optional[gr.OAuthToken]) -> Iterator[pd.DataFrame]:
         input_preview_rows = [{k: json.loads(v) for k, v in row.items()} for row in pretty_input_preview_df.to_dict(orient="records")]
         output_preview_rows = [{k: json.loads(v) for k, v in row.items()} for row in pretty_output_preview_df.to_dict(orient="records")]
         format = output_format_df.to_dict(orient="records")
         format = {"properties": {x["column"]: json.loads(x["type"]) for x in format}, "required": [x["column"] for x in format]}
         num_examples = dataset_info["splits"][split]["num_examples"]
-        total = min(num_examples, MAX_NUM_ROWS_TO_REWRITE)
+        total = min(num_examples, max_num_rows)
         print(f"ReWriting {dataset} with instruction '{prompt}'")
         yield {full_dataset_generation_label: gr.Label({f"⚙️ ReWriting {dataset}": 0.}, visible=True)}
         yield {pretty_full_dataset_generation_output: empty_dataframe}
@@ -387,7 +401,7 @@ with gr.Blocks(css=css) as demo:
         print(f"Done ReWriting {dataset} with instruction '{prompt}'")
 
         output_rows = [{k: json.loads(row[k]) for k in output_format_df["column"]} for rows in parallel_output_rows for row in rows]
-        new_dataset = find_new_name(dataset + (PARTIAL_SUFFIX if num_examples > total else ""), prompt)
+        new_dataset = find_new_name(dataset + (PARTIAL_SUFFIX[max_num_rows] if num_examples > total else ""), prompt)
         repo_id = namespace + "/" + new_dataset
         yield {full_dataset_generation_label: gr.Label({f"✅ ReWriting {dataset}": len(output_rows) / total, f"⚙️ Saving to {repo_id}": 0.})}
         token = oauth_token.token if oauth_token else save_dataset_hf_token
@@ -396,7 +410,11 @@ with gr.Blocks(css=css) as demo:
         ds.push_to_hub(repo_id, config_name=subset, split=split, token=token)
         DatasetCard(DATASET_CARD_CONTENT.format(new_dataset=new_dataset, dataset=dataset, model_id=model_id, prompt=prompt, url=URL)).push_to_hub(repo_id=repo_id, repo_type="dataset", token=token)
         yield {full_dataset_generation_label: gr.Label({f"✅ ReWriting {dataset}": len(output_rows) / total, f"✅ Saving to {repo_id}": 1.})}
-        yield {full_dataset_generation_success_markdown: f"# Open the ReWriten dataset in a new tab: [{repo_id}](https://huggingface.co/datasets/{repo_id})"}
+        yield {full_dataset_generation_success_html: (
+            f'<a href="https://huggingface.co/datasets/{repo_id}" target="_blank">'
+            '<img src="https://huggingface.co/datasets/huggingface/badges/resolve/main/dataset-on-hf-xl.svg" alt="Dataset on HF", style="margin-right: auto; margin-left: auto; max-width: fit-content;">'
+            '</a>'
+        )}
         print(f"Saved {repo_id}")