try catch for zenodo api and move tqdm import

jetnet/datasets/utils.py

@@ -79,9 +79,17 @@ def _getZenodoFileURL(record_id: int, file_name: str) -> str:
 
     records_url = f"https://zenodo.org/api/records/{record_id}"
     r = requests.get(records_url).json()
-    file = next(item for item in r["files"] if item["filename"] == file_name)
-    file_url = file["links"]["download"]
-    md5 = file["checksum"]
+
+    # Zenodo API seems to be switching back and forth between these at the moment... so trying both
+    try:
+        file = next(item for item in r["files"] if item["filename"] == file_name)
+        file_url = file["links"]["download"]
+        md5 = file["checksum"]
+    except KeyError:
+        file = next(item for item in r["files"] if item["key"] == file_name)
+        file_url = file["links"]["self"]
+        md5 = file["checksum"].split("md5:")[1]
+
     return file_url, md5
 
 

jetnet/evaluation/gen_metrics.py

@@ -14,7 +14,6 @@
 from scipy.stats import iqr, wasserstein_distance
 from torch import Tensor
 from torch.utils.data import DataLoader
-from tqdm import tqdm
 
 from jetnet import utils
 from jetnet.datasets import JetNet
@@ -40,6 +39,8 @@ def _check_get_ndarray(*arrs):
 
 def _optional_tqdm(iter_obj, use_tqdm, total=None, desc=None):
     if use_tqdm:
+        from tqdm import tqdm
+
         return tqdm(iter_obj, total=total, desc=desc)
     else:
         return iter_obj