add samples

app.py

@@ -47,12 +47,17 @@ def predict(texts: list[str]):
 def run_quality_check(dataset, column, n_samples):
     config = "default"
     data = pl.read_parquet(f"hf://datasets/{dataset}@~parquet/{config}/train/0000.parquet", columns=[column])
-    texts = data[column].to_list()
+    texts = data[column].to_list()[:n_samples]
     predictions = predict(texts[:n_samples])
+    texts_df = pd.DataFrame({"quality": predictions, "text": texts})
     counts = pd.DataFrame({"quality": predictions}).value_counts().to_frame()
     counts.reset_index(inplace=True)
-    return gr.BarPlot(counts, x="quality", y="count")
-
+    return (
+        gr.BarPlot(counts, x="quality", y="count"),
+        texts_df[texts_df["quality"] == "Low"][:20],
+        texts_df[texts_df["quality"] == "Medium"][:20],
+        texts_df[texts_df["quality"] == "High"][:20],
+    )
 
 with gr.Blocks() as demo:
     gr.Markdown("# 💫 Dataset Quality Checker 💫")
@@ -62,12 +67,6 @@ with gr.Blocks() as demo:
             search_type="dataset",
             value="fka/awesome-chatgpt-prompts",
         )
-    # dataset_name = HuggingfaceHubSearch(
-    #     label="Hub Dataset ID",
-    #     placeholder="Search for dataset id on Huggingface",
-    #     search_type="dataset",
-    #     value="HuggingFaceFW/fineweb",
-    # )
     # config_name = "default"  # TODO: user input
     @gr.render(inputs=dataset_name)
     def embed(name):
@@ -84,7 +83,9 @@ with gr.Blocks() as demo:
     n_samples = gr.Number(label="Num first samples to run check")
     gr_check_btn = gr.Button("Check Dataset")
     plot = gr.BarPlot()
-    # df = gr.DataFrame(visible=False)
-    gr_check_btn.click(run_quality_check, inputs=[dataset_name, text_column, n_samples], outputs=[plot])
+
+    with gr.Accordion("Explore some individual examples for each class", open=False):
+        df_low, df_medium, df_high = gr.DataFrame(), gr.DataFrame(), gr.DataFrame()
+    gr_check_btn.click(run_quality_check, inputs=[dataset_name, text_column, n_samples], outputs=[plot, df_low, df_medium, df_high])
 
 demo.launch()