From 9b5da8b4d9c6dec3f445b28c5cfdbecf5a599e0d Mon Sep 17 00:00:00 2001
From: RaulPPealez <raulppelaez@gmail.com>
Date: Fri, 28 Jun 2024 10:58:37 +0200
Subject: [PATCH] Add hash checking of downloaded files in SPICE

---
 torchmdnet/datasets/spice.py | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/torchmdnet/datasets/spice.py b/torchmdnet/datasets/spice.py
index 620e6ced..45a28f89 100644
--- 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"]