app.py

import os
import json
import time
from itertools import count, islice
from multiprocessing.pool import ThreadPool
from queue import Queue, Empty
from typing import Any, Callable, Iterable, Iterator, Optional, TypeVar

import gradio as gr
import ijson
import pandas as pd
import requests
from datasets import Dataset, Features, Value, Sequence
from datasets.fingerprint import Hasher
from gradio_huggingfacehub_search import HuggingfaceHubSearch
from huggingface_hub import DatasetCard, InferenceClient

from utils import StringIteratorIO


model_id = "meta-llama/Meta-Llama-3.1-8B-Instruct"
client = InferenceClient(model_id, token=os.environ.get("INFERENCE_API_HF_TOKEN"))

save_dataset_hf_token = os.environ.get("SAVE_DATASET_HF_TOKEN")
session = requests.Session()
empty_dataframe = pd.DataFrame({"1": [], "2": [], "3": []})

NAMESPACE = "dataset-rewriter"
URL = "https://huggingface.co/spaces/dataset-rewrite/dataset-rewrite"

NUM_ROWS_PREVIEW = 3
MAX_NUM_ROWS_TO_REWRITE = 1000
PARTIAL_SUFFIX = "-1k"
NUM_PARALLEL_CALLS = 10
NUM_ROWS_PER_CALL = 10
MAX_PROGRESS_UPDATES_PER_SECOND = 4
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 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'. "
    "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}\"'"
)


DATASET_CARD_CONTENT = """
---
license: mit
tags:
- dataset-rewriter
- synthetic
---

# {new_dataset}

_Note: This is an AI-generated dataset so its content may be inaccurate or false_

**Source of the data:**
The dataset was generated using the [Dataset ReWriter]({url}) and {model_id} from the dataset {dataset} and using the prompt '{prompt}':
- **Original Dataset**: https://huggingface.co/datasets/{dataset}
- **Model**: https://huggingface.co/{model_id}
- **More Datasets**: https://huggingface.co/datasets?other=dataset-rewriter
"""

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=10):
            gr.Markdown(
                "# 🤗 WIP Dataset ReWriter ✍️✨\n\n"
                "Adjust, translate or transform completely existing datasets.\n\n"
            )
            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("### Sample")
            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")
            rewrite_full_dataset_button = gr.Button("ReWrite Full Dataset", interactive=False)
            gr.Markdown("#### Output")
            full_dataset_generation_label = gr.Label(visible=False, show_label=False)
            pretty_output_preview = gr.DataFrame(interactive=False)
            full_dataset_generation_success_markdown = gr.Markdown("")
            pretty_full_dataset_generation_output = gr.DataFrame(interactive=False, visible=False)
        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)


    ############
    #
    #  Utils
    #
    ###########


    def stream_rows(dataset: str, subset: str, split: str, batch_size: int = 100) -> Iterable[dict[str, Any]]:
        for i in count():
            rows_resp = session.get(f"https://datasets-server.huggingface.co/rows?dataset={dataset}&config={subset}&split={split}&offset={i * batch_size}&length={batch_size}", timeout=10).json()
            if "error" in rows_resp:
                raise RuntimeError(rows_resp["error"])
            if not rows_resp["rows"]:
                break
            for row_item in rows_resp["rows"]:
                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_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"]}}
        yield from ijson.items(StringIteratorIO(stream_reponse(messages, response_format=response_format)), "data.item", buf_size=4)


    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", "items": format}}, "required": ["data"]}}
        yield from ijson.items(StringIteratorIO(stream_reponse(messages, response_format=response_format)), "data.item", buf_size=4)


    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:
            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]

    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
    #
    ###########

    def _resolve_dataset_selection(dataset: str, default_subset: str, default_split: str) -> dict:
        if "/" not in dataset.strip().strip("/"):
            return None, None, {
                subset_dropdown: gr.Dropdown(visible=False),
                split_dropdown: gr.Dropdown(visible=False),
            }
        info_resp = session.get(f"https://datasets-server.huggingface.co/info?dataset={dataset}", timeout=3).json()
        if "error" in info_resp:
            return None, None, {
                subset_dropdown: gr.Dropdown(visible=False),
                split_dropdown: gr.Dropdown(visible=False),
            }
        subsets: list[str] = list(info_resp["dataset_info"])
        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]
        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()])
        }


    def _show_input_preview(dataset: str, default_subset: str, default_split: str) -> dict:
        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 {
            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=[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, 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, 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_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, pretty_full_dataset_generation_output])
    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"(preview) ReWriting {dataset} with instruction '{prompt}'")
        yield {rewrite_full_dataset_button: gr.Button(interactive=False), full_dataset_generation_label: gr.Label(visible=False)}
        yield {
            pretty_output_preview: gr.DataFrame(visible=True),
            pretty_full_dataset_generation_output: gr.DataFrame(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)}
        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], outputs=[full_dataset_generation_label, full_dataset_generation_success_markdown, 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, 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]}
        num_examples = dataset_info["splits"][split]["num_examples"]
        total = min(num_examples, MAX_NUM_ROWS_TO_REWRITE)
        print(f"ReWriting {dataset} with instruction '{prompt}'")
        yield {full_dataset_generation_label: gr.Label({f"⚙️ ReWriting {dataset}": 0.}, visible=True)}
        yield {pretty_full_dataset_generation_output: empty_dataframe}
        yield {
            pretty_output_preview: gr.DataFrame(visible=False),
            pretty_full_dataset_generation_output: gr.DataFrame(visible=True),
        }

        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_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"]})
                    yield 1

        current = 0
        _last_time = time.time()
        for step in iflatmap_unordered(run, kwargs_iterable=[{"i": i} for i in range(num_parallel_calls)]):
            current += step
            if _last_time + 1 / MAX_PROGRESS_UPDATES_PER_SECOND < time.time():
                _last_time = time.time()
                yield {
                    full_dataset_generation_label: gr.Label({f"⚙️ ReWriting {dataset}": current / total}),
                    pretty_full_dataset_generation_output: gr.DataFrame(pd.DataFrame([row for rows in parallel_output_rows for row in rows]))
                }
        yield {
            full_dataset_generation_label: gr.Label({f"⚙️ ReWriting {dataset}": current / total}),
            pretty_full_dataset_generation_output: gr.DataFrame(pd.DataFrame([row for rows in parallel_output_rows for row in rows]))
        }
        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 if num_examples > total else ""), prompt)
        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
        print(f"Saving {repo_id}")
        ds = Dataset.from_list(output_rows)
        ds.push_to_hub(repo_id, config_name=subset, split=split, token=token)
        DatasetCard(DATASET_CARD_CONTENT.format(new_dataset=new_dataset, dataset=dataset, model_id=model_id, prompt=prompt, url=URL)).push_to_hub(repo_id=repo_id, repo_type="dataset", 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})"}
        print(f"Saved {repo_id}")


demo.launch()