Create app.py

app.py

@@ -0,0 +1,299 @@
+"""
+BigBIO Dataset Explorer Demo
+"""
+
+from collections import Counter
+from collections import defaultdict
+import string
+
+from datasets import load_dataset
+from loguru import logger
+import numpy as np
+import pandas as pd
+import plotly.express as px
+import spacy
+from spacy import displacy
+import streamlit as st
+
+from bigbio.dataloader import BigBioConfigHelpers
+from bigbio.hf_maps import BATCH_MAPPERS_TEXT_FROM_SCHEMA
+from sklearn.feature_extraction.text import CountVectorizer
+
+
+st.set_page_config(layout="wide")
+
+
+IBM_COLORS = [
+    "#648fff",
+    "#dc267f",
+    "#ffb000",
+    "#fe6100",
+    "#785ef0",
+    "#000000",
+    "#ffffff",
+]
+
+
+def get_html(html: str):
+    """Convert HTML so it can be rendered."""
+    WRAPPER = """<div style="overflow-x: auto; border: 1px solid #e6e9ef; border-radius: 0.25rem; padding: 1rem;\
+ margin-bottom: 2.5rem">{}</div>"""
+    # Newlines seem to mess with the rendering
+    html = html.replace("\n", " ")
+    return WRAPPER.format(html)
+
+
+@st.cache()
+def load_conhelps():
+    conhelps = BigBioConfigHelpers()
+    logger.info(conhelps)
+    conhelps = conhelps.filtered(lambda x: not x.is_large)
+    conhelps = conhelps.filtered(lambda x: x.is_bigbio_schema)
+    conhelps = conhelps.filtered(lambda x: not x.is_local)
+    return conhelps
+
+
+def update_axis_font(fig):
+    fig.update_layout(
+        xaxis = dict(title_font = dict(size=20)),
+        yaxis = dict(title_font = dict(size=20)),
+    )
+    return fig
+
+
+def draw_histogram(hist_data, col_name, histnorm=None, nbins=25, xmax=None, loc=st):
+    fig = px.histogram(
+        hist_data,
+        x=col_name,
+        color="split",
+        color_discrete_sequence=IBM_COLORS,
+        marginal="box",  # or violin, rug
+        barmode="group",
+        hover_data=hist_data.columns,
+        histnorm=histnorm,
+        nbins=nbins,
+        range_x=(0, xmax) if xmax else None,
+    )
+    fig = update_axis_font(fig)
+    loc.plotly_chart(fig, use_container_width=True)
+
+
+def draw_bar(bar_data, x, y, loc=st):
+    fig = px.bar(
+        bar_data,
+        x=x,
+        y=y,
+        color="split",
+        color_discrete_sequence=IBM_COLORS,
+        barmode="group",
+        hover_data=bar_data.columns,
+    )
+    fig = update_axis_font(fig)
+    loc.plotly_chart(fig, use_container_width=True)
+
+
+def parse_metrics(metadata, loc):
+    for split, meta in metadata.items():
+        for key, val in meta.__dict__.items():
+            if isinstance(val, int):
+                loc.metric(label=f"{split}-{key}", value=val)
+
+
+def parse_counters(metadata):
+    meta = metadata["train"]  # using the training counter to fetch the names
+    counters = []
+    for k, v in meta.__dict__.items():
+        if "counter" in k and len(v) > 0:
+            counters.append(k)
+    return counters
+
+
+# generate the df for histogram
+def parse_label_counter(metadata, counter_type):
+    hist_data = []
+    for split, m in metadata.items():
+        metadata_counter = getattr(m, counter_type)
+        for k, v in metadata_counter.items():
+            row = {}
+            row["labels"] = k
+            row[counter_type] = v
+            row["split"] = split
+            hist_data.append(row)
+    return pd.DataFrame(hist_data)
+
+
+
+
+# load BigBioConfigHelpers
+#==================================
+logger.info("about to call load_conhelps")
+conhelps = load_conhelps()
+logger.info("exiting call load_conhelps")
+config_name_to_conhelp = {ch.config.name: ch for ch in conhelps}
+ds_display_names = sorted(list(set([ch.display_name for ch in conhelps])))
+ds_display_name_to_config_names = defaultdict(list)
+for ch in conhelps:
+    ds_display_name_to_config_names[ch.display_name].append(ch.config.name)
+
+
+# dataset selection
+#==================================
+
+st.sidebar.title("Dataset Selection")
+ds_display_name = st.sidebar.selectbox("dataset name", ds_display_names, index=0)
+
+config_names = ds_display_name_to_config_names[ds_display_name]
+config_name = st.sidebar.selectbox("config name", config_names)
+conhelp = config_name_to_conhelp[config_name]
+
+
+st.header(f"Dataset stats for {ds_display_name}")
+
+
+@st.cache()
+def load_data(conhelp):
+    metadata = conhelp.get_metadata()
+    dsd = conhelp.load_dataset()
+    dsd = dsd.map(
+        BATCH_MAPPERS_TEXT_FROM_SCHEMA[conhelp.bigbio_schema_caps.lower()],
+        batched=True)
+
+    return dsd, metadata
+
+@st.cache()
+def count_vectorize(dsd):
+    cv = CountVectorizer()
+    xcvs = {}
+    dfs_tok_per_samp = []
+    for split, ds in dsd.items():
+        xcv = cv.fit_transform(ds['text'])
+        token_counts = np.asarray(xcv.sum(axis=1)).flatten()
+        df = pd.DataFrame(token_counts, columns=["tokens per sample"])
+        df["split"] = split
+        dfs_tok_per_samp.append(df)
+        xcvs[split] = xcv
+    df_tok_per_samp = pd.concat(dfs_tok_per_samp)
+    return xcvs, df_tok_per_samp
+
+
+dsd_load_state = st.info(f"Loading {ds_display_name} - {config_name} ...")
+dsd, metadata = load_data(conhelp)
+dsd_load_state.empty()
+
+cv_load_state = st.info(f"Count Vectorizing {ds_display_name} - {config_name} ...")
+xcvs, df_tok_per_samp = count_vectorize(dsd)
+cv_load_state.empty()
+
+
+st.sidebar.subheader(f"BigBIO Schema = {conhelp.bigbio_schema_caps}")
+
+st.sidebar.subheader("Tasks Supported by Dataset")
+tasks = conhelp.tasks
+tasks = [string.capwords(task.replace("_", " ")) for task in tasks]
+st.sidebar.markdown(
+    """
+    {}
+    """.format(
+        "\n".join([
+            f"- {task}" for task in tasks
+        ]))
+)
+
+st.sidebar.subheader("Languages")
+langs = conhelp.languages
+st.sidebar.markdown(
+    """
+    {}
+    """.format("\n".join([f"- {lang}" for lang in langs]))
+)
+
+st.sidebar.subheader("Home Page")
+st.sidebar.write(conhelp.homepage)
+
+st.sidebar.subheader("Description")
+st.sidebar.write(conhelp.description)
+
+st.sidebar.subheader("Citation")
+st.sidebar.markdown(f"""\
+```
+{conhelp.citation}
+````
+"""
+                    )
+st.sidebar.subheader("Counts")
+parse_metrics(metadata, st.sidebar)
+
+
+
+# dataframe display
+#if "train" in dsd.keys():
+#    st.subheader("Sample Preview")
+#    df = pd.DataFrame.from_dict(dsd["train"])
+#    st.write(df.head(10))
+
+
+
+# draw token distribution
+st.subheader("Sample Length Distribution")
+max_xmax = int(df_tok_per_samp["tokens per sample"].max())
+xmax = st.slider("xmax", min_value=0, max_value=max_xmax, value=max_xmax)
+histnorms = ['percent', 'probability', 'density', 'probability density', None]
+histnorm = st.selectbox("histnorm", histnorms)
+draw_histogram(df_tok_per_samp, "tokens per sample", histnorm=histnorm, xmax=xmax, loc=st)
+
+
+
+st.subheader("Counter Distributions")
+counters = parse_counters(metadata)
+counter_type = st.selectbox("counter_type", counters)
+label_df = parse_label_counter(metadata, counter_type)
+label_max = int(label_df[counter_type].max() - 1)
+label_min = int(label_df[counter_type].min())
+filter_value = st.slider("minimum cutoff", label_min, label_max)
+label_df = label_df[label_df[counter_type] >= filter_value]
+# draw bar chart for counter
+draw_bar(label_df, "labels", counter_type, st)
+
+
+st.subheader("Sample Explorer")
+split = st.selectbox("split", list(dsd.keys()))
+sample_index = st.number_input(
+    "sample index",
+    min_value=0,
+    max_value=len(dsd[split])-1,
+    value=0,
+)
+
+sample = dsd[split][sample_index]
+
+
+if conhelp.bigbio_schema_caps == "KB":
+    nlp = spacy.blank("en")
+    text = sample["text"]
+    doc = nlp(text)
+    spans = []
+    for bb_ent in sample["entities"]:
+        span = doc.char_span(
+            bb_ent["offsets"][0][0],
+            bb_ent["offsets"][0][1],
+            label=bb_ent["type"],
+        )
+        spans.append(span)
+    doc.spans["sc"] = spans
+    html = displacy.render(
+        doc,
+        style="span",
+        options={
+            "colors": {
+                et: clr for et,clr in zip(
+                    metadata[split].entities_type_counter.keys(),
+                    IBM_COLORS*10
+                )
+            }
+        },
+    )
+    style = "<style>mark.entity { display: inline-block }</style>"
+    st.write(f"{style}{get_html(html)}", unsafe_allow_html=True)
+
+
+st.write(sample)