support images

app.py

@@ -7,18 +7,72 @@ 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
-from fastapi import FastAPI
+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)
@@ -48,26 +102,17 @@ def get_parquet_splits(dataset: str, config: str) -> List[str]:
     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)
+
+#####################################################
+# 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, str]:
+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") + fs.glob(f"{config}/*-{split}-*-of-*.parquet")
     if not sources:
@@ -78,12 +123,10 @@ def index(dataset: str, config: str, split: str) -> Tuple[np.ndarray, RowGroupRe
     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())
-    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)]
+    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, info
+    return rg_offsets, rg_readers, max_page, features
 
 
 def query(page: int, page_size: int, rg_offsets: np.ndarray, rg_readers: RowGroupReaders) -> pd.DataFrame:
@@ -97,30 +140,70 @@ def query(page: int, page_size: int, rg_offsets: np.ndarray, rg_readers: RowGrou
     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(dataset: str, config: str, split: str, page: str) -> Tuple[str, int, str]:
+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, info = index(dataset, config, split)
+    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)
-    buf = StringIO()
-    df.to_json(buf, lines=True, orient="records")
-    return buf.getvalue(), max_page, info
+    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)")
-    cp_dataset = gr.Textbox("squad", label="Pick a dataset", placeholder="squad")
+    cp_dataset = gr.Textbox("competitions/aiornot", 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)
-    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_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 {
@@ -131,15 +214,25 @@ with gr.Blocks() as demo:
 
     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)
+            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=f"```json\n{jsonl_result}\n```"),
+                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:
@@ -147,6 +240,7 @@ with gr.Blocks() as demo:
                 **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))
@@ -179,18 +273,12 @@ with gr.Blocks() as demo:
         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]
+    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)
-
-
-app = FastAPI()
-
-@app.get("/lol")
-def read_main():
-    return {"message": "This is your main app"}
+    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__":