finalize save dataset

app.py

@@ -11,6 +11,7 @@ import ijson
 import pandas as pd
 import requests
 from datasets import Dataset, Features, Value, Sequence
+from datasets.fingerprint import Hasher
 from gradio_huggingfacehub_search import HuggingfaceHubSearch
 from huggingface_hub import InferenceClient
 
@@ -18,16 +19,19 @@ from utils import StringIteratorIO
 
 
 model_id = "meta-llama/Meta-Llama-3.1-8B-Instruct"
-client = InferenceClient(model_id, token=os.environ.get("HF_TOKEN_INFERENCE_API"))
+client = InferenceClient(model_id, token=os.environ.get("INFERENCE_API_HF_TOKEN"))
 
 save_dataset_hf_token = os.environ.get("SAVE_DATASET_HF_TOKEN")
 session = requests.Session()
 empty_dataframe = pd.DataFrame({"1": [], "2": [], "3": []})
 
-NAMESPACE = "lhoestq"
+NAMESPACE = "dataset-rewriter"
 
 NUM_ROWS_PREVIEW = 3
-MAX_NUM_ROWS_TO_REWRITE = 10
+MAX_NUM_ROWS_TO_REWRITE = 1000
+NUM_PARALLEL_CALLS = 10
+NUM_ROWS_PER_CALL = 10
+MAX_PROGRESS_UPDATES_PER_SECOND = 4
 REWRITE_DATASET_PREVIEW = (
     "A Machine Learning practitioner is looking for a dataset similar to '{dataset}' but slightly different. "
     "They want you to rewrite the dataset and apply this instruction, which can be about transforming, translating or filtering the rows: {prompt}."
@@ -77,7 +81,7 @@ with gr.Blocks(css=css) as demo:
                 subset_dropdown = gr.Dropdown(info="Subset", show_label=False, visible=False)
                 split_dropdown = gr.Dropdown(info="Split", show_label=False, visible=False)
 
-            gr.Markdown("### Input")
+            gr.Markdown("### Sample")
             pretty_input_preview = gr.DataFrame(interactive=False)
 
             gr.Markdown("### ReWrite")
@@ -86,10 +90,12 @@ with gr.Blocks(css=css) as demo:
                 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")
-            pretty_output_preview = gr.DataFrame(interactive=False)
             rewrite_full_dataset_button = gr.Button("ReWrite Full Dataset", interactive=False)
+            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)
         with gr.Column(scale=4, min_width="200px"):
             with gr.Accordion("Settings", open=False, elem_classes="settings"):
                 gr.Markdown("Save datasets to your account")
@@ -156,9 +162,9 @@ with gr.Blocks(css=css) as demo:
             prompt=prompt,
         )}]
         response_format = {"type": "json", "value": {"properties": {"data": {"type": "array", "items": format}}, "required": ["data"]}}
-        print("streaming preview")
+        print(f"Streaming preview of {dataset} with instruction '{prompt}'")
         yield from ijson.items(StringIteratorIO(stream_reponse(messages, response_format=response_format)), "data.item", buf_size=4)
-        print("done")
+        print(f"Done streaming preview of {dataset} with instruction '{prompt}'")
 
 
     def stream_rewrite_dataset_row_by_row(dataset: str, rows: list[dict[str, str]], prompt: str, format: str, input_preview_rows: list[dict[str, str]], output_preview_rows: list[dict[str, str]]) -> Iterator[dict[str, str]]:
@@ -171,18 +177,17 @@ with gr.Blocks(css=css) as demo:
             output_preview_rows=json.dumps({"data": output_preview_rows}),
         )}]
         response_format = {"type": "json", "value": {"properties": {"data": {"type": "array", "items": format}}, "required": ["data"]}}
-        print("streaming results")
         yield from ijson.items(StringIteratorIO(stream_reponse(messages, response_format=response_format)), "data.item", buf_size=4)
-        print("done")
 
 
     def find_new_name(dataset: str, prompt: str) -> str:
         messages = [{"role": "user", "content": FIND_NEW_NAME.format(prompt=prompt)}]
         out = "".join(stream_reponse(messages))
         if "should be" in out:
-            return dataset.split("/")[-1] + out.split("should be", 1)[1].replace(" ", "-").replace(".", "").replace(",", "")
+            out = dataset.split("/")[-1] + out.split("should be", 1)[1].replace(" ", "-").replace(".", "").replace(",", "")
         else:
-            return dataset.split("/")[-1] + prompt.replace(" ", "-")
+            out = dataset.split("/")[-1] + prompt.replace(" ", "-")
+        return out[:90] + "-" + Hasher.hash(prompt)[:4]
 
     def _write_generator_to_queue(queue: Queue, func: Callable[..., Iterable], kwargs: dict) -> None:
         for i, result in enumerate(func(**kwargs)):
@@ -293,7 +298,7 @@ with gr.Blocks(css=css) as demo:
         return {rewrite_full_dataset_button: gr.Button(interactive=False)}
 
 
-    @rewrite_preview_button.click(inputs=[dataset_search, pretty_input_preview, input_prompt, output_format_dataframe], outputs=[pretty_output_preview, rewrite_full_dataset_button, full_dataset_generation_label])
+    @rewrite_preview_button.click(inputs=[dataset_search, pretty_input_preview, input_prompt, output_format_dataframe], outputs=[pretty_output_preview, rewrite_full_dataset_button, full_dataset_generation_label, pretty_full_dataset_generation_output])
     def rewrite_preview(dataset: str, pretty_input_preview_df: pd.DataFrame, prompt: str, output_format_df: pd.DataFrame) -> Iterator[pd.DataFrame]:
         rows = [{k: json.loads(v) for k, v in row.items()} for row in pretty_input_preview_df.to_dict(orient="records")]
         format = output_format_df.to_dict(orient="records")
@@ -301,34 +306,68 @@ with gr.Blocks(css=css) as demo:
         output_rows = []
         print(f"ReWriting {dataset} preview with instruction '{prompt}'")
         yield {rewrite_full_dataset_button: gr.Button(interactive=False), full_dataset_generation_label: gr.Label(visible=False)}
+        yield {
+            pretty_output_preview: gr.DataFrame(visible=True),
+            pretty_full_dataset_generation_output: gr.DataFrame(visible=False),
+        }
         for row in stream_rewrite_dataset_preview_row_by_row(dataset=dataset, rows=rows, prompt=prompt, format=format):
             output_rows.append({k: json.dumps(row[k], ensure_ascii=False) for k in output_format_df["column"]})
             yield {pretty_output_preview: gr.DataFrame(pd.DataFrame(output_rows))}
         yield {rewrite_full_dataset_button: gr.Button(interactive=True)}
 
 
-    @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])
+    @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]:
         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]}
-        output_rows = []
         num_examples = dataset_info["splits"][split]["num_examples"]
         total = min(num_examples, MAX_NUM_ROWS_TO_REWRITE)
         print(f"ReWriting {dataset} (full dataset) with instruction '{prompt}'")
         yield {full_dataset_generation_label: gr.Label({f"⚙️ ReWriting {dataset}": 0.}, visible=True)}
-        for rows in batched(islice(stream_rows(dataset=dataset, subset=subset, split=split), total), n=10):
-            for row in stream_rewrite_dataset_row_by_row(dataset=dataset, rows=rows, prompt=prompt, format=format, input_preview_rows=input_preview_rows, output_preview_rows=output_preview_rows):
-                print(row)
-                output_rows.append({k: json.dumps(row[k], ensure_ascii=False) for k in output_format_df["column"]})
-                yield {full_dataset_generation_label: gr.Label({f"⚙️ ReWriting {dataset}": len(output_rows) / total})}
-        # repo_id = namespace + "/" + find_new_name(dataset, prompt)
-        # 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
-        # Dataset.from_list(output_rows).push_to_hub(repo_id, config_name=subset, split=split, 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 {pretty_full_dataset_generation_output: empty_dataframe}
+        yield {
+            pretty_output_preview: gr.DataFrame(visible=False),
+            pretty_full_dataset_generation_output: gr.DataFrame(visible=True),
+        }
+
+        num_parallel_calls = max(1, min(total // NUM_ROWS_PER_CALL, NUM_PARALLEL_CALLS))
+        parallel_input_rows = list(batched(islice(stream_rows(dataset=dataset, subset=subset, split=split), total), n=total // num_parallel_calls))
+        parallel_output_rows = [[] for _ in range(num_parallel_calls)]
+
+        def run(i):
+            for batch_rows in batched(parallel_input_rows[i], n=NUM_ROWS_PER_CALL):
+                for row in stream_rewrite_dataset_row_by_row(dataset=dataset, rows=batch_rows, prompt=prompt, format=format, input_preview_rows=input_preview_rows, output_preview_rows=output_preview_rows):
+                    parallel_output_rows[i].append({k: json.dumps(row[k], ensure_ascii=False) for k in output_format_df["column"]})
+                    yield 1
+
+        current = 0
+        _last_time = time.time()
+        for step in iflatmap_unordered(run, kwargs_iterable=[{"i": i} for i in range(num_parallel_calls)]):
+            current += step
+            if _last_time + 1 / MAX_PROGRESS_UPDATES_PER_SECOND < time.time():
+                _last_time = time.time()
+                yield {
+                    full_dataset_generation_label: gr.Label({f"⚙️ ReWriting {dataset}": current / total}),
+                    pretty_full_dataset_generation_output: gr.DataFrame(pd.DataFrame([row for rows in parallel_output_rows for row in rows]))
+                }
+        yield {
+            full_dataset_generation_label: gr.Label({f"⚙️ ReWriting {dataset}": current / total}),
+            pretty_full_dataset_generation_output: gr.DataFrame(pd.DataFrame([row for rows in parallel_output_rows for row in rows]))
+        }
+        print(f"Done ReWriting {dataset} (full 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]
+        repo_id = namespace + "/" + find_new_name(dataset, prompt)
+        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
+        print(f"Saving {repo_id}")
+        ds = Dataset.from_list(output_rows)
+        ds.push_to_hub(repo_id, config_name=subset, split=split, 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})"}
+        print(f"Saved {repo_id}")
 
 
 demo.launch()