fetch parquet filename via dataset-viewer api

app.py

@@ -81,20 +81,28 @@ def plot_and_df(texts, preds):
         )
 
 
+def get_first_parquet_filename(dataset, config, split):
+    parquet_resp = session.get(f"https://datasets-server.huggingface.co/parquet?dataset={dataset}&config={config}", timeout=3).json()
+    if "error" in parquet_resp:
+        raise ValueError(parquet_resp["error"])
+    first_parquet_file_url = [file for file in parquet_resp["parquet_files"] if file["split"] == split][0]["url"]
+    return "/".join(first_parquet_file_url.split("/")[-3:])
+
+
 @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:
-        data = pl.read_parquet(f"hf://datasets/{dataset}@~parquet/{config}/{split}/0000.parquet", columns=[column])
-    except pl.exceptions.ComputeError:
-        try:
-            data = pl.read_parquet(f"hf://datasets/{dataset}@~parquet/{config}/partial-{split}/0000.parquet", columns=[column])
-        except pl.exceptions.ComputeError:
-            try:
-                data = pl.read_parquet(f"hf://datasets/{dataset}@~parquet/{config}/{split}-part0/0000.parquet", columns=[column])
-            except Exception as error:
-                yield f"❌ {error}", gr.BarPlot(), pd.DataFrame(), pd.DataFrame(), pd.DataFrame(), 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:
+        data = pl.read_parquet(f"hf://datasets/{dataset}@~parquet/{filename}", columns=[column])
+    except Exception as error:
+        yield f"❌ {error}", gr.BarPlot(), pd.DataFrame(), pd.DataFrame(), pd.DataFrame(), pd.DataFrame(),
+        return
     logging.info("Data fetched.")
 
     data_sample = data.sample(num_examples, seed=16) if data.shape[0] > num_examples else data