limit max string length and fix input columns

app.py

@@ -37,21 +37,14 @@ assert MAX_NUM_ROWS_TO_REWRITE in PARTIAL_SUFFIX, "allowed max num rows are 100,
 NUM_PARALLEL_CALLS = 10
 NUM_ROWS_PER_CALL = 3
 MAX_PROGRESS_UPDATES_PER_SECOND = 4
-REWRITE_DATASET_PREVIEW = (
+MAX_STRING_LENGTH = 1000
+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}."
     "The first 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}'."
 )
-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'. "
@@ -186,6 +179,13 @@ with gr.Blocks(css=css, js=js) as demo:
     class ContextTooLongError(ValueError):
         pass
 
+    def crop_text(obj: Any) -> str:
+        if isinstance(obj, str):
+            return obj[:MAX_STRING_LENGTH]
+        else:
+            raise TypeError()
+
+
     def stream_reponse(messages: list[dict[str: str]], response_format=None, max_tokens=5000) -> Iterator[str]:
         for _ in range(3):
             message = None
@@ -212,23 +212,21 @@ with gr.Blocks(css=css, js=js) as demo:
 
     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(
+        messages = [{"role": "user", "content": REWRITE_DATASET.format(
             dataset=dataset,
-            rows=json.dumps({"data": rows}, ensure_ascii=False),
+            rows=json.dumps({"data": rows}, ensure_ascii=False, default=crop_text),
             prompt=prompt,
         )}]
         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, use_float=True)
 
 
-    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]]:
+    def stream_rewrite_dataset_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.format(
             dataset=dataset,
-            rows=json.dumps({"data": rows}, ensure_ascii=False),
+            rows=json.dumps({"data": rows}, ensure_ascii=False, default=crop_text),
             prompt=prompt,
-            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, "minItems": len(rows), "maxItems": len(rows)}}, "required": ["data"]}}
         try:
@@ -333,7 +331,7 @@ with gr.Blocks(css=css, js=js) as demo:
         print(f"Showing {dataset}")
         rows = list(islice((stream_rows(dataset, subset, split, batch_size=NUM_ROWS_PREVIEW)), NUM_ROWS_PREVIEW))
         return {
-            pretty_input_preview: gr.DataFrame(pd.DataFrame([{k: json.dumps(v, ensure_ascii=False) for k, v in row.items()} for row in rows])),
+            pretty_input_preview: gr.DataFrame(pd.DataFrame([{k: json.dumps(v, ensure_ascii=False, default=crop_text) for k, v in row.items()} for row in rows])),
             **output
         }
 
@@ -379,19 +377,17 @@ with gr.Blocks(css=css, js=js) as demo:
             full_dataset_generation_success_html: "",
         }
         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"]})
+            output_rows.append({k: json.dumps(row[k], ensure_ascii=False, default=crop_text) 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)}
         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, 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]:
+    @rewrite_full_dataset_button.click(inputs=[dataset_search, subset_dropdown, split_dropdown, 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, 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]:
         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]}
-        input_preview_rows = [{k: json.loads(row[k]) for k in output_format_df["column"] if k in row} 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")]
         num_examples = dataset_info["splits"][split]["num_examples"]
         total = min(num_examples, max_num_rows)
         print(f"ReWriting {dataset} with instruction '{prompt}'")
@@ -404,13 +400,13 @@ with gr.Blocks(css=css, js=js) as demo:
         }
 
         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_input_rows = list(batched(islice(({k: row[k] for k in output_format_df["column"] if k in row} for row in 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"]})
+                for row in stream_rewrite_dataset_row_by_row(dataset=dataset, rows=batch_rows, prompt=prompt, format=format):
+                    parallel_output_rows[i].append({k: json.dumps(row[k], ensure_ascii=False, default=crop_text) for k in output_format_df["column"]})
                     yield 1
 
         current = 0