Add hash checking of downloaded files in SPICE

torchmd · Jun 28, 2024 · 9b5da8b · 9b5da8b
1 parent cb7c50e
commit 9b5da8b
Showing 1 changed file with 9 additions and 0 deletions.
diff --git a/torchmdnet/datasets/spice.py b/torchmdnet/datasets/spice.py
@@ -56,22 +56,27 @@ class SPICE(MemmappedDataset):
         "1.1.1": {
             "url": "https://zenodo.org/record/7258940/files",
             "file": "SPICE-1.1.1.hdf5",
+            "hash": "5411e7014c6d18ff07d108c9ad820b53",
         },
         "1.1.2": {
             "url": "https://zenodo.org/record/7338495/files",
             "file": "SPICE-1.1.2.hdf5",
+            "hash": "a2b5ae2d1f72581040e1cceb20a79a33",
         },
         "1.1.3": {
             "url": "https://zenodo.org/record/7606550/files",
             "file": "SPICE-1.1.3.hdf5",
+            "hash": "be93706b3bb2b2e327b690b185905856",
         },
         "1.1.4": {
             "url": "https://zenodo.org/records/8222043/files",
             "file": "SPICE-1.1.4.hdf5",
+            "hash": "f27d4c81da0e37d6547276bf6b4ae6a1",
         },
         "2.0.1": {
             "url": "https://zenodo.org/records/10975225/files",
             "file": "SPICE-2.0.1.hdf5",
+            "hash": "bfba2224b6540e1390a579569b475510",
         },
     }
 
@@ -178,3 +183,7 @@ def sample_iter(self, mol_ids=False):
 
     def download(self):
         download_url(self.raw_url, self.raw_dir)
+        if "hash" in self.VERSIONS[self.version]:
+            with open(self.raw_paths[0], "rb") as f:
+                file_hash = hashlib.md5(f.read()).hexdigest()
+            assert file_hash == self.VERSIONS[self.version]["hash"]