fixes

code/utils/dataset_utils.py

@@ -0,0 +1,36 @@
+import os
+import pickle
+
+
+def get_dataset_path(lmdb_base_dir: str,
+                     dataset_name: str) -> str:
+    """
+    Get the path to the dataset in the lmdb directory
+    Args:
+        lmdb_base_dir: base directory of the lmdb dataset
+        dataset_name: name of the dataset
+
+    Returns:
+        path to the dataset
+    """
+
+    dataset_path = os.path.join(lmdb_base_dir, dataset_name)
+    if not os.path.exists(dataset_path):
+        raise FileNotFoundError(f"Dataset {dataset_name} not found in {lmdb_base_dir}")
+    return dataset_path
+
+
+def load_key_stats(lmdb_dir: str) -> dict:
+    """
+    Load key statistics from the LMDB database.
+    :param lmdb_dir: Directory of the LMDB database.
+    :return: Dictionary with statistics.
+    """
+    key_stats_path = f"{lmdb_dir}/key_stats.pickle"
+
+    if not os.path.exists(key_stats_path):
+        return dict()
+
+    with open(key_stats_path, "rb") as f:
+        key_stats = pickle.load(f)
+    return key_stats