update

app.py

@@ -2,7 +2,7 @@ import gradio as gr
 import polars as pl
 from gradio_huggingfacehub_search import HuggingfaceHubSearch
 import torch
-import spaces
+# import spaces
 from torch import nn
 from transformers import AutoModel, AutoTokenizer, AutoConfig
 from huggingface_hub import PyTorchModelHubMixin
@@ -31,7 +31,7 @@ model = QualityModel.from_pretrained("nvidia/quality-classifier-deberta").to(dev
 model.eval()
 
 
-@spaces.GPU
+# @spaces.GPU
 def predict(texts: list[str]):
     inputs = tokenizer(
         texts, return_tensors="pt", padding="longest", truncation=True
@@ -46,26 +46,26 @@ 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].tolist()
+    data = pl.read_parquet(f"hf://datasets/{dataset}@~parquet/{config}/train/0000.parquet", columns=[column])
+    texts = data[column].to_list()
     predictions = predict(texts[:n_samples])
-    return pd.DataFrame({"quality": predictions}).value_counts()
+    return pd.DataFrame({"quality": predictions})
 
 
 with gr.Blocks() as demo:
     gr.Markdown("# 💫 Dataset Quality Checker 💫")
-    gr_dataset_name = HuggingfaceHubSearch(
+    dataset_name = HuggingfaceHubSearch(
             label="Hub Dataset ID",
             placeholder="Search for dataset id on Huggingface",
             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",
-    )
+    # 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):
@@ -82,6 +82,8 @@ 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)
+    df = gr.DataFrame()
     gr_check_btn.click(run_quality_check, inputs=[dataset_name, text_column, n_samples], outputs=[df])
-    gr.BarPlot(df)
+    # gr.BarPlot(df)
+
+demo.launch()