column feature

app.py

@@ -90,6 +90,12 @@ a {
     color: var(--body-text-color-subdued);
 }
 """
+js = """
+function load() {
+    Array.from(document.getElementsByClassName("secondary")).filter(e => (e.innerText.includes("New row")))[0].innerText = "New column"
+    return 'done';
+}
+"""
 
 examples = [
     ["fka/awesome-chatgpt-prompts", "make the prompt 6 words long maximum"],
@@ -97,7 +103,7 @@ examples = [
     ["infinite-dataset-hub/SmallTalkDialogues", "translate to proper French"],
 ]
 
-with gr.Blocks(css=css) as demo:
+with gr.Blocks(css=css, js=js) as demo:
     dataset_info_json = gr.JSON(visible=False)
     with gr.Row():
         with gr.Column(scale=10):
@@ -123,6 +129,12 @@ with gr.Blocks(css=css) as demo:
                 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"])
+                    column_ro_remove_dropdown = gr.Dropdown(info="Select a column to remove", show_label=False)
+                    with gr.Row():
+                        with gr.Column(scale=99):
+                            pass
+                        with gr.Column(scale=1, min_width=88):
+                            remove_column_button = gr.Button("Remove", size="sm", elem_id="remove_column_button")
             rewrite_preview_button = gr.Button("Preview Results", variant="primary")
             rewrite_full_dataset_button = gr.Button("ReWrite Full Dataset", interactive=False)
             gr.Markdown("#### Output")
@@ -226,14 +238,14 @@ with gr.Blocks(css=css) as demo:
             print("Warning: Some rows were missing during ReWriting.")
 
 
-    def find_new_name(dataset: str, prompt: str) -> str:
+    def find_new_name(dataset: str, prompt: str, format: dict) -> str:
         messages = [{"role": "user", "content": FIND_NEW_NAME.format(prompt=prompt)}]
         out = "".join(stream_reponse(messages))
         if "should be" in out:
             out = dataset.split("/")[-1] + out.split("should be", 1)[1].replace(" ", "-").replace(".", "").replace(",", "")
         else:
             out = dataset.split("/")[-1] + prompt.replace(" ", "-")
-        return out[:80] + "-" + Hasher.hash(prompt)[:4]
+        return out[:80] + "-" + Hasher.hash(prompt + str(format))[:4]
 
     def _write_generator_to_queue(queue: Queue, func: Callable[..., Iterable], kwargs: dict) -> None:
         for i, result in enumerate(func(**kwargs)):
@@ -342,6 +354,14 @@ with gr.Blocks(css=css) as demo:
     @input_prompt.change(outputs=[rewrite_full_dataset_button])
     def disable_rewrite_full_dataset() -> dict:
         return {rewrite_full_dataset_button: gr.Button(interactive=False)}
+    
+    @output_format_dataframe.change(inputs=[output_format_dataframe], outputs=[column_ro_remove_dropdown])
+    def update_columns_to_remove_dropdown(output_format_df: pd.DataFrame) -> dict:
+        return gr.Dropdown(choices=output_format_df["column"].tolist())
+    
+    @remove_column_button.click(inputs=[column_ro_remove_dropdown, output_format_dataframe], outputs=[output_format_dataframe])
+    def update_output_format_dataframe(column: str, output_format_df: pd.DataFrame) -> pd.DataFrame:
+        return output_format_df[output_format_df["column"] != column]
 
 
     @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, full_dataset_generation_success_html, pretty_full_dataset_generation_output])
@@ -368,6 +388,7 @@ with gr.Blocks(css=css) as demo:
     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")]
+        output_format_df = output_format_df[output_format_df["column"] != ""]
         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"]
@@ -411,7 +432,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[max_num_rows] if num_examples > total else ""), prompt)
+        new_dataset = find_new_name(dataset + (PARTIAL_SUFFIX[max_num_rows] if num_examples > total else ""), prompt, format)
         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