add feature dropdown

app.py

@@ -83,14 +83,7 @@ def plot_and_df(texts, preds):
 
 @spaces.GPU
 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}")
+    logging.info(f"Fetching data for {dataset=} {config=} {split=} {column=}")
     try:
         data = pl.read_parquet(f"hf://datasets/{dataset}@~parquet/{config}/{split}/0000.parquet", columns=[column])
     except pl.exceptions.ComputeError:
@@ -244,7 +237,6 @@ with gr.Blocks() as demo:
                     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)
@@ -263,40 +255,47 @@ with gr.Blocks() as demo:
                 """
             return gr.HTML(value=html_code)
 
+    text_column_dropdown = gr.Dropdown(label="Text column name", info="Text colum name to check (only non-nested texts are supported)")
+
     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),
+                text_column_dropdown: gr.Dropdown(info="Text colum name to check (only non-nested texts are supported)"),
             }
         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),
+                text_column_dropdown: gr.Dropdown(label="Text column name", info="Text colum name to check (only non-nested texts are supported)")
             }
         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]
+        features = info_resp["dataset_info"][subset]["features"]
+        text_features = [feature_name for feature_name, feature in features.items() if isinstance(feature, dict) and feature.get("dtype") == "string"] # and feature.get("_type") == "Value"]
         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),
+            text_column_dropdown: gr.Dropdown(choices=text_features, label="Text column name", info="Text colum name to check (only non-nested texts are supported)"),
         }
 
-    @dataset_name.change(inputs=[dataset_name], outputs=[subset_dropdown, split_dropdown])
+    @dataset_name.change(inputs=[dataset_name], outputs=[subset_dropdown, split_dropdown, text_column_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])
+    @subset_dropdown.change(inputs=[dataset_name, subset_dropdown], outputs=[subset_dropdown, split_dropdown, text_column_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])
+    @split_dropdown.change(inputs=[dataset_name, subset_dropdown, split_dropdown], outputs=[subset_dropdown, split_dropdown, text_column_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!)")
+    # 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, 64, 32, step=4, label="Inference batch size (set this to smaller value if this space crashes.)")
@@ -317,17 +316,17 @@ with gr.Blocks() as demo:
     texts_df = gr.DataFrame(visible=False)
     gr_check_btn.click(
         run_quality_check,
-        inputs=[dataset_name, subset_dropdown, split_dropdown, text_column, batch_size, num_examples],
+        inputs=[dataset_name, subset_dropdown, split_dropdown, text_column_dropdown, batch_size, num_examples],
         outputs=[progress_bar, plot, df_low, df_medium, df_high, texts_df]
     )
 
-    gr.Markdown("""## Compute text quality measures
-                * proportion of non-ascii characters
-                * #TODO""")
-    gr_ascii_btn = gr.Button("Data measures")
-    non_ascii_hist = gr.Plot()
-
-    gr_ascii_btn.click(non_ascii_check, inputs=[texts_df, text_column], outputs=[non_ascii_hist])
+    # gr.Markdown("""## Compute text quality measures
+    #             * proportion of non-ascii characters
+    #             * #TODO""")
+    # gr_ascii_btn = gr.Button("Data measures")
+    # non_ascii_hist = gr.Plot()
+    #
+    # gr_ascii_btn.click(non_ascii_check, inputs=[texts_df, text_column], outputs=[non_ascii_hist])
 
     gr.Markdown("## Explore toxicity")
     checkbox = gr.Checkbox(value=False, label="Run on full first parquet data (better not)")
@@ -338,7 +337,7 @@ with gr.Blocks() as demo:
         toxicity_df = gr.DataFrame()
     gr_toxicity_btn.click(
         call_perspective_api,
-        inputs=[texts_df, text_column, checkbox],
+        inputs=[texts_df, text_column_dropdown, checkbox],
         outputs=[toxicity_progress_bar, toxicity_hist, toxicity_df]
     )