get config and split with api, include partial datasets

app.py

@@ -1,13 +1,22 @@
+import requests
+from collections import Counter
+from requests.adapters import HTTPAdapter, Retry
+
 import gradio as gr
+import pandas as pd
 import polars as pl
+import spaces
 from gradio_huggingfacehub_search import HuggingfaceHubSearch
+from huggingface_hub import PyTorchModelHubMixin
 import torch
-import spaces
 from torch import nn
 from transformers import AutoModel, AutoTokenizer, AutoConfig
-from huggingface_hub import PyTorchModelHubMixin
-import pandas as pd
-from collections import Counter
+
+
+
+session = requests.Session()
+retries = Retry(total=5, backoff_factor=1, status_forcelist=[502, 503, 504])
+session.mount('http://', HTTPAdapter(max_retries=retries))
 
 
 class QualityModel(nn.Module, PyTorchModelHubMixin):
@@ -64,8 +73,22 @@ def plot_and_df(texts, preds):
 
 
 def run_quality_check(dataset, column, batch_size, num_examples):
-    config = "default"
-    data = pl.read_parquet(f"hf://datasets/{dataset}@~parquet/{config}/train/0000.parquet", columns=[column])
+    # config = "default"
+    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()
+        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"]))
+    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 Exception as error:
+            yield f"❌ {error}", gr.BarPlot(), pd.DataFrame(), pd.DataFrame(), pd.DataFrame()
+            return
     texts = data[column].to_list()
     # batch_size = 100
     predictions, texts_processed = [], []
@@ -106,8 +129,8 @@ with gr.Blocks() as demo:
             return gr.HTML(value=html_code)
 
     text_column = gr.Textbox(placeholder="text", label="Text colum name to check (data must be non-nested, raw texts!)")
-    batch_size = gr.Slider(0, 128, 64, step=8, label="Inference batch size (set this to smaller value if this space crashes.)")
-    num_examples = gr.Number(1000, label="Number of first examples to check")
+    batch_size = gr.Slider(0, 128, 32, step=8, 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)
     plot = gr.BarPlot()