update

app.py

@@ -88,10 +88,14 @@ def run_inference(
         )
         df["site"] = df["site"].apply(lambda x: site_mapper.get(x, x))
         df["cell_line"] = df["cell_line"].apply(lambda x: x.split("_")[0])
-        if not confidence:
+        if (not confidence) and "aleatoric_confidence" in df.columns:
             df.drop(
                 ["aleatoric_confidence", "epistemic_confidence"], axis=1, inplace=True
             )
+        if (not confidence) and "aleatoric_confidence" in predicted_df.columns:
+            predicted_df.drop(
+                ["aleatoric_confidence", "epistemic_confidence"], axis=1, inplace=True
+            )
         # else:
         #     df = pd.read_csv(omic_path, low_memory=False)