Merge pull request #158 from MolecularAI/remove-download-of-hdf5-models

Remove support to download public .hdf5 model files
MolecularAI · Jun 3, 2024 · c18cde1 · c18cde1
2 parents 1caf8e2 + 917a6be
commit c18cde1
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 14 deletions.
diff --git a/aizynthfinder/tools/download_public_data.py b/aizynthfinder/tools/download_public_data.py
@@ -8,10 +8,6 @@
 import tqdm
 
 FILES_TO_DOWNLOAD = {
-    "policy_model": {
-        "filename": "uspto_model.hdf5",
-        "url": "https://zenodo.org/record/7341155/files/uspto_keras_model.hdf5",
-    },
     "policy_model_onnx": {
         "filename": "uspto_model.onnx",
         "url": "https://zenodo.org/record/7797465/files/uspto_model.onnx",
@@ -20,10 +16,6 @@
         "filename": "uspto_templates.csv.gz",
         "url": "https://zenodo.org/record/7341155/files/uspto_unique_templates.csv.gz",
     },
-    "ringbreaker_model": {
-        "filename": "uspto_ringbreaker_model.hdf5",
-        "url": "https://zenodo.org/record/7341155/files/uspto_ringbreaker_keras_model.hdf5",
-    },
     "ringbreaker_model_onnx": {
         "filename": "uspto_ringbreaker_model.onnx",
         "url": "https://zenodo.org/record/7797465/files/uspto_ringbreaker_model.onnx",
@@ -36,10 +28,6 @@
         "filename": "zinc_stock.hdf5",
         "url": "https://ndownloader.figshare.com/files/23086469",
     },
-    "filter_policy": {
-        "filename": "uspto_filter_model.hdf5",
-        "url": "https://ndownloader.figshare.com/files/25584743",
-    },
     "filter_policy_onnx": {
         "filename": "uspto_filter_model.onnx",
         "url": "https://zenodo.org/record/7797465/files/uspto_filter_model.onnx",
@@ -52,7 +40,7 @@
     - {}
   ringbreaker:
     - {}
-    - {}     
+    - {}
 filter:
   uspto: {}
 stock:

diff --git a/tests/test_cli.py b/tests/test_cli.py
@@ -419,8 +419,12 @@ def test_download_public_data(tmpdir, mocker, add_cli_arguments):
 
     download_main()
 
+    filenames = glob.glob(str(tmpdir / "*.onnx"))
+    assert len(filenames) == 3
+
     filenames = glob.glob(str(tmpdir / "*.hdf5"))
-    assert len(filenames) == 4
+    assert len(filenames) == 1
+
     for filename in filenames:
         with open(filename, "r") as fileobj:
             assert fileobj.read() == "abcdef"