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: diff --git a/tests/test_cli.py b/tests/test_cli.py index c3d580a..b363a13 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 / "*.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"