add nested texts

app.py

@@ -107,7 +107,9 @@ def run_quality_check(dataset, config, split, column, batch_size, num_examples):
     logging.info("Data fetched.")
 
     data_sample = data.sample(num_examples, seed=16) if data.shape[0] > num_examples else data
-    texts = [text[:10000] for text in data_sample[column].to_list()]
+    texts = data_sample[column].to_list()
+    if nested_column:
+        texts = [text[nested_column] for text in texts]
     predictions, texts_processed = [], []
     num_examples = min(len(texts), num_examples)
     for i in range(0, num_examples, batch_size):
@@ -144,7 +146,7 @@ def plot_toxicity(scores):
 
     return fig
 
-def call_perspective_api(texts_df, column_name, dataset, config, split):#, full_check=False):
+def call_perspective_api(texts_df, column_name, nested_column_name, dataset, config, split):#, full_check=False):
     headers = {
         "content-type": "application/json",
     }
@@ -154,21 +156,23 @@ def call_perspective_api(texts_df, column_name, dataset, config, split):#, full_
     if texts_df.values.tolist() == [['', '', '']]:
         logging.info(f"Fetching data for {dataset=} {config=} {split=} {column_name=}")
         try:
-            texts_df = pl.read_parquet(f"hf://datasets/{dataset}@~parquet/{config}/{split}/0000.parquet", columns=[column_name])
-        except pl.exceptions.ComputeError:
-            try:
-                texts_df = pl.read_parquet(f"hf://datasets/{dataset}@~parquet/{config}/partial-{split}/0000.parquet", columns=[column_name])
-            except pl.exceptions.ComputeError:
-                try:
-                    texts_df = pl.read_parquet(f"hf://datasets/{dataset}@~parquet/{config}/{split}-part0/0000.parquet", columns=[column_name])
-                except Exception as error:
-                    yield f"❌ {error}", plt.gcf(), pd.DataFrame(),
-                    return
+            filename = get_first_parquet_filename(dataset, config, split)
+        except Exception as error:
+            yield f"❌ {error}", gr.BarPlot(), pd.DataFrame(), pd.DataFrame(), pd.DataFrame(), pd.DataFrame()
+            return
+
+        try:
+            logging.info(f"Loading hf://datasets/{dataset}@~parquet/{filename}")
+            texts_df = pl.read_parquet(f"hf://datasets/{dataset}@~parquet/{filename}", columns=[column_name])
+        except Exception as error:
+            yield f"❌ {error}", gr.BarPlot(), pd.DataFrame(), pd.DataFrame(), pd.DataFrame(), pd.DataFrame()
+            return
         logging.info("Data fetched.")
         texts_df = texts_df.to_pandas()
 
-    # texts = texts_df.sample(100, seed=16)[column_name].values if not full_check else texts_df[column_name].values
     texts = texts_df.sample(100, random_state=16)[column_name].values if texts_df.shape[0] > 100 else texts_df[column_name].values
+    if nested_column_name:
+        texts = [text[nested_column_name] for text in texts]
 
     n_samples = len(texts)
     for i, text in tqdm(enumerate(texts), desc="scanning with perspective"):
@@ -245,45 +249,80 @@ 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. ")
+    with gr.Row():
+        text_column_dropdown = gr.Dropdown(label="Text column name", info="Text colum name to check. ")
+        nested_text_column_dropdown = gr.Dropdown(label="Nested text key")#, visible=False)
 
-    def _resolve_dataset_selection(dataset: str, default_subset: str, default_split: str):
+    def _resolve_dataset_selection(dataset: str, default_subset: str, default_split: str, text_feature):
         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)"),
+                text_column_dropdown: gr.Dropdown(info="Text colum name to check"),
+                nested_text_column_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),
-                text_column_dropdown: gr.Dropdown(label="Text column name", info="Text colum name to check (only non-nested texts are supported)")
+                text_column_dropdown: gr.Dropdown(label="Text column name", info="Text colum name to check"),
+                nested_text_column_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]
         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"]
+
+        def _is_string_feature(feature):
+            return isinstance(feature, dict) and feature.get("dtype") == "string"
+
+        text_features = [feature_name for feature_name, feature in features.items() if _is_string_feature(feature)]
+        nested_features = [feature_name for feature_name, feature in features.items() if isinstance(feature, dict) and isinstance(next(iter(feature.values())), dict)]
+        nested_text_features = [feature_name for feature_name in nested_features if any(_is_string_feature(nested_feature) for nested_feature in features[feature_name].values())]
+        if not text_feature:
+            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 + nested_text_features, label="Text column name",
+                                                  info="Text colum name to check"),
+                nested_text_column_dropdown: gr.Dropdown(visible=False),
+            }
+        logging.info(nested_text_features)
+        if text_feature in nested_text_features:
+            nested_keys = [feature_name for feature_name, feature in features[text_feature].items() if _is_string_feature(feature)]
+            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 + nested_text_features,
+                                                  label="Text column name",
+                                                  info="Text colum name to check (only non-nested texts are supported)"),
+                nested_text_column_dropdown: gr.Dropdown(value=nested_keys[0], choices=nested_keys,
+                                                         label="Nested text column name", visible=True)
+            }
         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)"),
+            text_column_dropdown: gr.Dropdown(choices=text_features + nested_text_features, label="Text column name", info="Text colum name to check (only non-nested texts are supported)"),
+            nested_text_column_dropdown: gr.Dropdown(visible=False),
         }
 
-    @dataset_name.change(inputs=[dataset_name], outputs=[subset_dropdown, split_dropdown, text_column_dropdown])
+    @dataset_name.change(inputs=[dataset_name], outputs=[subset_dropdown, split_dropdown, text_column_dropdown, nested_text_column_dropdown])
     def show_input_from_subset_dropdown(dataset: str) -> dict:
-        return _resolve_dataset_selection(dataset, default_subset="default", default_split="train")
+        return _resolve_dataset_selection(dataset, default_subset="default", default_split="train", text_feature=None)
 
-    @subset_dropdown.change(inputs=[dataset_name, subset_dropdown], outputs=[subset_dropdown, split_dropdown, text_column_dropdown])
+    @subset_dropdown.change(inputs=[dataset_name, subset_dropdown], outputs=[subset_dropdown, split_dropdown, text_column_dropdown, nested_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")
+        return _resolve_dataset_selection(dataset, default_subset=subset, default_split="train", text_feature=None)
 
-    @split_dropdown.change(inputs=[dataset_name, subset_dropdown, split_dropdown], outputs=[subset_dropdown, split_dropdown, text_column_dropdown])
+    @split_dropdown.change(inputs=[dataset_name, subset_dropdown, split_dropdown], outputs=[subset_dropdown, split_dropdown, text_column_dropdown, nested_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)
+        return _resolve_dataset_selection(dataset, default_subset=subset, default_split=split, text_feature=None)
+
+    @text_column_dropdown.change(inputs=[dataset_name, subset_dropdown, split_dropdown, text_column_dropdown], outputs=[subset_dropdown, split_dropdown, text_column_dropdown, nested_text_column_dropdown])
+    def show_input_from_text_column_dropdown(dataset: str, subset: str, split: str, text_column) -> dict:
+        return _resolve_dataset_selection(dataset, default_subset=subset, default_split=split, text_feature=text_column)
 
     gr.Markdown("## Run nvidia quality classifier")
     batch_size = gr.Slider(0, 64, 32, step=4, label="Inference batch size", info="(set this to smaller value if this space crashes.)")
@@ -305,13 +344,13 @@ with gr.Blocks() as demo:
 
     gr.Examples(
         [
-            ["HuggingFaceFW/fineweb-edu", "default", "train", "text", 16, 500],
-            ["fka/awesome-chatgpt-prompts", "default", "train", "prompt", 64, 200],
-            ["proj-persona/PersonaHub", "instruction", "train", "synthesized text",  32, 1000],
-            ["argilla/FinePersonas-v0.1", "default", "train", "persona", 64, 1000],
-            ["Open-Orca/OpenOrca", "default", "train", "response", 16, 500],
+            ["HuggingFaceFW/fineweb-edu", "default", "train", "text", None, 16, 500],
+            # ["fka/awesome-chatgpt-prompts", "default", "train", "prompt", 64, 200],
+            # ["proj-persona/PersonaHub", "instruction", "train", "synthesized text",  32, 1000],
+            ["argilla/FinePersonas-v0.1", "default", "train", "persona", None, 64, 1000],
+            ["allenai/real-toxicity-prompts", "default", "train", "continuation", "text", 64, 1000],
         ],
-        [dataset_name, subset_dropdown, split_dropdown, text_column_dropdown, batch_size, num_examples],
+        [dataset_name, subset_dropdown, split_dropdown, text_column_dropdown, nested_text_column_dropdown, batch_size, num_examples],
         [progress_bar, plot, df_low, df_medium, df_high, texts_df],
         fn=run_quality_check,
         run_on_click=False,
@@ -320,7 +359,7 @@ with gr.Blocks() as demo:
 
     gr_check_btn.click(
         run_quality_check,
-        inputs=[dataset_name, subset_dropdown, split_dropdown, text_column_dropdown, batch_size, num_examples],
+        inputs=[dataset_name, subset_dropdown, split_dropdown, text_column_dropdown, nested_text_column_dropdown, batch_size, num_examples],
         outputs=[progress_bar, plot, df_low, df_medium, df_high, texts_df]
     )
 
@@ -335,7 +374,7 @@ with gr.Blocks() as demo:
         toxicity_df = gr.DataFrame()
     gr_toxicity_btn.click(
         call_perspective_api,
-        inputs=[texts_df, text_column_dropdown, dataset_name, subset_dropdown, split_dropdown],#, checkbox],
+        inputs=[texts_df, text_column_dropdown, nested_text_column_dropdown, dataset_name, subset_dropdown, split_dropdown],#, checkbox],
         outputs=[toxicity_progress_bar, toxicity_hist, toxicity_df]
     )