try except prediction

app.py

@@ -95,14 +95,14 @@ def run_quality_check(dataset, config, split, column, batch_size, num_examples):
     try:
         filename = get_first_parquet_filename(dataset, config, split)
     except Exception as error:
-        yield f"❌ {error}", gr.BarPlot(), pd.DataFrame(), pd.DataFrame(), pd.DataFrame(), pd.DataFrame(),
+        yield f"❌ {error}", gr.BarPlot(), pd.DataFrame(), pd.DataFrame(), pd.DataFrame(), pd.DataFrame()
         return
 
     try:
         logging.info(f"Loading hf://datasets/{dataset}@~parquet/{filename}")
         data = pl.read_parquet(f"hf://datasets/{dataset}@~parquet/{filename}", columns=[column])
     except Exception as error:
-        yield f"❌ {error}", gr.BarPlot(), pd.DataFrame(), pd.DataFrame(), pd.DataFrame(), pd.DataFrame(),
+        yield f"❌ {error}", gr.BarPlot(), pd.DataFrame(), pd.DataFrame(), pd.DataFrame(), pd.DataFrame()
         return
     logging.info("Data fetched.")
 
@@ -112,7 +112,11 @@ def run_quality_check(dataset, config, split, column, batch_size, num_examples):
     num_examples = min(len(texts), num_examples)
     for i in range(0, num_examples, batch_size):
         batch_texts = texts[i:i+batch_size]
-        batch_predictions = predict(batch_texts)
+        try:
+            batch_predictions = predict(batch_texts)
+        except Exception as error:
+            yield f"❌ {error}", gr.BarPlot(), pd.DataFrame(), pd.DataFrame(), pd.DataFrame(), pd.DataFrame()
+            return
         predictions.extend(batch_predictions)
         texts_processed.extend(batch_texts)
         yield {"quality check in progress...": i / num_examples}, *plot_and_df(texts_processed, predictions), pd.DataFrame()