fetch data for toxicity if it doesn't exist yet

app.py

@@ -10,7 +10,7 @@ import gradio as gr
 import pandas as pd
 import polars as pl
 import matplotlib.pyplot as plt
-import spaces
+# import spaces
 from gradio_huggingfacehub_search import HuggingfaceHubSearch
 from huggingface_hub import PyTorchModelHubMixin
 import torch
@@ -50,7 +50,7 @@ model = QualityModel.from_pretrained("nvidia/quality-classifier-deberta").to(dev
 model.eval()
 
 
-@spaces.GPU
+# @spaces.GPU
 def predict(texts: list[str]):
     inputs = tokenizer(
         texts, return_tensors="pt", padding="longest", truncation=True
@@ -81,7 +81,11 @@ def plot_and_df(texts, preds):
         )
 
 
-@spaces.GPU
+# def download_data(dataset, config, split, column):
+#
+
+
+# @spaces.GPU
 def run_quality_check(dataset, config, split, column, batch_size, num_examples):
     logging.info(f"Fetching data for {dataset=} {config=} {split=} {column=}")
     try:
@@ -97,9 +101,8 @@ def run_quality_check(dataset, config, split, column, batch_size, num_examples):
                 return
     logging.info("Data fetched.")
 
-    texts = [text[:10000] for text in data[column].to_list()]
-    # texts_sample = data.sample(100, shuffle=True, seed=16).to_pandas()
-    # batch_size = 100
+    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()]
     predictions, texts_processed = [], []
     num_examples = min(len(texts), num_examples)
     for i in range(0, num_examples, batch_size):
@@ -118,7 +121,7 @@ def run_quality_check(dataset, config, split, column, batch_size, num_examples):
     # plt.xlabel('Proportion of non-ASCII characters')
     # plt.ylabel('Number of texts')
 
-    yield {"finished": 1.}, *plot_and_df(texts_processed, predictions), data
+    yield {"finished": 1.}, *plot_and_df(texts_processed, predictions), data_sample
 
 
 PERSPECTIVE_API_KEY = os.environ.get("PERSPECTIVE_API_KEY")
@@ -141,13 +144,31 @@ def plot_toxicity(scores):
 
     return fig
 
-def call_perspective_api(texts_df, column_name, full_check=False):
+def call_perspective_api(texts_df, column_name, dataset, config, split):#, full_check=False):
     headers = {
         "content-type": "application/json",
     }
     req_att_scores = {attr: [] for attr in REQUESTED_ATTRIBUTES}
 
-    texts = texts_df.sample(100, random_state=16)[column_name].values if not full_check else texts_df[column_name].values
+    # fetch data if it doesn't exist yet
+    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
+        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
 
     n_samples = len(texts)
     for i, text in tqdm(enumerate(texts), desc="scanning with perspective"):
@@ -165,8 +186,6 @@ def call_perspective_api(texts_df, column_name, full_check=False):
 
         if req_response.ok:
             response = req_response.json()
-            # logger.info("Perspective API response is:")
-            # logger.info(response)
             if ATT_SCORE in response:
                 for req_att in REQUESTED_ATTRIBUTES:
                     if req_att in response[ATT_SCORE]:
@@ -175,15 +194,12 @@ def call_perspective_api(texts_df, column_name, full_check=False):
                     else:
                         req_att_scores[req_att].append(0)
             else:
-                # logger.error(
-                #     "Unexpected response format from Perspective API."
-                # )
                 raise ValueError(req_response)
         else:
             try:
                 req_response.raise_for_status()
             except Exception as e:
-                print(e)
+                logging.info(e)
                 return req_att_scores
         if i % 10 == 0:
             plot_toxicity(req_att_scores)
@@ -295,11 +311,9 @@ with gr.Blocks() as demo:
     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, 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")
+    num_examples = gr.Slider(0, 1000, 500, step=10, label="Number of random examples to check")
     gr_check_btn = gr.Button("Check Dataset")
     progress_bar = gr.Label(show_label=False)
     plot = gr.BarPlot()
@@ -329,7 +343,7 @@ with gr.Blocks() as demo:
     # 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)")
+    # checkbox = gr.Checkbox(value=False, label="Run on full first parquet data (better not)")
     gr_toxicity_btn = gr.Button("Run perpspective API to check toxicity of random samples.")
     toxicity_progress_bar = gr.Label(show_label=False)
     toxicity_hist = gr.Plot()
@@ -337,7 +351,7 @@ with gr.Blocks() as demo:
         toxicity_df = gr.DataFrame()
     gr_toxicity_btn.click(
         call_perspective_api,
-        inputs=[texts_df, text_column_dropdown, checkbox],
+        inputs=[texts_df, text_column_dropdown, dataset_name, subset_dropdown, split_dropdown],#, checkbox],
         outputs=[toxicity_progress_bar, toxicity_hist, toxicity_df]
     )