implement revrite preview

.gitignore

@@ -0,0 +1 @@
+__*

app.py

@@ -1,15 +1,35 @@
+import json
+import time
 from itertools import count, islice
-from typing import Any, Iterable
+from multiprocessing.pool import ThreadPool
+from queue import Queue, Empty
+from typing import Any, Callable, Iterable, Iterator, TypeVar
 
 import gradio as gr
+import ijson
 import pandas as pd
 import requests
+from datasets import Features, Value, Sequence
 from gradio_huggingfacehub_search import HuggingfaceHubSearch
+from huggingface_hub import InferenceClient
 
+from utils import StringIteratorIO
+
+
+model_id = "meta-llama/Meta-Llama-3.1-8B-Instruct"
+client = InferenceClient(model_id)
 
 session = requests.Session()
 empty_dataframe = pd.DataFrame({"1": [], "2": [], "3": []})
-NUM_ROWS_PREVIEW = 5
+
+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). "
+    "Try to keep some of the text or meaning intact, and apply the requested transformation '{prompt}'."
+)
 
 
 with gr.Blocks() as demo:
@@ -27,15 +47,17 @@ with gr.Blocks() as demo:
         subset_dropdown = gr.Dropdown(info="Subset", show_label=False, visible=False)
         split_dropdown = gr.Dropdown(info="Split", show_label=False, visible=False)
 
-    input_query = gr.Textbox(label="Enter the adjustment or transformation to apply to the dataset:")
-    rewrite_button = gr.Button("ReWrite Dataset", variant="primary")
-
     gr.Markdown("### Input")
-    input_preview = gr.DataFrame(interactive=False, wrap=True)
+    input_preview = gr.DataFrame(visible=False)
+    pretty_input_preview = gr.DataFrame(interactive=False, wrap=True)
 
-    gr.Markdown("### Output")
+    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)
+    rewrite_button = gr.Button("ReWrite Dataset", variant="primary")
     output_preview = gr.DataFrame(interactive=False, wrap=True)
-    save_button = gr.Button("Save ReWriten Dataset", interactive=False)
+    save_button = gr.Button("ReWrite Full Dataset", interactive=False)
 
 
     ############
@@ -56,6 +78,98 @@ with gr.Blocks() as demo:
                 yield row_item["row"]
 
 
+    T = TypeVar("T")
+
+
+    def batched(it: Iterable[T], n: int) -> Iterator[list[T]]:
+        it = iter(it)
+        while batch := list(islice(it, n)):
+            yield batch
+
+
+    def stream_reponse(messages: list[dict[str: str]], response_format=None) -> Iterator[str]:
+        for _ in range(3):
+            message = None
+            try:
+                for message in client.chat_completion(
+                    messages=messages,
+                    max_tokens=5000,
+                    stream=True,
+                    top_p=0.8,
+                    seed=42,
+                    response_format=response_format
+                ):
+                    yield message.choices[0].delta.content
+            except requests.exceptions.ConnectionError as e:
+                if message:
+                    raise
+                print(e + "\n\nRetrying in 1sec")
+                time.sleep(1)
+                continue
+            break
+
+
+    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}),
+            prompt=prompt,
+        )}]
+        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))
+        print("done")
+
+
+    def _write_generator_to_queue(queue: Queue, func: Callable[..., Iterable], kwargs: dict) -> None:
+        for i, result in enumerate(func(**kwargs)):
+            queue.put(result)
+        return None
+
+
+    def iflatmap_unordered(
+        func: Callable[..., Iterable[T]],
+        *,
+        kwargs_iterable: Iterable[dict],
+    ) -> Iterable[T]:
+        queue = Queue()
+        with ThreadPool() as pool:
+            async_results = [pool.apply_async(_write_generator_to_queue, (queue, func, kwargs)) for kwargs in kwargs_iterable]
+            try:
+                while True:
+                    try:
+                        yield queue.get(timeout=0.05)
+                    except Empty:
+                        if all(async_result.ready() for async_result in async_results) and queue.empty():
+                            break
+            finally:  # in case there's an error to raise
+                [async_result.get(timeout=0.05) for async_result in async_results]
+
+
+    def features_to_format(features: Features) -> dict:
+        def feature_to_format(feature):
+            if isinstance(feature, Value):
+                if "int" in feature.dtype:
+                    return {"type": "integer"} 
+                elif "float" in feature.dtype:
+                    return {"type": "number"}
+                else:
+                    return {"type": "string"}
+            elif isinstance(feature, list):
+                return {"type": "array", "items": feature_to_format(feature[0])}
+            elif isinstance(feature, dict):
+                return {"properties": {k: feature_to_format(v) for k, v in feature.items()}, "required": list(feature)}
+            elif isinstance(feature, Sequence):
+                if isinstance(feature.feature, dict):
+                    return {"properties": {k: {"type": "array", "items": v } for k, v in feature_to_format(feature.feature).items()}, "required": list(feature)}
+                else:
+                    return {"type": "array", "items": feature_to_format(feature.feature)}
+            else:
+                return {"type": "string"}
+        return feature_to_format(features)
+
+
     ############
     #
     #  Events
@@ -78,9 +192,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)
         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)
         }
 
 
@@ -88,32 +204,34 @@ 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
+        rows = list(islice((stream_rows(dataset, subset, split, batch_size=NUM_ROWS_PREVIEW)), NUM_ROWS_PREVIEW))
         return {
-            input_preview: pd.DataFrame(islice(({
-                k: str(v) for k, v in row.items()}
-                for row in stream_rows(dataset, subset, split, batch_size=NUM_ROWS_PREVIEW)
-            ), NUM_ROWS_PREVIEW)),
+            input_preview: pd.DataFrame(rows),
+            pretty_input_preview: pd.DataFrame([{k: str(v) for k, v in row.items()} for row in rows]),
             **output
         }
 
 
-    @dataset_search.change(inputs=[dataset_search], outputs=[input_preview, subset_dropdown, split_dropdown])
+    @dataset_search.change(inputs=[dataset_search], outputs=[input_preview, pretty_input_preview, subset_dropdown, split_dropdown, output_format])
     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, subset_dropdown, split_dropdown])
+    @subset_dropdown.change(inputs=[dataset_search, subset_dropdown], outputs=[input_preview, pretty_input_preview, subset_dropdown, split_dropdown, output_format])
     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, subset_dropdown, split_dropdown])
+    @split_dropdown.change(inputs=[dataset_search, subset_dropdown, split_dropdown], outputs=[input_preview, pretty_input_preview, subset_dropdown, split_dropdown, output_format])
     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], outputs=[output_preview])
-    def rewrite(dataset: str, subset: str, split: str, input_preview_df: pd.DataFrame) -> dict:
-        # TODO: implement
-        return {output_preview: pd.DataFrame([{"TODO": ["implement"]}])}
+    @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")
+        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)
+            yield pd.DataFrame(output_rows)
 
 
 demo.launch()

requirements.txt

@@ -1,3 +1,4 @@
 requests
 pandas
 gradio_huggingfacehub_search
+datasets

utils.py

@@ -0,0 +1,60 @@
+import io
+import logging
+
+
+logger = logging.getLogger(__name__)
+
+
+class StringIteratorIO(io.TextIOBase):
+    """From: https://stackoverflow.com/a/12604375"""
+
+    def __init__(self, iter):
+        self._iter = iter
+        self._left = ''
+
+    def readable(self):
+        return True
+
+    def _read1(self, n=None):
+        while not self._left:
+            try:
+                self._left = next(self._iter)
+            except StopIteration:
+                break
+        ret = self._left[:n]
+        self._left = self._left[len(ret):]
+        return ret
+
+    def read(self, n=None):
+        buf = []
+        if n is None or n < 0:
+            while True:
+                m = self._read1()
+                if not m:
+                    break
+                buf.append(m)
+        else:
+            while n > 0:
+                m = self._read1(n)
+                if not m:
+                    break
+                n -= len(m)
+                buf.append(m)
+        return ''.join(buf)
+
+    def readline(self):
+        buf = []
+        while True:
+            i = self._left.find('\n')
+            if i == -1:
+                buf.append(self._left)
+                try:
+                    self._left = next(self._iter)
+                except StopIteration:
+                    self._left = ''
+                    break
+            else:
+                buf.append(self._left[:i+1])
+                self._left = self._left[i+1:]
+                break
+        return ''.join(buf)