update app

app.py

@@ -24,7 +24,7 @@ class QualityModel(nn.Module, PyTorchModelHubMixin):
         outputs = self.fc(dropped)
         return torch.softmax(outputs[:, 0, :], dim=1)
 
-device = "cuda"
+device = "cuda" if torch.cuda.is_available() else "cpu"
 config = AutoConfig.from_pretrained("nvidia/quality-classifier-deberta")
 tokenizer = AutoTokenizer.from_pretrained("nvidia/quality-classifier-deberta")
 model = QualityModel.from_pretrained("nvidia/quality-classifier-deberta").to(device)
@@ -44,7 +44,8 @@ def predict(texts: list[str]):
     return predicted_domains
 
 
-def run_quality_check(dataset, config, column, n_samples):
+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()
     predictions = predict(texts[:n_samples])
@@ -65,12 +66,12 @@ with gr.Blocks() as demo:
         search_type="dataset",
         value="HuggingFaceFW/fineweb",
     )
-    config_name = "default"
+    # config_name = "default"  # TODO: user input
     @gr.render(inputs=dataset_name)
     def embed(name):
         html_code = f"""
         <iframe
-          src="https://huggingface.co/datasets/{name}/embed/viewer/{config_name}/train"
+          src="https://huggingface.co/datasets/{name}/embed/viewer/default/train"
           frameborder="0"
           width="100%"
           height="700px"
@@ -82,5 +83,5 @@ with gr.Blocks() as demo:
     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, config_name, text_column, n_samples], outputs=[df])
+    gr_check_btn.click(run_quality_check, inputs=[dataset_name, text_column, n_samples], outputs=[df])
     gr.BarPlot(df)