add config and split dropdown

app.py

@@ -19,7 +19,7 @@ from transformers import AutoModel, AutoTokenizer, AutoConfig
 from tqdm import tqdm
 
 
-logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
+logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s")
 
 
 session = requests.Session()
@@ -74,7 +74,7 @@ def plot_and_df(texts, preds):
     )
     # counts.reset_index(inplace=True)
     return (
-            gr.BarPlot(counts_df, x="quality", y="count"),
+            gr.BarPlot(counts_df, x="quality", y="count", sort=None),
             texts_df[texts_df["quality"] == "Low"][["text"]][:20],
             texts_df[texts_df["quality"] == "Medium"][["text"]][:20],
             texts_df[texts_df["quality"] == "High"][["text"]][:20],
@@ -82,14 +82,14 @@ def plot_and_df(texts, preds):
 
 
 @spaces.GPU
-def run_quality_check(dataset, column, batch_size, num_examples):
-    info_resp = session.get(f"https://datasets-server.huggingface.co/info?dataset={dataset}", timeout=3).json()
-    if "error" in info_resp:
-        yield "❌ " + info_resp["error"], gr.BarPlot(), pd.DataFrame(), pd.DataFrame(), pd.DataFrame(), pd.DataFrame(),
-        return
-    config = "default" if "default" in info_resp["dataset_info"] else next(iter(info_resp["dataset_info"]))
-    split = "train" if "train" in info_resp["dataset_info"][config]["splits"] else next(
-        iter(info_resp["dataset_info"][config]["splits"]))
+def run_quality_check(dataset, config, split, column, batch_size, num_examples):
+    # info_resp = session.get(f"https://datasets-server.huggingface.co/info?dataset={dataset}", timeout=3).json()
+    # if "error" in info_resp:
+    #     yield "❌ " + info_resp["error"], gr.BarPlot(), pd.DataFrame(), pd.DataFrame(), pd.DataFrame(), pd.DataFrame(),
+    #     return
+    # config = "default" if "default" in info_resp["dataset_info"] else next(iter(info_resp["dataset_info"]))
+    # split = "train" if "train" in info_resp["dataset_info"][config]["splits"] else next(
+    #     iter(info_resp["dataset_info"][config]["splits"]))
     logging.info(f"Fetching data for {dataset} {config} {split}")
     try:
         data = pl.read_parquet(f"hf://datasets/{dataset}@~parquet/{config}/{split}/0000.parquet", columns=[column])
@@ -238,30 +238,68 @@ with gr.Blocks() as demo:
         ## Select dataset and text column
         """
     )
-    dataset_name = HuggingfaceHubSearch(
-            label="Hub Dataset ID",
-            placeholder="Search for dataset id on Huggingface",
-            search_type="dataset",
-            # value="fka/awesome-chatgpt-prompts",
-        )
+    with gr.Row():
+        with gr.Column(scale=3):
+            dataset_name = HuggingfaceHubSearch(
+                    label="Hub Dataset ID",
+                    placeholder="Search for dataset id on Huggingface",
+                    search_type="dataset",
+                    # value="fka/awesome-chatgpt-prompts",
+                )
+        subset_dropdown = gr.Dropdown(info="Subset", show_label=False, visible=False)
+        split_dropdown = gr.Dropdown(info="Split", show_label=False, visible=False)
+
     # config_name = "default"  # TODO: user input
     with gr.Accordion("Dataset preview", open=False):
-        @gr.render(inputs=dataset_name)
-        def embed(name):
+        @gr.render(inputs=[dataset_name, subset_dropdown, split_dropdown])
+        def embed(name, subset, split):
             html_code = f"""
             <iframe
-              src="https://huggingface.co/datasets/{name}/embed/viewer/default/train"
+              src="https://huggingface.co/datasets/{name}/embed/viewer/{subset}/{split}"
               frameborder="0"
               width="100%"
-              height="700px"
+              height="600px"
             ></iframe>
                 """
             return gr.HTML(value=html_code)
 
+    def _resolve_dataset_selection(dataset: str, default_subset: str, default_split: str):
+        if "/" not in dataset.strip().strip("/"):
+            return {
+                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 {
+                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]
+        return {
+            subset_dropdown: gr.Dropdown(value=subset, choices=subsets, visible=len(subsets) > 1),
+            split_dropdown: gr.Dropdown(value=split, choices=splits, visible=len(splits) > 1),
+        }
+
+    @dataset_name.change(inputs=[dataset_name], outputs=[subset_dropdown, split_dropdown])
+    def show_input_from_subset_dropdown(dataset: str) -> dict:
+        return _resolve_dataset_selection(dataset, default_subset="default", default_split="train")
+
+    @subset_dropdown.change(inputs=[dataset_name, subset_dropdown], outputs=[subset_dropdown, split_dropdown])
+    def show_input_from_subset_dropdown(dataset: str, subset: str) -> dict:
+        return _resolve_dataset_selection(dataset, default_subset=subset, default_split="train")
+
+    @split_dropdown.change(inputs=[dataset_name, subset_dropdown, split_dropdown], outputs=[subset_dropdown, split_dropdown])
+    def show_input_from_split_dropdown(dataset: str, subset: str, split: str) -> dict:
+        return _resolve_dataset_selection(dataset, default_subset=subset, default_split=split)
+
     text_column = gr.Textbox(placeholder="text", label="Text colum name to check (data must be non-nested, raw texts!)")
 
     gr.Markdown("## Run nvidia quality classifier")
-    batch_size = gr.Slider(0, 128, 32, step=8, label="Inference batch size (set this to smaller value if this space crashes.)")
+    batch_size = gr.Slider(0, 64, 32, step=4, label="Inference batch size (set this to smaller value if this space crashes.)")
     num_examples = gr.Number(500, label="Number of first examples to check")
     gr_check_btn = gr.Button("Check Dataset")
     progress_bar = gr.Label(show_label=False)
@@ -279,7 +317,7 @@ with gr.Blocks() as demo:
     texts_df = gr.DataFrame(visible=False)
     gr_check_btn.click(
         run_quality_check,
-        inputs=[dataset_name, text_column, batch_size, num_examples],
+        inputs=[dataset_name, subset_dropdown, split_dropdown, text_column, batch_size, num_examples],
         outputs=[progress_bar, plot, df_low, df_medium, df_high, texts_df]
     )