add full generation

README.md

@@ -1,12 +1,17 @@
 ---
 title: Dataset Rewriter
-emoji: 🏃
+emoji: ✍️✨
 colorFrom: purple
 colorTo: green
 sdk: gradio
 sdk_version: 4.42.0
 app_file: app.py
 pinned: false
+hf_oauth: true
+hf_oauth_scopes:
+  - read-repos
+  - write-repos
+  - manage-repos
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -4,13 +4,13 @@ import time
 from itertools import count, islice
 from multiprocessing.pool import ThreadPool
 from queue import Queue, Empty
-from typing import Any, Callable, Iterable, Iterator, TypeVar
+from typing import Any, Callable, Iterable, Iterator, Optional, TypeVar
 
 import gradio as gr
 import ijson
 import pandas as pd
 import requests
-from datasets import Features, Value, Sequence
+from datasets import Dataset, Features, Value, Sequence
 from gradio_huggingfacehub_search import HuggingfaceHubSearch
 from huggingface_hub import InferenceClient
 
@@ -20,45 +20,83 @@ 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"))
 
+save_dataset_hf_token = os.environ.get("SAVE_DATASET_HF_TOKEN")
 session = requests.Session()
 empty_dataframe = pd.DataFrame({"1": [], "2": [], "3": []})
 
+NAMESPACE = "lhoestq"
+
 NUM_ROWS_PREVIEW = 3
-REWRITE_DATASET = (
+MAX_NUM_ROWS_TO_REWRITE = 10
+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 transformation: {prompt}."
+    "They want you to rewrite the dataset and apply this instruction, which can be about transforming, translating or filtering the rows: {prompt}."
     "The first rows of the dataset are below in JSON format:\n\n{rows}\n\n"
-    "Rewrite those rows from the '{dataset}' dataset using the same JSON format. "
-    "Try to keep some of the text or meaning intact, and apply the requested transformation '{prompt}'."
+    "Apply the instruction to those rows from the '{dataset}' dataset and output the resulting rows using the same JSON format. "
+    "Try to keep some of the text or meaning intact, and apply the requested instruction '{prompt}'."
+)
+REWRITE_DATASET= (
+    "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}."
+    "Here is an example:\n\nOriginal rows:\n{input_preview_rows}\n\Resulting rows:\n{output_preview_rows}\n\n"
+    "The rows of the dataset are below in JSON format:\n\n{rows}\n\n"
+    "Apply the instruction to those rows from the '{dataset}' dataset and output the resulting rows using the same JSON format. "
+    "Try to keep some of the text or meaning intact, and apply the requested instruction '{prompt}'."
+)
+FIND_NEW_NAME = (
+    "You are a helpful assistant specialized in transforming english sentences for machine learning practitioners."
+    "Your job is to take input sentences like 'Take this dataset and apply the instruction xxx' and rephrase them them as 'The dataset should be yyy'. "
+    "You shoud use adjectives and exactly follow the output formula 'The dataset should be yyy'. "
+    "Here is your first job: rephrase the sentence 'Take this dataset and apply the instruction \"{prompt}\"'"
 )
 
-
-with gr.Blocks() as demo:
-    gr.Markdown(
-        "# 🤗 WIP Dataset ReWriter ✍️✨\n\n"
-        "Adjust, translate or transform completely existing datasets.\n\n"
-    )
+css = """
+.settings {
+    background: transparent;
+}
+.settings button span {
+    color: var(--body-text-color-subdued);
+}
+"""
+
+with gr.Blocks(css=css) as demo:
+    dataset_info_json = gr.JSON(visible=False)
     with gr.Row():
-        with gr.Column(scale=3):
-            dataset_search =  HuggingfaceHubSearch(
-                label="Hub Dataset ID",
-                placeholder="Search for dataset id on Huggingface",
-                search_type="dataset",
+        with gr.Column(scale=10):
+            gr.Markdown(
+                "# 🤗 WIP Dataset ReWriter ✍️✨\n\n"
+                "Adjust, translate or transform completely existing datasets.\n\n"
             )
-        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")
-    pretty_input_preview = gr.DataFrame(interactive=False)
-
-    gr.Markdown("### ReWrite")
-    with gr.Group():
-        input_prompt = gr.Textbox(label="Enter the 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_button = gr.Button("ReWrite Dataset", variant="primary")
-    pretty_output_preview = gr.DataFrame(interactive=False)
-    save_button = gr.Button("ReWrite Full Dataset", interactive=False)
+            with gr.Row():
+                with gr.Column(scale=3):
+                    dataset_search =  HuggingfaceHubSearch(
+                        label="Hub Dataset ID",
+                        placeholder="Search for dataset id on Huggingface",
+                        search_type="dataset",
+                    )
+                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")
+            pretty_input_preview = gr.DataFrame(interactive=False)
+
+            gr.Markdown("### ReWrite")
+            with gr.Group():
+                input_prompt = gr.Textbox(label="Enter the 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")
+            pretty_output_preview = gr.DataFrame(interactive=False)
+            rewrite_full_dataset_button = gr.Button("ReWrite Full Dataset", interactive=False)
+            full_dataset_generation_label = gr.Label(visible=False, show_label=False)
+            full_dataset_generation_success_markdown = gr.Markdown("")
+        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)
 
 
     ############
@@ -110,19 +148,42 @@ with gr.Blocks() as demo:
             break
 
 
-    def stream_rewrite_dataset_row_by_row(dataset: str, rows: list[dict[str, str]], prompt: str, format: str) -> Iterator[dict[str, str]]:
+    def stream_rewrite_dataset_preview_row_by_row(dataset: str, rows: list[dict[str, str]], prompt: str, format: str) -> Iterator[dict[str, str]]:
+        prompt = prompt[:1000] if prompt.strip() else ""
+        messages = [{"role": "user", "content": REWRITE_DATASET_PREVIEW.format(
+            dataset=dataset,
+            rows=json.dumps({"data": rows}),
+            prompt=prompt,
+        )}]
+        response_format = {"type": "json", "value": {"properties": {"data": {"type": "array", "items": format}}, "required": ["data"]}}
+        print("streaming preview")
+        yield from ijson.items(StringIteratorIO(stream_reponse(messages, response_format=response_format)), "data.item", buf_size=4)
+        print("done")
+
+
+    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]]:
         prompt = prompt[:1000] if prompt.strip() else ""
         messages = [{"role": "user", "content": REWRITE_DATASET.format(
             dataset=dataset,
             rows=json.dumps({"data": rows}),
             prompt=prompt,
+            input_preview_rows=json.dumps({"data": input_preview_rows}),
+            output_preview_rows=json.dumps({"data": output_preview_rows}),
         )}]
-        response_format = {"type": "json", "value": {"properties": {"data": {"type": "array", "maxItems": len(rows), "minItems": len(rows), "items": format}}, "required": ["data"]}}
-        print("go")
-        yield from islice(ijson.items(StringIteratorIO(stream_reponse(messages, response_format=response_format)), "data.item", buf_size=4), len(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(",", "")
+        else:
+            return dataset.split("/")[-1] + prompt.replace(" ", "-")
+
     def _write_generator_to_queue(queue: Queue, func: Callable[..., Iterable], kwargs: dict) -> None:
         for i, result in enumerate(func(**kwargs)):
             queue.put(result)
@@ -195,6 +256,7 @@ with gr.Blocks() as demo:
         split = default_split if default_split in splits else splits[0]
         dict_format = features_to_format(Features.from_dict(info_resp["dataset_info"][subset]["features"]))
         return subset, split, {
+            dataset_info_json: info_resp["dataset_info"][subset],
             subset_dropdown: gr.Dropdown(value=subset, choices=subsets, visible=len(subsets) > 1),
             split_dropdown: gr.Dropdown(value=split, choices=splits, visible=len(splits) > 1),
             output_format_dataframe: pd.DataFrame([{"column": col, "type": json.dumps(format_type)} for col, format_type in dict_format["properties"].items()])
@@ -213,29 +275,60 @@ with gr.Blocks() as demo:
         }
 
 
-    @dataset_search.change(inputs=[dataset_search], outputs=[pretty_input_preview, subset_dropdown, split_dropdown, output_format_dataframe])
+    @dataset_search.change(inputs=[dataset_search], outputs=[pretty_input_preview, subset_dropdown, split_dropdown, output_format_dataframe, dataset_info_json])
     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])
+    @subset_dropdown.change(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])
+    @split_dropdown.change(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)
+    
+
+    @input_prompt.change(outputs=[rewrite_full_dataset_button])
+    def disable_rewrite_full_dataset() -> dict:
+        return {rewrite_full_dataset_button: gr.Button(interactive=False)}
 
 
-    @rewrite_button.click(inputs=[dataset_search, subset_dropdown, split_dropdown, pretty_input_preview, input_prompt, output_format_dataframe], outputs=[pretty_output_preview])
-    def rewrite(dataset: str, subset: str, split: str, pretty_input_preview_df: pd.DataFrame, prompt: str, output_format_df: pd.DataFrame) -> Iterator[pd.DataFrame]:
+    @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])
+    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")
         format = {"properties": {x["column"]: json.loads(x["type"]) for x in format}, "required": [x["column"] for x in format]}
         output_rows = []
-        print(f"ReWriting {dataset} with instructions '{prompt}'")
-        for row in stream_rewrite_dataset_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 pretty_input_preview_df.columns})
-            yield pd.DataFrame(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)}
+        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])
+    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})"}
 
 
 demo.launch()