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 tqdm.contrib.concurrent import thread_map from datasets import Features, Image, Audio from fastapi import FastAPI, Response import uvicorn import os class AppError(RuntimeError): pass APP_URL = "http://127.0.0.1:7860" if os.getenv("DEV") else "https://lhoestq-datasets-explorer.hf.space" PAGE_SIZE = 20 MAX_CACHED_BLOBS = PAGE_SIZE * 10 _blobs_cache = {} ##################################################### # Define routes for image and audio files ##################################################### app = FastAPI() @app.get( "/image", responses={200: {"content": {"image/png": {}}}}, response_class=Response, ) def image(id: str): blob = get_blob(id) return Response(content=blob, media_type="image/png") @app.get( "/audio", responses={200: {"content": {"audio/wav": {}}}}, response_class=Response, ) def audio(id: str): blob = get_blob(id) return Response(content=blob, media_type="audio/wav") def push_blob(blob: bytes, blob_id: str) -> str: global _blobs_cache if blob_id in _blobs_cache: del _blobs_cache[blob_id] _blobs_cache[blob_id] = blob if len(_blobs_cache) > MAX_CACHED_BLOBS: del _blobs_cache[next(iter(_blobs_cache))] return blob_id def get_blob(blob_id: str) -> bytes: global _blobs_cache return _blobs_cache[blob_id] def blobs_to_urls(blobs: List[bytes], type: str, prefix: str) -> List[str]: image_blob_ids = [push_blob(blob, f"{prefix}-{i}") for i, blob in enumerate(blobs)] return [APP_URL + f"/{type}?id={blob_id}" for blob_id in image_blob_ids] ##################################################### # List configs, splits and parquet files ##################################################### @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) return [path.split("/")[1] for path in fs.ls(config) if fs.isdir(path)] ##################################################### # Index and query Parquet data ##################################################### RowGroupReaders = List[Callable[[], pa.Table]] @lru_cache(maxsize=128) def index(dataset: str, config: str, split: str) -> Tuple[np.ndarray, RowGroupReaders, int, Features]: fs = get_parquet_fs(dataset) sources = fs.glob(f"{config}/{split}/*.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))}.") desc = f"{dataset}/{config}/{split}" all_pf: List[pq.ParquetFile] = thread_map(partial(pq.ParquetFile, filesystem=fs), sources, desc=desc, unit="pq") features = Features.from_arrow_schema(all_pf[0].schema.to_arrow_schema()) 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) for pf in all_pf for i in range(pf.metadata.num_row_groups)] max_page = 1 + (rg_offsets[-1] - 1) // PAGE_SIZE return rg_offsets, rg_readers, max_page, features def query(page: int, page_size: int, rg_offsets: np.ndarray, rg_readers: RowGroupReaders) -> pd.DataFrame: start_row, end_row = (page - 1) * page_size, min(page * page_size, rg_offsets[-1] - 1) # both included # rg_offsets[start_rg - 1] <= start_row < rg_offsets[start_rg] # rg_offsets[end_rg - 1] <= end_row < rg_offsets[end_rg] start_rg, end_rg = np.searchsorted(rg_offsets, [start_row, end_row], side="right") # both included 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 > 0 else 0) pa_table = pa_table.slice(offset, page_size) return pa_table.to_pandas() 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) @lru_cache(maxsize=128) def get_page_df(dataset: str, config: str, split: str, page: str) -> Tuple[pd.DataFrame, int, Features]: dataset, config, split, page = sanitize_inputs(dataset, config, split, page) rg_offsets, rg_readers, max_page, features = index(dataset, config, split) if page > max_page: raise AppError(f"Page {page} does not exist") df = query(page, PAGE_SIZE, rg_offsets=rg_offsets, rg_readers=rg_readers) return df, max_page, features ##################################################### # Format results ##################################################### def get_page(dataset: str, config: str, split: str, page: str) -> Tuple[str, int, str]: df, max_page, features = get_page_df(dataset, config, split, page) unsupported_columns = [] for column, feature in features.items(): if isinstance(feature, Image): blob_type = "image" # TODO: support audio - right now it seems that the markdown renderer in gradio doesn't support audio and shows nothing blob_urls = blobs_to_urls([item.get("bytes") if isinstance(item, dict) else None for item in df[column]], blob_type, prefix=f"{dataset}-{config}-{split}-{page}-{column}") df = df.drop([column], axis=1) df[column] = [f"![]({url})" for url in blob_urls] elif any(bad_type in str(feature) for bad_type in ["Image(", "Audio(", "'binary'"]): unsupported_columns.append(column) df = df.drop([column], axis=1) info = "" if not unsupported_columns else f"Some columns are not supported yet: {unsupported_columns}" return df.to_markdown(index=False), max_page, info ##################################################### # Gradio app ##################################################### with gr.Blocks() as demo: gr.Markdown("# 📖 Datasets Explorer\n\nAccess any slice of data of any dataset on the [Hugging Face Dataset Hub](https://huggingface.co/datasets)") gr.Markdown("This is the dataset viewer from parquet export demo before the feature was added on the Hugging Face website.") cp_dataset = gr.Textbox("frgfm/imagenette", label="Pick a dataset", placeholder="competitions/aiornot") 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) cp_goto_next_page = gr.Button("Next page", visible=False) cp_error = gr.Markdown("", visible=False) cp_info = gr.Markdown("", visible=False) cp_result = gr.Markdown("", 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) 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: markdown_result, max_page, info = get_page(dataset, config, split, page) info = f"({info})" if info else "" return { cp_result: gr.update(visible=True, value=markdown_result), 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_and_next_page(dataset: str, config: str, split: str, page: str) -> dict: try: next_page = str(int(page) + 1) return { **show_dataset_at_config_and_split_and_page(dataset, config, split, next_page), cp_page: gr.update(value=next_page, visible=True), } 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), cp_goto_next_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_goto_next_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) cp_goto_next_page.click(show_dataset_at_config_and_split_and_next_page, inputs=[cp_dataset, cp_config, cp_split, cp_page], outputs=all_outputs) if __name__ == "__main__": app = gr.mount_gradio_app(app, demo, path="/") uvicorn.run(app, host="0.0.0.0", port=7860)