add proper structured generation

app.py

@@ -27,8 +27,8 @@ NUM_ROWS_PREVIEW = 3
 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 transformation: {prompt}."
-    "The first rows of the dataset are below in JSON format (one JSON object per line):\n\n{rows}\n\n"
-    "Rewrite those rows from the '{dataset}' dataset using the same format (one JSON object per line). "
+    "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}'."
 )
 
@@ -49,15 +49,15 @@ with gr.Blocks() as demo:
         split_dropdown = gr.Dropdown(info="Split", show_label=False, visible=False)
 
     gr.Markdown("### Input")
-    input_preview = gr.DataFrame(visible=False)
-    pretty_input_preview = gr.DataFrame(interactive=False, wrap=True)
+    pretty_input_preview = gr.DataFrame(interactive=False)
 
     gr.Markdown("### ReWrite")
-    input_prompt = gr.Textbox(label="Enter the adjustment or transformation to apply to the dataset:")
-    with gr.Accordion("Modify Format", open=False):
-        output_format = gr.Textbox(interactive=True, show_label=False, container=False)
+    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")
-    output_preview = gr.DataFrame(interactive=False, wrap=True)
+    pretty_output_preview = gr.DataFrame(interactive=False)
     save_button = gr.Button("ReWrite Full Dataset", interactive=False)
 
 
@@ -193,11 +193,11 @@ with gr.Blocks() as demo:
         subset = default_subset if default_subset in subsets else subsets[0]
         splits: list[str] = info_resp["dataset_info"][subset]["splits"]
         split = default_split if default_split in splits else splits[0]
-        json_format = json.dumps(features_to_format(Features.from_dict(info_resp["dataset_info"][subset]["features"])), indent=2)
+        dict_format = features_to_format(Features.from_dict(info_resp["dataset_info"][subset]["features"]))
         return subset, split, {
             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: gr.Textbox(json_format, lines=json_format.count("\n") + 1)
+            output_format_dataframe: pd.DataFrame([{"column": col, "type": json.dumps(format_type)} for col, format_type in dict_format["properties"].items()])
         }
 
 
@@ -205,33 +205,36 @@ with gr.Blocks() as demo:
         subset, split, output = _resolve_dataset_selection(dataset, default_subset=default_subset, default_split=default_split)
         if subset is None or split is None:
             return output
+        print(f"Showing {dataset}")
         rows = list(islice((stream_rows(dataset, subset, split, batch_size=NUM_ROWS_PREVIEW)), NUM_ROWS_PREVIEW))
         return {
-            input_preview: pd.DataFrame(rows),
-            pretty_input_preview: pd.DataFrame([{k: str(v) for k, v in row.items()} for row in rows]),
+            pretty_input_preview: gr.DataFrame(pd.DataFrame([{k: json.dumps(v, ensure_ascii=False) for k, v in row.items()} for row in rows])),
             **output
         }
 
 
-    @dataset_search.change(inputs=[dataset_search], outputs=[input_preview, pretty_input_preview, subset_dropdown, split_dropdown, output_format])
+    @dataset_search.change(inputs=[dataset_search], outputs=[pretty_input_preview, subset_dropdown, split_dropdown, output_format_dataframe])
     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=[input_preview, pretty_input_preview, subset_dropdown, split_dropdown, output_format])
+    @subset_dropdown.change(inputs=[dataset_search, subset_dropdown], outputs=[pretty_input_preview, subset_dropdown, split_dropdown, output_format_dataframe])
     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=[input_preview, pretty_input_preview, subset_dropdown, split_dropdown, output_format])
+    @split_dropdown.change(inputs=[dataset_search, subset_dropdown, split_dropdown], outputs=[pretty_input_preview, subset_dropdown, split_dropdown, output_format_dataframe])
     def show_input_from_split_dropdown(dataset: str, subset: str, split: str) -> dict:
         return _show_input_preview(dataset, default_subset=subset, default_split=split)
 
 
-    @rewrite_button.click(inputs=[dataset_search, subset_dropdown, split_dropdown, input_preview, input_prompt, output_format], outputs=[output_preview])
-    def rewrite(dataset: str, subset: str, split: str, input_preview_df: pd.DataFrame, prompt: str, json_format: str) -> Iterator[pd.DataFrame]:
-        rows = input_preview_df.to_dict(orient="records")
+    @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]:
+        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 = []
-        for row in stream_rewrite_dataset_row_by_row(dataset=dataset, rows=rows, prompt=prompt, format=json.loads(json_format)):
-            output_rows.append(row)
+        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)