initial commit

README.md

@@ -1,6 +1,6 @@
 ---
-title: Datasets Slices
-emoji: 👀
+title: Datasets Explorer
+emoji: 📖 
 colorFrom: pink
 colorTo: purple
 sdk: gradio
@@ -9,4 +9,12 @@ app_file: app.py
 pinned: false
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# 📖 Dataset Explorer
+
+Access any slice of data of any dataset on the [Hugging Face Dataset Hub](https://huggingface.co/datasets)
+
+Run:
+
+```python
+gradio app.py
+```

app.py

@@ -0,0 +1,186 @@
+import gradio as gr
+from functools import lru_cache
+from hffs.fs import HfFileSystem
+from typing import List, Tuple, Callable
+import pandas as pd
+import numpy as np
+import pyarrow as pa
+import pyarrow.parquet as pq
+from functools import partial
+from io import StringIO
+from tqdm.contrib.concurrent import thread_map
+from datasets import Features
+
+
+
+class AppError(RuntimeError):
+    pass
+
+
+PAGE_SIZE = 20
+
+
+@lru_cache(maxsize=128)
+def get_parquet_fs(dataset: str) -> HfFileSystem:
+    try:
+        fs = HfFileSystem(dataset, repo_type="dataset", revision="refs/convert/parquet")
+        if any(fs.isfile(path) for path in fs.ls("") if not path.startswith(".")):
+            raise AppError(f"Parquet export doesn't exist for '{dataset}'.")
+        return fs
+    except:
+        raise AppError(f"Parquet export doesn't exist for '{dataset}'.")
+
+
+@lru_cache(maxsize=128)
+def get_parquet_configs(dataset: str) -> List[str]:
+    fs = get_parquet_fs(dataset)
+    return [path for path in fs.ls("") if fs.isdir(path)]
+
+
+def _sorted_split_key(split: str) -> str:
+    return split if not split.startswith("train") else chr(0) + split  # always "train" first
+
+
+@lru_cache(maxsize=128)
+def get_parquet_splits(dataset: str, config: str) -> List[str]:
+    fs = get_parquet_fs(dataset)
+    all_parts = [path.rsplit(".", 1)[0].split("-") for path in fs.glob(f"{config}/*.parquet")]
+    return sorted(set(parts[-4] if len(parts) > 3 and parts[-2] == "of" else parts[-1] for parts in all_parts), key=_sorted_split_key)
+
+def sanitize_inputs(dataset: str, config: str, split: str, page:  str) -> Tuple[str, str, str, int]:
+    try:
+        page = int(page)
+        assert page > 0
+    except:
+        raise AppError(f"Bad page: {page}")
+    if not dataset:
+        raise AppError("Empty dataset name")
+    if not config:
+        raise AppError(f"Empty config. Available configs are: {', '.join(get_parquet_configs(dataset))}.")
+    if not split:
+        raise AppError(f"Empty split. Available splits are: {', '.join(get_parquet_splits(dataset, config))}.")
+    return dataset, config, split, int(page)
+
+
+RowGroupReaders = List[Callable[[], pa.Table]]
+
+
+@lru_cache(maxsize=128)
+def index(dataset: str, config: str, split: str) -> Tuple[np.ndarray, RowGroupReaders, int, str]:
+    fs = get_parquet_fs(dataset)
+    sources = fs.glob(f"{config}/*-{split}.parquet") + fs.glob(f"{config}/*-{split}-*-of-*.parquet")
+    if not sources:
+        if config not in get_parquet_configs(dataset):
+            raise AppError(f"Invalid config {config}. Available configs are: {', '.join(get_parquet_configs(dataset))}.")
+        else:
+            raise AppError(f"Invalid split {split}. Available splits are: {', '.join(get_parquet_splits(dataset, config))}.")
+    all_pf: List[pq.ParquetFile] = thread_map(partial(pq.ParquetFile, filesystem=fs), sources)
+    features = Features.from_arrow_schema(all_pf[0].schema.to_arrow_schema())
+    columns = [col for col in features if all(bad_type not in str(features[col]) for bad_type in ["Image(", "Audio(", "'binary'"])]
+    info = "" if len(columns) == len(features) else f"Some columns are not supported yet: {sorted(set(features) - set(columns))}"
+    rg_offsets = np.cumsum([pf.metadata.row_group(i).num_rows for pf in all_pf for i in range(pf.metadata.num_row_groups)])
+    rg_readers = [partial(pf.read_row_group, i, columns=columns) for pf in all_pf for i in range(pf.metadata.num_row_groups)]
+    max_page = rg_offsets[-1] // PAGE_SIZE
+    return rg_offsets, rg_readers, max_page, info
+
+
+def query(page: int, page_size: int, rg_offsets: np.ndarray, rg_readers: RowGroupReaders) -> pd.DataFrame:
+    start_row, end_row = (page - 1) * page_size, page * page_size
+    start_rg, end_rg = np.searchsorted(rg_offsets, [start_row, end_row], side="right")
+    if page < 1 or end_rg >= len(rg_readers):
+        raise AppError(f"Page {page} does not exist")
+    pa_table = pa.concat_tables([rg_readers[i]() for i in range(start_rg, end_rg + 1)])
+    offset = start_row - rg_offsets[start_rg - 1] if start_rg else start_row
+    pa_table = pa_table.slice(offset, end_row - start_row)
+    return pa_table.to_pandas()
+
+
+@lru_cache(maxsize=128)
+def get_page(dataset: str, config: str, split: str, page: str) -> Tuple[str, int, str]:
+    dataset, config, split, page = sanitize_inputs(dataset, config, split, page)
+    rg_offsets, rg_readers, max_page, info = index(dataset, config, split)
+    df = query(page, PAGE_SIZE, rg_offsets=rg_offsets, rg_readers=rg_readers)
+    buf = StringIO()
+    df.to_json(buf, lines=True, orient="records")
+    return buf.getvalue(), max_page, info
+
+
+with gr.Blocks() as demo:
+    gr.Markdown("# 📖 Dataset Explorer\n\nAccess any slice of data of any dataset on the [Hugging Face Dataset Hub](https://huggingface.co/datasets)")
+    cp_dataset = gr.Textbox("squad", label="Pick a dataset", placeholder="squad")
+    cp_go = gr.Button("Explore")
+    cp_config = gr.Dropdown(["plain_text"], value="plain_text", label="Config", visible=False)
+    cp_split = gr.Dropdown(["train", "validation"], value="train", label="Split", visible=False)
+    with gr.Row():
+        cp_page = gr.Textbox("1", label="Page", placeholder="1", visible=False)
+        cp_goto_page = gr.Button("Go to page", visible=False)
+    cp_error = gr.Markdown("", visible=False)
+    cp_info = gr.Markdown("", visible=False)
+    cp_result = gr.Markdown("", visible=False)
+
+    def show_error(message: str) -> dict():
+        return {
+            cp_error: gr.update(visible=True, value=f"## ❌ Error:\n\n{message}"),
+            cp_info: gr.update(visible=False, value=""),
+            cp_result: gr.update(visible=False, value=""),
+        }
+        
+    def show_dataset_at_config_and_split_and_page(dataset: str, config: str, split: str, page: str) -> dict:
+        try:
+            jsonl_result, max_page, info = get_page(dataset, config, split, page)
+            info = f"({info})" if info else ""
+            return {
+                cp_result: gr.update(visible=True, value=f"```json\n{jsonl_result}\n```"),
+                cp_info: gr.update(visible=True, value=f"Page {page}/{max_page}) {info}"),
+                cp_error: gr.update(visible=False, value="")
+            }
+        except AppError as err:
+            return show_error(str(err))
+
+    def show_dataset_at_config_and_split(dataset: str, config: str, split: str) -> dict:
+        try:
+            return {
+                **show_dataset_at_config_and_split_and_page(dataset, config, split, "1"),
+                cp_page: gr.update(value="1", visible=True),
+                cp_goto_page: gr.update(visible=True),
+            }
+        except AppError as err:
+            return show_error(str(err))
+
+    def show_dataset_at_config(dataset: str, config: str) -> dict:
+        try:
+            splits = get_parquet_splits(dataset, config)
+            if not splits:
+                raise AppError(f"Dataset {dataset} with config {config} has no splits.")
+            else:
+                split = splits[0]
+            return {
+                **show_dataset_at_config_and_split(dataset, config, split),
+                cp_split: gr.update(value=split, choices=splits, visible=len(splits) > 1),
+            }
+        except AppError as err:
+            return show_error(str(err))
+
+    def show_dataset(dataset: str) -> dict:
+        try:
+            configs = get_parquet_configs(dataset)
+            if not configs:
+                raise AppError(f"Dataset {dataset} has no configs.")
+            else:
+                config = configs[0]
+            return {
+                **show_dataset_at_config(dataset, config),
+                cp_config: gr.update(value=config, choices=configs, visible=len(configs) > 1),
+            }
+        except AppError as err:
+            return show_error(str(err))
+
+    all_outputs = [cp_config, cp_split, cp_page, cp_goto_page, cp_result, cp_info, cp_error]
+    cp_go.click(show_dataset, inputs=[cp_dataset], outputs=all_outputs)
+    cp_config.change(show_dataset_at_config, inputs=[cp_dataset, cp_config], outputs=all_outputs)
+    cp_split.change(show_dataset_at_config_and_split, inputs=[cp_dataset, cp_config, cp_split], outputs=all_outputs)
+    cp_goto_page.click(show_dataset_at_config_and_split_and_page, inputs=[cp_dataset, cp_config, cp_split, cp_page], outputs=all_outputs)
+
+
+if __name__ == "__main__":
+    demo.launch()

requirements.txt

@@ -0,0 +1,5 @@
+pyarrow
+fsspec[http]
+tqdm
+datasets
+git+https://github.com/huggingface/hffs.git