initial commit

.example.env

@@ -0,0 +1 @@
+TOKEN=example_token

.gitignore

@@ -0,0 +1,8 @@
+tox21/*
+predictions/*
+hp_search/trials/*
+hp_search/logs/*
+__pycache__
+.env
+pretrained/*
+grover_base.pt

LICENSE

@@ -0,0 +1,400 @@
+Attribution-NonCommercial 4.0 International
+
+=======================================================================
+
+Creative Commons Corporation ("Creative Commons") is not a law firm and
+does not provide legal services or legal advice. Distribution of
+Creative Commons public licenses does not create a lawyer-client or
+other relationship. Creative Commons makes its licenses and related
+information available on an "as-is" basis. Creative Commons gives no
+warranties regarding its licenses, any material licensed under their
+terms and conditions, or any related information. Creative Commons
+disclaims all liability for damages resulting from their use to the
+fullest extent possible.
+
+Using Creative Commons Public Licenses
+
+Creative Commons public licenses provide a standard set of terms and
+conditions that creators and other rights holders may use to share
+original works of authorship and other material subject to copyright
+and certain other rights specified in the public license below. The
+following considerations are for informational purposes only, are not
+exhaustive, and do not form part of our licenses.
+
+     Considerations for licensors: Our public licenses are
+     intended for use by those authorized to give the public
+     permission to use material in ways otherwise restricted by
+     copyright and certain other rights. Our licenses are
+     irrevocable. Licensors should read and understand the terms
+     and conditions of the license they choose before applying it.
+     Licensors should also secure all rights necessary before
+     applying our licenses so that the public can reuse the
+     material as expected. Licensors should clearly mark any
+     material not subject to the license. This includes other CC-
+     licensed material, or material used under an exception or
+     limitation to copyright. More considerations for licensors:
+    wiki.creativecommons.org/Considerations_for_licensors
+
+     Considerations for the public: By using one of our public
+     licenses, a licensor grants the public permission to use the
+     licensed material under specified terms and conditions. If
+     the licensor's permission is not necessary for any reason--for
+     example, because of any applicable exception or limitation to
+     copyright--then that use is not regulated by the license. Our
+     licenses grant only permissions under copyright and certain
+     other rights that a licensor has authority to grant. Use of
+     the licensed material may still be restricted for other
+     reasons, including because others have copyright or other
+     rights in the material. A licensor may make special requests,
+     such as asking that all changes be marked or described.
+     Although not required by our licenses, you are encouraged to
+     respect those requests where reasonable. More considerations
+     for the public:
+    wiki.creativecommons.org/Considerations_for_licensees
+=======================================================================
+Creative Commons Attribution-NonCommercial 4.0 International Public
+License
+By exercising the Licensed Rights (defined below), You accept and agree
+to be bound by the terms and conditions of this Creative Commons
+Attribution-NonCommercial 4.0 International Public License ("Public
+License"). To the extent this Public License may be interpreted as a
+contract, You are granted the Licensed Rights in consideration of Your
+acceptance of these terms and conditions, and the Licensor grants You
+such rights in consideration of benefits the Licensor receives from
+making the Licensed Material available under these terms and
+conditions.
+Section 1 -- Definitions.
+  a. Adapted Material means material subject to Copyright and Similar
+     Rights that is derived from or based upon the Licensed Material
+     and in which the Licensed Material is translated, altered,
+     arranged, transformed, or otherwise modified in a manner requiring
+     permission under the Copyright and Similar Rights held by the
+     Licensor. For purposes of this Public License, where the Licensed
+     Material is a musical work, performance, or sound recording,
+     Adapted Material is always produced where the Licensed Material is
+     synched in timed relation with a moving image.
+  b. Adapter's License means the license You apply to Your Copyright
+     and Similar Rights in Your contributions to Adapted Material in
+     accordance with the terms and conditions of this Public License.
+
+  c. Copyright and Similar Rights means copyright and/or similar rights
+     closely related to copyright including, without limitation,
+     performance, broadcast, sound recording, and Sui Generis Database
+     Rights, without regard to how the rights are labeled or
+     categorized. For purposes of this Public License, the rights
+     specified in Section 2(b)(1)-(2) are not Copyright and Similar
+     Rights.
+  d. Effective Technological Measures means those measures that, in the
+     absence of proper authority, may not be circumvented under laws
+     fulfilling obligations under Article 11 of the WIPO Copyright
+     Treaty adopted on December 20, 1996, and/or similar international
+     agreements.
+
+  e. Exceptions and Limitations means fair use, fair dealing, and/or
+     any other exception or limitation to Copyright and Similar Rights
+     that applies to Your use of the Licensed Material.
+
+  f. Licensed Material means the artistic or literary work, database,
+     or other material to which the Licensor applied this Public
+     License.
+
+  g. Licensed Rights means the rights granted to You subject to the
+     terms and conditions of this Public License, which are limited to
+     all Copyright and Similar Rights that apply to Your use of the
+     Licensed Material and that the Licensor has authority to license.
+
+  h. Licensor means the individual(s) or entity(ies) granting rights
+     under this Public License.
+
+  i. NonCommercial means not primarily intended for or directed towards
+     commercial advantage or monetary compensation. For purposes of
+     this Public License, the exchange of the Licensed Material for
+     other material subject to Copyright and Similar Rights by digital
+     file-sharing or similar means is NonCommercial provided there is
+     no payment of monetary compensation in connection with the
+     exchange.
+
+  j. Share means to provide material to the public by any means or
+     process that requires permission under the Licensed Rights, such
+     as reproduction, public display, public performance, distribution,
+     dissemination, communication, or importation, and to make material
+     available to the public including in ways that members of the
+     public may access the material from a place and at a time
+     individually chosen by them.
+
+  k. Sui Generis Database Rights means rights other than copyright
+     resulting from Directive 96/9/EC of the European Parliament and of
+     the Council of 11 March 1996 on the legal protection of databases,
+     as amended and/or succeeded, as well as other essentially
+     equivalent rights anywhere in the world.
+
+  l. You means the individual or entity exercising the Licensed Rights
+     under this Public License. Your has a corresponding meaning.
+
+
+Section 2 -- Scope.
+
+  a. License grant.
+
+       1. Subject to the terms and conditions of this Public License,
+          the Licensor hereby grants You a worldwide, royalty-free,
+          non-sublicensable, non-exclusive, irrevocable license to
+          exercise the Licensed Rights in the Licensed Material to:
+
+            a. reproduce and Share the Licensed Material, in whole or
+               in part, for NonCommercial purposes only; and
+
+            b. produce, reproduce, and Share Adapted Material for
+               NonCommercial purposes only.
+
+       2. Exceptions and Limitations. For the avoidance of doubt, where
+          Exceptions and Limitations apply to Your use, this Public
+          License does not apply, and You do not need to comply with
+          its terms and conditions.
+
+       3. Term. The term of this Public License is specified in Section
+          6(a).
+
+       4. Media and formats; technical modifications allowed. The
+          Licensor authorizes You to exercise the Licensed Rights in
+          all media and formats whether now known or hereafter created,
+          and to make technical modifications necessary to do so. The
+          Licensor waives and/or agrees not to assert any right or
+          authority to forbid You from making technical modifications
+          necessary to exercise the Licensed Rights, including
+          technical modifications necessary to circumvent Effective
+          Technological Measures. For purposes of this Public License,
+          simply making modifications authorized by this Section 2(a)
+          (4) never produces Adapted Material.
+
+       5. Downstream recipients.
+
+            a. Offer from the Licensor -- Licensed Material. Every
+               recipient of the Licensed Material automatically
+               receives an offer from the Licensor to exercise the
+               Licensed Rights under the terms and conditions of this
+               Public License.
+
+            b. No downstream restrictions. You may not offer or impose
+               any additional or different terms or conditions on, or
+               apply any Effective Technological Measures to, the
+               Licensed Material if doing so restricts exercise of the
+               Licensed Rights by any recipient of the Licensed
+               Material.
+
+       6. No endorsement. Nothing in this Public License constitutes or
+          may be construed as permission to assert or imply that You
+          are, or that Your use of the Licensed Material is, connected
+          with, or sponsored, endorsed, or granted official status by,
+          the Licensor or others designated to receive attribution as
+          provided in Section 3(a)(1)(A)(i).
+
+  b. Other rights.
+
+       1. Moral rights, such as the right of integrity, are not
+          licensed under this Public License, nor are publicity,
+          privacy, and/or other similar personality rights; however, to
+          the extent possible, the Licensor waives and/or agrees not to
+          assert any such rights held by the Licensor to the limited
+          extent necessary to allow You to exercise the Licensed
+          Rights, but not otherwise.
+
+       2. Patent and trademark rights are not licensed under this
+          Public License.
+
+       3. To the extent possible, the Licensor waives any right to
+          collect royalties from You for the exercise of the Licensed
+          Rights, whether directly or through a collecting society
+          under any voluntary or waivable statutory or compulsory
+          licensing scheme. In all other cases the Licensor expressly
+          reserves any right to collect such royalties, including when
+          the Licensed Material is used other than for NonCommercial
+          purposes.
+
+
+Section 3 -- License Conditions.
+
+Your exercise of the Licensed Rights is expressly made subject to the
+following conditions.
+
+  a. Attribution.
+
+       1. If You Share the Licensed Material (including in modified
+          form), You must:
+
+            a. retain the following if it is supplied by the Licensor
+               with the Licensed Material:
+
+                 i. identification of the creator(s) of the Licensed
+                    Material and any others designated to receive
+                    attribution, in any reasonable manner requested by
+                    the Licensor (including by pseudonym if
+                    designated);
+
+                ii. a copyright notice;
+
+               iii. a notice that refers to this Public License;
+
+                iv. a notice that refers to the disclaimer of
+                    warranties;
+
+                 v. a URI or hyperlink to the Licensed Material to the
+                    extent reasonably practicable;
+
+            b. indicate if You modified the Licensed Material and
+               retain an indication of any previous modifications; and
+
+            c. indicate the Licensed Material is licensed under this
+               Public License, and include the text of, or the URI or
+               hyperlink to, this Public License.
+
+       2. You may satisfy the conditions in Section 3(a)(1) in any
+          reasonable manner based on the medium, means, and context in
+          which You Share the Licensed Material. For example, it may be
+          reasonable to satisfy the conditions by providing a URI or
+          hyperlink to a resource that includes the required
+          information.
+
+       3. If requested by the Licensor, You must remove any of the
+          information required by Section 3(a)(1)(A) to the extent
+          reasonably practicable.
+
+       4. If You Share Adapted Material You produce, the Adapter's
+          License You apply must not prevent recipients of the Adapted
+          Material from complying with this Public License.
+
+
+Section 4 -- Sui Generis Database Rights.
+
+Where the Licensed Rights include Sui Generis Database Rights that
+apply to Your use of the Licensed Material:
+
+  a. for the avoidance of doubt, Section 2(a)(1) grants You the right
+     to extract, reuse, reproduce, and Share all or a substantial
+     portion of the contents of the database for NonCommercial purposes
+     only;
+
+  b. if You include all or a substantial portion of the database
+     contents in a database in which You have Sui Generis Database
+     Rights, then the database in which You have Sui Generis Database
+     Rights (but not its individual contents) is Adapted Material; and
+
+  c. You must comply with the conditions in Section 3(a) if You Share
+     all or a substantial portion of the contents of the database.
+
+For the avoidance of doubt, this Section 4 supplements and does not
+replace Your obligations under this Public License where the Licensed
+Rights include other Copyright and Similar Rights.
+
+
+Section 5 -- Disclaimer of Warranties and Limitation of Liability.
+
+  a. UNLESS OTHERWISE SEPARATELY UNDERTAKEN BY THE LICENSOR, TO THE
+     EXTENT POSSIBLE, THE LICENSOR OFFERS THE LICENSED MATERIAL AS-IS
+     AND AS-AVAILABLE, AND MAKES NO REPRESENTATIONS OR WARRANTIES OF
+     ANY KIND CONCERNING THE LICENSED MATERIAL, WHETHER EXPRESS,
+     IMPLIED, STATUTORY, OR OTHER. THIS INCLUDES, WITHOUT LIMITATION,
+     WARRANTIES OF TITLE, MERCHANTABILITY, FITNESS FOR A PARTICULAR
+     PURPOSE, NON-INFRINGEMENT, ABSENCE OF LATENT OR OTHER DEFECTS,
+     ACCURACY, OR THE PRESENCE OR ABSENCE OF ERRORS, WHETHER OR NOT
+     KNOWN OR DISCOVERABLE. WHERE DISCLAIMERS OF WARRANTIES ARE NOT
+     ALLOWED IN FULL OR IN PART, THIS DISCLAIMER MAY NOT APPLY TO YOU.
+
+  b. TO THE EXTENT POSSIBLE, IN NO EVENT WILL THE LICENSOR BE LIABLE
+     TO YOU ON ANY LEGAL THEORY (INCLUDING, WITHOUT LIMITATION,
+     NEGLIGENCE) OR OTHERWISE FOR ANY DIRECT, SPECIAL, INDIRECT,
+     INCIDENTAL, CONSEQUENTIAL, PUNITIVE, EXEMPLARY, OR OTHER LOSSES,
+     COSTS, EXPENSES, OR DAMAGES ARISING OUT OF THIS PUBLIC LICENSE OR
+     USE OF THE LICENSED MATERIAL, EVEN IF THE LICENSOR HAS BEEN
+     ADVISED OF THE POSSIBILITY OF SUCH LOSSES, COSTS, EXPENSES, OR
+     DAMAGES. WHERE A LIMITATION OF LIABILITY IS NOT ALLOWED IN FULL OR
+     IN PART, THIS LIMITATION MAY NOT APPLY TO YOU.
+
+  c. The disclaimer of warranties and limitation of liability provided
+     above shall be interpreted in a manner that, to the extent
+     possible, most closely approximates an absolute disclaimer and
+     waiver of all liability.
+
+
+Section 6 -- Term and Termination.
+
+  a. This Public License applies for the term of the Copyright and
+     Similar Rights licensed here. However, if You fail to comply with
+     this Public License, then Your rights under this Public License
+     terminate automatically.
+
+  b. Where Your right to use the Licensed Material has terminated under
+     Section 6(a), it reinstates:
+
+       1. automatically as of the date the violation is cured, provided
+          it is cured within 30 days of Your discovery of the
+          violation; or
+
+       2. upon express reinstatement by the Licensor.
+
+     For the avoidance of doubt, this Section 6(b) does not affect any
+     right the Licensor may have to seek remedies for Your violations
+     of this Public License.
+
+  c. For the avoidance of doubt, the Licensor may also offer the
+     Licensed Material under separate terms or conditions or stop
+     distributing the Licensed Material at any time; however, doing so
+     will not terminate this Public License.
+
+  d. Sections 1, 5, 6, 7, and 8 survive termination of this Public
+     License.
+
+
+Section 7 -- Other Terms and Conditions.
+
+  a. The Licensor shall not be bound by any additional or different
+     terms or conditions communicated by You unless expressly agreed.
+
+  b. Any arrangements, understandings, or agreements regarding the
+     Licensed Material not stated herein are separate from and
+     independent of the terms and conditions of this Public License.
+
+
+Section 8 -- Interpretation.
+
+  a. For the avoidance of doubt, this Public License does not, and
+     shall not be interpreted to, reduce, limit, restrict, or impose
+     conditions on any use of the Licensed Material that could lawfully
+     be made without permission under this Public License.
+
+  b. To the extent possible, if any provision of this Public License is
+     deemed unenforceable, it shall be automatically reformed to the
+     minimum extent necessary to make it enforceable. If the provision
+     cannot be reformed, it shall be severed from this Public License
+     without affecting the enforceability of the remaining terms and
+     conditions.
+
+  c. No term or condition of this Public License will be waived and no
+     failure to comply consented to unless expressly agreed to by the
+     Licensor.
+
+  d. Nothing in this Public License constitutes or may be interpreted
+     as a limitation upon, or waiver of, any privileges and immunities
+     that apply to the Licensor or You, including from the legal
+     processes of any jurisdiction or authority.
+
+=======================================================================
+
+Creative Commons is not a party to its public
+licenses. Notwithstanding, Creative Commons may elect to apply one of
+its public licenses to material it publishes and in those instances
+will be considered the “Licensor.” The text of the Creative Commons
+public licenses is dedicated to the public domain under the CC0 Public
+Domain Dedication. Except for the limited purpose of indicating that
+material is shared under a Creative Commons public license or as
+otherwise permitted by the Creative Commons policies published at
+creativecommons.org/policies, Creative Commons does not authorize the
+use of the trademark "Creative Commons" or any other trademark or logo
+of Creative Commons without its prior written consent including,
+without limitation, in connection with any unauthorized modifications
+to any of its public licenses or any other arrangements,
+understandings, or agreements concerning use of licensed material. For
+the avoidance of doubt, this paragraph does not form part of the
+public licenses.
+
+Creative Commons may be contacted at creativecommons.org.

LICENSE_GROVER

@@ -0,0 +1,53 @@
+MIT License
+
+Copyright (c) 2021 Tencent AI Lab.  All rights reserved.
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+
+=============================================================================
+
+Other dependencies and licenses:
+
+-----------------------------------------------------------------------------
+	The MIT License (MIT)
+		    applies to:
+            - chemprop (https://github.com/chemprop/chemprop)
+-----------------------------------------------------------------------------
+
+MIT License
+
+Copyright (c) 2020 Wengong Jin, Kyle Swanson, Kevin Yang, Regina Barzilay, Tommi Jaakkola
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

MODEL_CARD.md

@@ -0,0 +1,29 @@
+# Model card - tox21_grover_classifier
+### Model details
+- Model name: Grover Tox21 Baseline
+- Developer: Tencent AI Lab (trained by JKU Linz)
+- Paper URL: https://arxiv.org/pdf/2007.02835
+- Model type / architecture:
+    - Grover implemented using the code accompanying the paper.
+    - The pretrained grover_base model is used as provided.
+    - A multitask network is finetuned for all Tox21 targets.
+- Inference: Access via FastAPI endpoint. Upon a Tox21 prediction request, the model generates and returns predictions for all Tox21 targets simultaneously.
+- Model version: v0
+- Model date: 10.12.2025
+- Reproducibility: Code for full training is available and enables refitting of TabPFN
+from scratch.
+- Reproducibility: Code for full training is available and enables retraining of the model from scratch.
+
+### Intended use
+This model serves as a baseline benchmark for evaluating and comparing toxicity prediction
+methods across the 12 pathway assays of the Tox21 dataset. It is not intended for clinical
+decision-making without experimental validation.
+
+### Metric
+Each Tox21 task is evaluated using the area under the receiver operating characteristic curve (AUC). Overall performance is reported as the mean AUC across all individual tasks.
+
+### Training data
+Tox21 training and validation sets.
+
+### Evaluation data
+Tox21 test set.

app.py

@@ -0,0 +1,77 @@
+"""
+This is the main entry point for the FastAPI application.
+The app handles the request to predict toxicity for a list of SMILES strings.
+"""
+
+# ---------------------------------------------------------------------------------------
+# Dependencies and global variable definition
+import os
+from typing import List, Dict, Optional
+from fastapi import FastAPI, Header, HTTPException
+from pydantic import BaseModel, Field
+
+from predict import predict as predict_func
+
+API_KEY = os.getenv("API_KEY")  # set via Space Secrets
+
+
+# ---------------------------------------------------------------------------------------
+class Request(BaseModel):
+    smiles: List[str] = Field(min_items=1, max_items=1000)
+
+
+class Response(BaseModel):
+    predictions: dict
+    model_info: Dict[str, str] = {}
+
+
+app = FastAPI(title="toxicity-api")
+
+
+@app.get("/")
+def root():
+    return {
+        "message": "Toxicity Prediction API",
+        "endpoints": {
+            "/metadata": "GET - API metadata and capabilities",
+            "/healthz": "GET - Health check",
+            "/predict": "POST - Predict toxicity for SMILES",
+        },
+        "usage": "Send POST to /predict with {'smiles': ['your_smiles_here']}",
+    }
+
+
+@app.get("/metadata")
+def metadata():
+    return {
+        "name": "Tox21 GROVER Classifier",
+        "version": "0.1.0",
+        "tox_endpoints": [
+            "NR-AR",
+            "NR-AR-LBD",
+            "NR-AhR",
+            "NR-Aromatase",
+            "NR-ER",
+            "NR-ER-LBD",
+            "NR-PPAR-gamma",
+            "SR-ARE",
+            "SR-ATAD5",
+            "SR-HSE",
+            "SR-MMP",
+            "SR-p53",
+        ],
+    }
+
+
+@app.get("/healthz")
+def healthz():
+    return {"ok": True}
+
+
+@app.post("/predict", response_model=Response)
+def predict(request: Request):
+    predictions = predict_func(request.smiles)
+    return {
+        "predictions": predictions,
+        "model_info": {"name": "Tox21 GROVER classifier", "version": "0.1.0"},
+    }

config/config.json

@@ -0,0 +1,15 @@
+{
+  "batch_size": 32,
+  "init_lr": 10,
+  "max_lr": 0.0001,
+  "final_lr": 9,
+  "dropout": 0.2,
+  "attn_hidden": 128,
+  "attn_out": 8,
+  "dist_coff": 0.15,
+  "bond_drop_rate": 0.2,
+  "ffn_num_layer": 2,
+  "ffn_hidden_size": 13,
+  "real_init_lr": 1e-05,
+  "real_final_lr": 1.1111111111111112e-05
+}

environment.yaml

@@ -0,0 +1,229 @@
+name: grover
+channels:
+  - conda-forge
+  - defaults
+  - pytorch
+  - https://repo.anaconda.com/pkgs/free
+  - rmg
+  - rdkit
+dependencies:
+  - _libgcc_mutex=0.1=conda_forge
+  - _openmp_mutex=4.5=2_gnu
+  - absl-py=2.3.1=pyhd8ed1ab_0
+  - aom=3.5.0=h27087fc_0
+  - blas=1.0=mkl
+  - boost=1.74.0=py39h5472131_5
+  - boost-cpp=1.74.0=h75c5d50_8
+  - brotli-python=1.1.0=py39hf88036b_3
+  - bzip2=1.0.8=hda65f42_8
+  - c-ares=1.34.5=hb9d3cd8_0
+  - ca-certificates=2025.11.12=hbd8a1cb_0
+  - cairo=1.18.4=h44eff21_0
+  - certifi=2025.8.3=pyhd8ed1ab_0
+  - cffi=1.17.1=py39h15c3d72_0
+  - charset-normalizer=3.4.3=pyhd8ed1ab_0
+  - colorama=0.4.6=pyhd8ed1ab_1
+  - cpython=3.9.23=py39hd8ed1ab_0
+  - cuda-cudart=12.4.127=he02047a_2
+  - cuda-cudart_linux-64=12.4.127=h85509e4_2
+  - cuda-cupti=12.4.127=he02047a_2
+  - cuda-libraries=12.4.1=ha770c72_1
+  - cuda-nvrtc=12.4.127=he02047a_2
+  - cuda-nvtx=12.4.127=he02047a_2
+  - cuda-opencl=12.4.127=he02047a_1
+  - cuda-runtime=12.4.1=ha804496_0
+  - cuda-version=12.4=h3060b56_3
+  - ffmpeg=5.1.2=gpl_h8dda1f0_106
+  - filelock=3.19.1=pyhd8ed1ab_0
+  - font-ttf-dejavu-sans-mono=2.37=hab24e00_0
+  - font-ttf-inconsolata=3.000=h77eed37_0
+  - font-ttf-source-code-pro=2.038=h77eed37_0
+  - font-ttf-ubuntu=0.83=h77eed37_3
+  - fontconfig=2.15.0=h7e30c49_1
+  - fonts-conda-ecosystem=1=0
+  - fonts-conda-forge=1=hc364b38_1
+  - freetype=2.14.1=ha770c72_0
+  - gmp=6.3.0=hac33072_2
+  - gmpy2=2.2.1=py39h7196dd7_0
+  - gnutls=3.7.9=hb077bed_0
+  - grpcio=1.74.1=py39h6482d51_0
+  - h2=4.2.0=pyhd8ed1ab_0
+  - hpack=4.1.0=pyhd8ed1ab_0
+  - hyperframe=6.1.0=pyhd8ed1ab_0
+  - icu=70.1=h27087fc_0
+  - idna=3.10=pyhd8ed1ab_1
+  - importlib-metadata=8.7.0=pyhe01879c_1
+  - intel-openmp=2023.0.0=h9e868ea_25371
+  - jinja2=3.1.6=pyhd8ed1ab_0
+  - joblib=1.5.1=pyhd8ed1ab_0
+  - lame=3.100=h166bdaf_1003
+  - lcms2=2.17=h717163a_0
+  - ld_impl_linux-64=2.45=default_hbd61a6d_104
+  - lerc=4.0.0=h0aef613_1
+  - libabseil=20250512.1=cxx17_hba17884_0
+  - libblas=3.9.0=1_h86c2bf4_netlib
+  - libcblas=3.9.0=13_h8e06fc2_netlib
+  - libcublas=12.4.5.8=he02047a_2
+  - libcufft=11.2.1.3=he02047a_2
+  - libcufile=1.9.1.3=he02047a_2
+  - libcurand=10.3.5.147=he02047a_2
+  - libcusolver=11.6.1.9=he02047a_2
+  - libcusparse=12.3.1.170=he02047a_2
+  - libdeflate=1.25=h17f619e_0
+  - libdrm=2.4.125=hb03c661_1
+  - libegl=1.7.0=ha4b6fd6_2
+  - libexpat=2.7.3=hecca717_0
+  - libffi=3.4.6=h2dba641_1
+  - libfreetype=2.14.1=ha770c72_0
+  - libfreetype6=2.14.1=h73754d4_0
+  - libgcc=15.2.0=he0feb66_13
+  - libgcc-ng=15.2.0=h69a702a_13
+  - libgfortran=15.2.0=h69a702a_13
+  - libgfortran-ng=15.2.0=h69a702a_13
+  - libgfortran5=15.2.0=h68bc16d_13
+  - libgl=1.7.0=ha4b6fd6_2
+  - libglib=2.86.0=h1fed272_0
+  - libglvnd=1.7.0=ha4b6fd6_2
+  - libglx=1.7.0=ha4b6fd6_2
+  - libgomp=15.2.0=he0feb66_13
+  - libgrpc=1.74.1=hebd82d6_0
+  - libhwloc=2.9.1=hd6dc26d_0
+  - libiconv=1.18=h3b78370_2
+  - libidn2=2.3.8=hfac485b_1
+  - libjpeg-turbo=3.1.2=hb03c661_0
+  - liblapack=3.9.0=13_h8876d29_netlib
+  - liblzma=5.8.1=hb9d3cd8_2
+  - liblzma-devel=5.8.1=hb9d3cd8_2
+  - libnpp=12.2.5.30=he02047a_2
+  - libnsl=2.0.1=hb9d3cd8_1
+  - libnvfatbin=12.4.127=he02047a_2
+  - libnvjitlink=12.4.127=he02047a_2
+  - libnvjpeg=12.3.1.117=he02047a_2
+  - libopus=1.5.2=hd0c01bc_0
+  - libpciaccess=0.18=hb9d3cd8_0
+  - libpng=1.6.51=h421ea60_0
+  - libprotobuf=6.31.1=h49aed37_2
+  - libre2-11=2025.11.05=h7b12aa8_0
+  - libsqlite=3.51.0=h0c1763c_0
+  - libstdcxx=15.2.0=h934c35e_13
+  - libstdcxx-ng=15.2.0=hdf11a46_13
+  - libtasn1=4.20.0=hb03c661_1
+  - libtiff=4.7.1=h9d88235_1
+  - libunistring=0.9.10=h7f98852_0
+  - libuuid=2.41.2=he9a06e4_0
+  - libva=2.22.0=h4f16b4b_2
+  - libvpx=1.11.0=h9c3ff4c_3
+  - libwebp-base=1.6.0=hd42ef1d_0
+  - libxcb=1.17.0=h8a09558_0
+  - libxcrypt=4.4.36=hd590300_1
+  - libxml2=2.10.3=hca2bb57_4
+  - libzlib=1.3.1=hb9d3cd8_2
+  - llvm-openmp=15.0.7=h0cdce71_0
+  - markdown=3.8.2=pyhd8ed1ab_0
+  - markupsafe=3.0.2=py39h9399b63_1
+  - mkl=2023.1.0=h213fc3f_46344
+  - mpc=1.3.1=h24ddda3_1
+  - mpfr=4.2.1=h90cbb55_3
+  - mpmath=1.3.0=pyhd8ed1ab_1
+  - ncurses=6.5=h2d0b736_3
+  - nettle=3.9.1=h7ab15ed_0
+  - networkx=3.2.1=pyhd8ed1ab_0
+  - numpy=1.26.4=py39h474f0d3_0
+  - ocl-icd=2.3.3=hb9d3cd8_0
+  - opencl-headers=2025.06.13=h5888daf_0
+  - openh264=2.3.1=hcb278e6_2
+  - openjpeg=2.5.4=h55fea9a_0
+  - openssl=3.6.0=h26f9b46_0
+  - p11-kit=0.24.1=hc5aa10d_0
+  - packaging=25.0=pyh29332c3_1
+  - pandas=2.3.1=py39h1b6b32d_0
+  - pcre2=10.46=h1321c63_0
+  - pillow=11.3.0=py39h15c0740_0
+  - pip=25.2=pyh8b19718_0
+  - pixman=0.46.4=h54a6638_1
+  - protobuf=6.31.1=py39h2f5525a_0
+  - pthread-stubs=0.4=hb9d3cd8_1002
+  - pycairo=1.28.0=py39ha09e8ab_0
+  - pycparser=2.22=pyh29332c3_1
+  - pysocks=1.7.1=pyha55dd90_7
+  - python=3.9.23=hc30ae73_0_cpython
+  - python-dateutil=2.9.0.post0=pyhe01879c_2
+  - python-tzdata=2025.2=pyhd8ed1ab_0
+  - python_abi=3.9=8_cp39
+  - pytorch=2.4.0=py3.9_cuda12.4_cudnn9.1.0_0
+  - pytorch-cuda=12.4=hc786d27_7
+  - pytorch-mutex=1.0=cuda
+  - pytz=2025.2=pyhd8ed1ab_0
+  - pyyaml=6.0.2=py39h9399b63_2
+  - re2=2025.11.05=h5301d42_0
+  - readline=8.2=h8c095d6_2
+  - requests=2.32.5=pyhd8ed1ab_0
+  - scikit-learn=1.6.1=py39h4b7350c_0
+  - setuptools=80.9.0=pyhff2d567_0
+  - six=1.17.0=pyhe01879c_1
+  - svt-av1=1.4.1=hcb278e6_0
+  - sympy=1.14.0=pyh2585a3b_105
+  - tbb=2021.9.0=hf52228f_0
+  - tensorboard=2.20.0=pyhe01879c_0
+  - tensorboard-data-server=0.7.0=py39h7170ec2_2
+  - threadpoolctl=3.6.0=pyhecae5ae_0
+  - tk=8.6.13=noxft_ha0e22de_103
+  - torchtriton=3.0.0=py39
+  - torchvision=0.19.0=py39_cu124
+  - tqdm=4.67.1=pyhd8ed1ab_1
+  - typing_extensions=4.14.1=pyhe01879c_0
+  - tzdata=2025b=h78e105d_0
+  - urllib3=2.5.0=pyhd8ed1ab_0
+  - wayland=1.24.0=h3e06ad9_0
+  - wayland-protocols=1.46=hd8ed1ab_0
+  - werkzeug=3.1.3=pyhd8ed1ab_1
+  - wheel=0.45.1=pyhd8ed1ab_1
+  - x264=1!164.3095=h166bdaf_2
+  - x265=3.5=h924138e_3
+  - xorg-libx11=1.8.12=h4f16b4b_0
+  - xorg-libxau=1.0.12=hb03c661_1
+  - xorg-libxdmcp=1.1.5=hb03c661_1
+  - xorg-libxext=1.3.6=hb9d3cd8_0
+  - xorg-libxfixes=6.0.2=hb03c661_0
+  - xorg-libxrender=0.9.12=hb9d3cd8_0
+  - xz=5.8.1=hbcc6ac9_2
+  - xz-gpl-tools=5.8.1=hbcc6ac9_2
+  - xz-tools=5.8.1=hb9d3cd8_2
+  - yaml=0.2.5=h280c20c_3
+  - zipp=3.23.0=pyhd8ed1ab_0
+  - zlib=1.3.1=hb9d3cd8_2
+  - zstandard=0.23.0=py39hd399759_3
+  - zstd=1.5.7=hb8e6e7a_2
+  - pip:
+      - aiohappyeyeballs==2.6.1
+      - aiohttp==3.13.2
+      - aiosignal==1.4.0
+      - anyio==4.12.0
+      - async-timeout==5.0.1
+      - attrs==25.4.0
+      - click==8.1.8
+      - datasets==4.4.1
+      - descriptastorus==2.8.0
+      - dill==0.4.0
+      - exceptiongroup==1.3.1
+      - frozenlist==1.8.0
+      - fsspec==2025.10.0
+      - h11==0.16.0
+      - hf-xet==1.2.0
+      - httpcore==1.0.9
+      - httpx==0.28.1
+      - huggingface-hub==1.1.7
+      - multidict==6.7.0
+      - multiprocess==0.70.18
+      - pandas-flavor==0.7.0
+      - propcache==0.4.1
+      - pyarrow==21.0.0
+      - rdkit==2025.9.1
+      - rdkit-pypi==2022.9.5
+      - scipy==1.10.1
+      - shellingham==1.5.4
+      - typer-slim==0.20.0
+      - xarray==2024.7.0
+      - xxhash==3.6.0
+      - yarl==1.22.0
+prefix: /system/apps/userenv/stopf/grover

evaluate.py

@@ -0,0 +1,17 @@
+import numpy as np
+
+from src.eval import compute_roc_auc_from_csv
+from src.commands import predict_from_csv
+
+data_path = "./tox21/tox21_test_clean.csv"
+features_path = data_path.replace(".csv", ".npz")
+checkpoint_dir = "checkpoints"
+output_path = "predictions/test_set_preds_best.csv"
+
+predict_from_csv(data_path, features_path, checkpoint_dir, output_path)
+
+valid_mask = np.load("./tox21/valid_mask_test.npy")
+auc_array, mean_auc = compute_roc_auc_from_csv(output_path, "./tox21/tox21_test.csv", valid_mask)
+
+print(auc_array)
+print(mean_auc)

generate_features.py

@@ -0,0 +1,9 @@
+import os
+
+from src.commands import generate_features
+
+# paths
+TRAIN_CSV = "./tox21/tox21_train_clean.csv"
+VAL_CSV = "./tox21/tox21_validation_clean.csv"
+
+generate_features(EXAMPLES_CSV, EXAMPLES_CSV.replace('.csv', '.npz'))

grover

@@ -0,0 +1 @@
+Subproject commit 3f280d7d3419a781d303b1500c7039e37a1d87a2

hp_search.py

@@ -0,0 +1,101 @@
+import json
+import os
+import numpy as np
+from datetime import datetime
+
+from src.hp_search import generate_random_search
+from src.commands import finetune, predict_from_csv
+from src.eval import compute_roc_auc_from_csv
+
+HYPERPARAM_GRID = {
+    "batch_size": [32],
+    "init_lr": [10],
+    "max_lr": [0.001, 0.0005, 0.0001],
+    "final_lr": [2, 3, 4, 5, 6, 7, 8, 9, 10],
+    "dropout": [0.0, 0.05, 0.1, 0.2],
+    "attn_hidden": [128],
+    "attn_out": [4, 8],
+    "dist_coff": [0.05, 0.1, 0.15],
+    "bond_drop_rate": [0.0, 0.2, 0.4, 0.6],
+    "ffn_num_layer": [2, 3],
+    "ffn_hidden_size": [5, 7, 13],
+}
+
+hp_grid = generate_random_search(HYPERPARAM_GRID, num_trials=300, seed=42)
+print("Total number of configs:", len(hp_grid))
+
+# general vars
+train_path = "tox21/tox21_train_clean.csv"
+val_path = "tox21/tox21_validation_clean.csv"
+train_features_path = train_path.replace(".csv", ".npz")
+val_features_path = val_path.replace(".csv", ".npz")
+checkpoint_path = "pretrained_models/grover_base.pt" 
+
+# Tracking best model
+best_mean_auc = -1
+best_config = None
+best_model_path = None
+
+# Create directory for logs
+timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+log_dir = f"hp_search/logs/{timestamp}"
+os.makedirs(log_dir, exist_ok=True)
+overall_log_path = f"{log_dir}/hp_search_results.txt"
+best_log_path = f"{log_dir}/best_result.txt"
+
+
+# iterate over configs
+for i, args in enumerate(hp_grid):
+    save_dir = f"hp_search/trials/Trial_{i+1}"
+    
+    print("\n=========================================")
+    print("Training with config:")
+    print(args)
+    print("Save dir:", save_dir)
+    print("=========================================\n")
+
+    # finetune model
+    finetune(train_path, val_path, train_features_path, val_features_path,
+        save_dir, checkpoint_path, args)
+
+    # predict on val set
+    finetuned_model_dir = save_dir + "/fold_0/model_0"
+    output_path = save_dir + "/predictions.csv"
+    predict_from_csv(val_path, val_features_path, finetuned_model_dir, output_path)
+
+    # evaluate model
+    preds_path = save_dir + "/predictions.csv"
+    labels_path = "tox21/tox21_validation.csv"
+    valid_mask = np.load("./tox21/valid_mask_val.npy")
+    auc_per_task, mean_auc = compute_roc_auc_from_csv(preds_path, labels_path, valid_mask)
+
+    # Save all experiment results
+    with open(overall_log_path, "a") as f:
+        f.write("\n===============================\n")
+        f.write(f"Trial Num: {i+1}\n")
+        f.write(f"Mean AUC: {mean_auc}\n")
+        f.write(f"Config: {args}\n")
+        f.write(f"Save dir: {save_dir}\n")
+        f.write(f"AUC per task: {auc_per_task}\n")
+
+    # Check if best model
+    if mean_auc > best_mean_auc:
+        print("New BEST model found!")
+        best_mean_auc = mean_auc
+        best_config = args
+        best_model_path = save_dir
+
+        with open(best_log_path, "w") as f:
+            f.write("==== BEST MODEL SO FAR ====\n")
+            f.write(f"Trial Num: {i+1}\n")
+            f.write(f"Mean AUC: {best_mean_auc}\n")
+            f.write(f"Config: {best_config}\n")
+            f.write(f"Saved at: {best_model_path}\n")
+
+print("\n============================")
+print("Hyperparameter Search DONE!")
+print("Trial Num: ", i+1)
+print("Best mean AUC:", best_mean_auc)
+print("Best model saved at:", best_model_path)
+print("Best config:", best_config)
+print("============================\n")

main.py

@@ -0,0 +1,55 @@
+import random
+
+import numpy as np
+import torch
+from rdkit import RDLogger
+
+from grover.util.parsing import parse_args, get_newest_train_args
+from grover.util.utils import create_logger
+from grover.task.cross_validate import cross_validate
+from grover.task.fingerprint import generate_fingerprints
+from grover.task.predict import make_predictions, write_prediction
+from grover.task.pretrain import pretrain_model
+from grover.data.torchvocab import MolVocab
+
+
+def setup(seed):
+    # frozen random seed
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    np.random.seed(seed)
+    random.seed(seed)
+    torch.backends.cudnn.deterministic = True
+
+
+if __name__ == '__main__':
+    # setup random seed
+    setup(seed=42)
+    # Avoid the pylint warning.
+    a = MolVocab
+    # supress rdkit logger
+    lg = RDLogger.logger()
+    lg.setLevel(RDLogger.CRITICAL)
+
+    # Initialize MolVocab
+    mol_vocab = MolVocab
+
+    args = parse_args()
+    if args.parser_name == 'finetune':
+        logger = create_logger(name='train', save_dir=args.save_dir, quiet=False)
+        cross_validate(args, logger)
+    elif args.parser_name == 'pretrain':
+        logger = create_logger(name='pretrain', save_dir=args.save_dir)
+        pretrain_model(args, logger)
+    elif args.parser_name == "eval":
+        logger = create_logger(name='eval', save_dir=args.save_dir, quiet=False)
+        cross_validate(args, logger)
+    elif args.parser_name == 'fingerprint':
+        train_args = get_newest_train_args()
+        logger = create_logger(name='fingerprint', save_dir=None, quiet=False)
+        feas = generate_fingerprints(args, logger)
+        np.savez_compressed(args.output_path, fps=feas)
+    elif args.parser_name == 'predict':
+        train_args = get_newest_train_args()
+        avg_preds, test_smiles = make_predictions(args, train_args)
+        write_prediction(avg_preds, test_smiles, args)

predict.py

@@ -0,0 +1,89 @@
+import torch
+import csv
+import subprocess
+import os
+
+from src.preprocess import create_clean_smiles
+from src.commands import generate_features, predict_from_csv
+
+
+
+def predict(smiles_list):
+    """
+    Predict toxicity targets for a list of SMILES strings.
+
+    Args:
+        smiles_list (list[str]): SMILES strings
+
+    Returns:
+        dict: {smiles: {target_name: prediction_prob}}
+    """
+    data_path = "tox21/predict_smiles.csv"
+    features_path = data_path.replace(".csv", ".npz")
+    checkpoint_dir = "checkpoints"
+    output_path = "predictions/smiles_predictions.csv"
+
+    # clean smiles
+    clean_smiles, valid_mask = create_clean_smiles(smiles_list)
+
+    # Mapping from cleaned to original for valid ones
+    originals_valid = [orig for orig, ok in zip(smiles_list, valid_mask) if ok]
+
+    # sanity check (optional but nice to have)
+    if len(originals_valid) != len(clean_smiles):
+        raise ValueError(
+            f"Mismatch: {len(originals_valid)} valid originals vs {len(clean_smiles)} cleaned SMILES"
+        )
+
+    # map cleaned → original
+    cleaned_to_original = dict(zip(clean_smiles, originals_valid))
+
+    # tox21 targets
+    TARGET_NAMES = [
+            "NR-AhR","NR-AR","NR-AR-LBD","NR-Aromatase","NR-ER","NR-ER-LBD","NR-PPAR-gamma","SR-ARE","SR-ATAD5","SR-HSE","SR-MMP","SR-p53"
+        ]
+    DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Received {len(smiles_list)} SMILES strings")
+
+    # put smiles into csv
+    with open(data_path, "w", newline="") as f:
+        writer = csv.writer(f)
+        writer.writerow(["smiles"] + TARGET_NAMES)  # header
+        for smi in clean_smiles:
+            writer.writerow([smi] + [""] * len(TARGET_NAMES))
+
+    # generate features
+    generate_features(data_path, features_path)
+
+    # predict
+    predict_from_csv(data_path, features_path, checkpoint_dir, output_path)
+
+    # create results dictionary from predictions
+    predictions = {}
+    with open(output_path, "r", newline="") as f:
+        reader = csv.DictReader(f)
+        rows = list(reader)
+
+        # Identify the SMILES column even if it is unnamed
+        fieldnames = reader.fieldnames
+        smiles_col = fieldnames[0]  # first column, even if empty string
+
+        target_names = fieldnames[1:]  # all columns except first
+
+        for row in rows:
+            clean_smi = row[smiles_col]  
+            original_smi = cleaned_to_original.get(clean_smi, clean_smi)
+
+            pred_dict = {t: float(row[t]) for t in target_names}
+            predictions[original_smi] = pred_dict
+
+    # Add placeholder predictions for invalid SMILES
+    for smi, is_valid in zip(smiles_list, valid_mask):
+        if not is_valid:
+            predictions[smi] = {t: 0.5 for t in TARGET_NAMES}
+        
+
+    return predictions
+
+preds = predict(["Oc1cc(O)cc(C=Cc2ccc(O)c(O)c2)c1"])
+print(preds)

prepare_data.py

@@ -0,0 +1,20 @@
+import numpy as np
+
+from src.preprocess import clean_smiles_in_csv
+
+TOX21_TASKS = [
+    "NR-AhR","NR-AR","NR-AR-LBD","NR-Aromatase","NR-ER","NR-ER-LBD","NR-PPAR-gamma","SR-ARE","SR-ATAD5","SR-HSE","SR-MMP","SR-p53"
+]
+
+def prepare_data(data_path, save_path_clean_data, save_path_valid_mask):
+    valid_mask_train = clean_smiles_in_csv(data_path, save_path_clean_data, "smiles", TOX21_TASKS)
+    np.save(save_path_valid_mask, valid_mask_train)
+
+
+train_path = "./tox21/tox21_train.csv"
+val_path = "./tox21/tox21_validation.csv"
+
+train_path_clean = "./tox21/tox21_train_clean.csv"
+val_path_clean = "./tox21/tox21_validation_clean.csv"
+
+prepare_data(test_path, test_path_clean, "./tox21/valid_mask_test.npy")

src/commands.py

@@ -0,0 +1,54 @@
+import os
+
+def run(cmd):
+    print("\n======================================")
+    print("Running command:")
+    print(cmd)
+    print("======================================\n")
+    os.system(cmd)
+
+def generate_features(data_path, save_path):
+    run(
+        f"python scripts/save_features.py "
+        f"--data_path {data_path} "
+        f"--save_path {save_path} "
+        f"--features_generator rdkit_2d_normalized "
+        f"--restart"
+    )
+
+
+def finetune(train_path, val_path, train_features_path, val_features_path, 
+        save_dir, checkpoint_path, args
+    ):
+    finetune_cmd = (
+        f"python main.py finetune "
+        f"--data_path {train_path} "
+        f"--split_type random "
+        f"--split_sizes 1 0 0 "
+        f"--separate_val_path {val_path} "
+        f"--separate_test_path {val_path} "
+        f"--features_path {train_features_path} "
+        f"--separate_val_features_path {val_features_path} "
+        f"--separate_test_features_path {val_features_path} "
+        f"--save_dir {save_dir} "
+        f"--checkpoint_path {checkpoint_path} "
+        f"--dataset_type classification "
+        f"--num_folds 1 "
+        f"--ensemble_size 1 "
+        f"--no_features_scaling "
+        f"--ffn_hidden_size {args['ffn_hidden_size']} "
+        f"--ffn_num_layers {args['ffn_num_layer']} "
+        f"--batch_size {args['batch_size']} "
+        f"--epochs 100 "
+        f"--init_lr {args['real_init_lr']} "
+        f"--final_lr {args['real_final_lr']} "
+        f"--max_lr {args['max_lr']} "
+        f"--dropout {args['dropout']} "
+        f"--attn_hidden {args['attn_hidden']} "
+        f"--attn_out {args['attn_out']} "
+        f"--dist_coff {args['dist_coff']} "
+        f"--bond_drop_rate {args['bond_drop_rate']} "
+    )
+    run(finetune_cmd)
+
+

src/eval.py

@@ -0,0 +1,47 @@
+import pandas as pd
+import numpy as np
+from sklearn.metrics import roc_auc_score
+
+def compute_roc_auc_from_csv(preds_csv: str, labels_csv: str, valid_mask):
+    """
+    Compute ROC AUC per class and overall mean, similar to the PyTorch-style function.
+    Handles missing labels (NaN) like y_mask.
+    """
+    preds = pd.read_csv(preds_csv)
+    labels = pd.read_csv(labels_csv)
+
+    smiles_cols = [c for c in preds.columns if "smiles" in c.lower()]
+    if smiles_cols:
+        print(f"🧪 Dropping SMILES columns: {smiles_cols}")
+        preds = preds.drop(columns=smiles_cols, errors="ignore")
+        labels = labels.drop(columns=smiles_cols, errors="ignore")
+
+    shared_cols = [c for c in preds.columns if c in labels.columns]
+    preds = preds[shared_cols].apply(pd.to_numeric, errors="coerce")
+    labels = labels[shared_cols].apply(pd.to_numeric, errors="coerce")
+
+    y_pred_clean = preds.to_numpy(dtype=float)
+    y_true = labels.to_numpy(dtype=float)
+    valid_mask = valid_mask[-y_true.shape[0]:]
+    #Re-expand to original size
+    y_pred = np.full((len(valid_mask), y_pred_clean.shape[1]), 0.5, dtype=float)
+    y_pred[valid_mask] = y_pred_clean
+
+    y_mask = ~np.isnan(y_true)  
+
+    auc_list = []
+    for i in range(y_true.shape[1]):
+        mask_i = y_mask[:, i]
+        if mask_i.sum() > 0:
+            try:
+                auc = roc_auc_score(y_true[mask_i, i], y_pred[mask_i, i])
+            except ValueError:
+                auc = np.nan 
+        else:
+            auc = np.nan
+        auc_list.append(auc)
+
+    auc_array = np.array(auc_list, dtype=np.float32)
+    mean_auc = np.nanmean(auc_array)  
+
+    return auc_array, mean_auc

src/hp_search.py

@@ -0,0 +1,20 @@
+from itertools import product
+import random
+
+def generate_random_search(grid, num_trials=300, seed=42):
+    random.seed(seed)
+
+    keys = list(grid.keys())
+    values = list(grid.values())
+    all_combinations = []
+    for combo in product(*values):
+        params = dict(zip(keys, combo))
+        # Convert ratios to actual LR values
+        params["real_init_lr"]  = params["max_lr"] / params["init_lr"]
+        params["real_final_lr"] = params["max_lr"] / params["final_lr"]
+        all_combinations.append(params)
+
+    # sample num_runs out of this
+    indices = random.sample(range(len(all_combinations)), num_trials)
+    hp_subset = [all_combinations[i] for i in indices]
+    return hp_subset

src/preprocess.py

@@ -0,0 +1,88 @@
+from rdkit import Chem
+from rdkit.Chem.MolStandardize import rdMolStandardize
+from rdkit import Chem
+import numpy as np
+import pandas as pd
+from datasets import load_dataset
+from typing import List, Optional
+
+TOX21_TASKS = [
+    "NR-AhR","NR-AR","NR-AR-LBD","NR-Aromatase","NR-ER","NR-ER-LBD","NR-PPAR-gamma","SR-ARE","SR-ATAD5","SR-HSE","SR-MMP","SR-p53"
+]
+
+def create_clean_smiles(smiles_list: list[str]) -> tuple[list[str], np.ndarray]:
+    """
+    Clean and canonicalize SMILES strings while staying in SMILES space.
+    Returns (list of cleaned SMILES, mask of valid SMILES).
+    """
+    clean_smis = []
+    valid_mask = []
+
+    cleaner = rdMolStandardize.CleanupParameters()
+    tautomer_enumerator = rdMolStandardize.TautomerEnumerator()
+
+    for smi in smiles_list:
+        try:
+            mol = Chem.MolFromSmiles(smi)
+            if mol is None:
+                valid_mask.append(False)
+                continue
+
+            # Cleanup and tautomer canonicalization
+            mol = rdMolStandardize.Cleanup(mol, cleaner)
+            mol = tautomer_enumerator.Canonicalize(mol)
+
+            # -------- Charge filtering (prevents GROVER crash) --------
+            allowed_charges = {-1, 0, 1}
+            bad_charge = False
+            for atom in mol.GetAtoms():
+                if atom.GetFormalCharge() not in allowed_charges:
+                    bad_charge = True
+                    break
+
+            if bad_charge:
+                valid_mask.append(False)
+                continue
+            # ----------------------------------------------------------
+
+            # Canonical SMILES output
+            clean_smi = Chem.MolToSmiles(mol, canonical=True)
+            clean_smis.append(clean_smi)
+            valid_mask.append(True)
+
+        except Exception as e:
+            print(f"Failed to clean {smi}: {e}")
+            valid_mask.append(False)
+
+    return clean_smis, np.array(valid_mask, dtype=bool)
+
+
+def clean_smiles_in_csv(input_csv: str, output_csv: str, smiles_col: str = "smiles", target_cols: Optional[List[str]] = None):
+    """
+    Reads a CSV, cleans SMILES, and saves only valid cleaned rows with all target columns to a new CSV.
+    """
+    # Load dataset
+    df = pd.read_csv(input_csv)
+    if smiles_col not in df.columns:
+        raise ValueError(f"'{smiles_col}' column not found in CSV.")
+
+    # Infer target columns if not specified
+    if target_cols is None:
+        target_cols = [c for c in df.columns if c != smiles_col]
+    keep_cols = target_cols
+    # Validate target columns
+    missing_targets = [c for c in target_cols if c not in df.columns]
+    if missing_targets:
+        raise ValueError(f"Missing target columns in CSV: {missing_targets}")
+
+    # Clean SMILES
+    clean_smis, valid_mask = create_clean_smiles(df[smiles_col].tolist())
+
+    # Keep only valid rows
+    df_clean = df.loc[valid_mask, keep_cols].copy()
+    df_clean.insert(0, smiles_col, clean_smis)  # smiles first column
+
+    # Save cleaned dataset
+    df_clean.to_csv(output_csv, index=False)
+    print(f"✅ Cleaned dataset saved to '{output_csv}' ({len(df_clean)} valid molecules).")
+    return valid_mask

train.py

@@ -0,0 +1,42 @@
+import os
+import json
+
+from src.commands import finetune, predict_from_csv
+from src.eval import compute_roc_auc_from_csv
+
+
+def load_config(path="config.json"):
+    with open(path, "r") as f:
+        config = json.load(f)
+    return config
+
+config = load_config()
+
+print(config)
+
+
+# Paths to custom split
+train_path = "tox21/tox21_train_clean.csv"
+val_path = "tox21/tox21_validation_clean.csv"
+
+train_features_path = train_path.replace(".csv", ".npz")
+val_features_path = val_path.replace(".csv", ".npz")
+checkpoint_path = "pretrained_models/grover_base.pt" 
+
+# Output directory for finetuned model
+save_dir = "finetune/"
+
+
+finetune(train_path, val_path, train_features_path, val_features_path,
+        save_dir, checkpoint_path, args)
+
+# predict on val set
+finetuned_model_dir = save_dir + "/fold_0/model_0"
+output_path = save_dir + "/predictions.csv"
+predict_from_csv(val_path, val_features_path, finetuned_model_dir, output_path)
+
+# evaluate model
+preds_path = save_dir + "/predictions.csv"
+labels_path = "tox21/tox21_validation.csv"
+valid_mask = np.load("./tox21/valid_mask_val.npy")
+auc_per_task, mean_auc = compute_roc_auc_from_csv(preds_path, labels_path, valid_mask)