From f10942915708a2acbf24bd0dd5c4fe1a420cde0a Mon Sep 17 00:00:00 2001
From: Lakshidaa <lakshidaa.saigiridharan@astrazeneca.com>
Date: Mon, 3 Jun 2024 15:01:23 +0200
Subject: [PATCH 1/3] Remove support to download public .hdf5 model files

---
 aizynthfinder/tools/download_public_data.py | 14 +-------------
 1 file changed, 1 insertion(+), 13 deletions(-)

diff --git a/aizynthfinder/tools/download_public_data.py b/aizynthfinder/tools/download_public_data.py
index e198012..7471b90 100644
--- 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:

From 7daa702b31811fc633da575cf79ea1904577bc7e Mon Sep 17 00:00:00 2001
From: Lakshidaa <lakshidaa.saigiridharan@astrazeneca.com>
Date: Mon, 3 Jun 2024 15:14:38 +0200
Subject: [PATCH 2/3] Fix pytest

---
 tests/test_cli.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/tests/test_cli.py b/tests/test_cli.py
index c3d580a..a5c49c6 100644
--- 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 / "*.hdf5"))
-    assert len(filenames) == 4
+    filenames = glob.glob(str(tmpdir / "*.onnx"))
+    assert len(filenames) == 3
+
+    filenames = glob.glob(str(tmpdir / "*.onnx"))
+    assert len(filenames) == 1
+
     for filename in filenames:
         with open(filename, "r") as fileobj:
             assert fileobj.read() == "abcdef"

From 917a6be963c28739d1a791c6d9e291e1e7871d41 Mon Sep 17 00:00:00 2001
From: Lakshidaa <lakshidaa.saigiridharan@astrazeneca.com>
Date: Mon, 3 Jun 2024 15:19:20 +0200
Subject: [PATCH 3/3] Fix pytest typo

---
 tests/test_cli.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/test_cli.py b/tests/test_cli.py
index a5c49c6..b363a13 100644
--- a/tests/test_cli.py
+++ b/tests/test_cli.py
@@ -422,7 +422,7 @@ def test_download_public_data(tmpdir, mocker, add_cli_arguments):
     filenames = glob.glob(str(tmpdir / "*.onnx"))
     assert len(filenames) == 3
 
-    filenames = glob.glob(str(tmpdir / "*.onnx"))
+    filenames = glob.glob(str(tmpdir / "*.hdf5"))
     assert len(filenames) == 1
 
     for filename in filenames: