Add more gnn baselines

[aditya0by0]

Changes

• Removed redundant code which which adds molecule properties to final classification linear layers. Another reason was due to an error: https://wandb.ai/chebai/chebai/runs/hmh3m22n/logs?nw=nwuseraditya0by0

• Updated number of atom and bond properties to reflect recently added new tokens in SDF-based dataset support #32

[Copilot]

Pull request overview

This PR expands the set of supported GNN baselines and augmented-graph pooling variants, while refactoring graph-dataset property handling and removing molecule-level feature concatenation from (most of) the classification heads.

Changes:

• Added new baseline model configurations (e.g., GAT, GINE) and augmented pooling variants (AA / AMG pools).
• Refactored dataset property-processing/merging into new preprocessing modules and added MoleculeNet “augmented graph” dataset wrappers.
• Updated atom/bond feature dimensionalities and removed n_molecule_properties from configs/models.

Reviewed changes

Copilot reviewed 47 out of 54 changed files in this pull request and generated 6 comments.

Show a summary per file

File	Description
requirements.txt	Adds a pinned environment-style requirements set.
README.md	Minor markdown formatting fix.
pyproject.toml	Bumps minimum Python version; updates optional deps and linter tooling.
configs/model/resgated_dynamic_gni.yml	Removes n_molecule_properties from model config.
configs/model/res_aug_amgpool.yml	Removes n_molecule_properties from model config.
configs/model/res_aug_aapool.yml	Removes n_molecule_properties from model config.
configs/model/gnn_resgated_pretrain.yml	Removes n_molecule_properties from model config.
configs/model/gnn_res_gated.yml	Removes n_molecule_properties from model config; updates atom feature count.
configs/model/gat_aug_aapool.yml	Removes n_molecule_properties from model config.
configs/model/baselines/rggcn.yml	Updates baseline feature dimensionalities and removes n_molecule_properties.
configs/model/baselines/gine.yml	Adds a new GINE baseline model config.
configs/model/baselines/gat.yml	Adds a new GAT baseline model config.
configs/model/augmented/pooling/no_pooling/gat.yml	Updates augmented “no pooling” GAT config feature dimensionalities.
configs/model/augmented/pooling/amg_pool/rggcn.yml	Adds AMG-pooling ResGated config for augmented graphs.
configs/model/augmented/pooling/amg_pool/gine.yml	Adds AMG-pooling GINE config for augmented graphs.
configs/model/augmented/pooling/amg_pool/gat.yml	Adds AMG-pooling GAT config for augmented graphs.
configs/model/augmented/pooling/aa_pool/rggcn.yml	Adds AA-pooling ResGated config for augmented graphs.
configs/model/augmented/pooling/aa_pool/gine.yml	Adds AA-pooling GINE config for augmented graphs.
configs/model/augmented/pooling/aa_pool/gat.yml	Updates AA-pooling GAT class path and removes n_molecule_properties.
configs/data/chebi50_baseline.yml	Adds baseline ChEBI50 “graph properties” dataset config.
configs/data/augmented/properties/chebi50_graph_node_prop_only.yml	Adds augmented dataset property subset config (graph node only).
configs/data/augmented/properties/chebi50_fg_prop_only.yml	Adds augmented dataset property subset config (FG only).
configs/data/augmented/properties/chebi50_fg_graph_node_prop_only.yml	Adds augmented dataset property subset config (FG + graph node).
configs/data/augmented/properties/chebi50_atom_prop_only.yml	Adds augmented dataset property subset config (atom only).
configs/data/augmented/properties/chebi50_atom_graph_node_prop_only.yml	Adds augmented dataset property subset config (atom + graph node).
configs/data/augmented/properties/chebi50_atom_fg_prop_only.yml	Adds augmented dataset property subset config (atom + FG).
configs/data/augmented/final/ToxCast_final_augmented.yml	Adds final augmented ToxCast dataset config.
configs/data/augmented/final/Tox21_final_augmented.yml	Adds final augmented Tox21 dataset config.
configs/data/augmented/final/SIDER_final_augmented.yml	Adds final augmented SIDER dataset config.
configs/data/augmented/final/PCBA_final_augmented.yml	Adds final augmented PCBA dataset config.
configs/data/augmented/final/MUV_final_augmented.yml	Adds final augmented MUV dataset config.
configs/data/augmented/final/HIV_final_augmented.yml	Adds final augmented HIV dataset config.
configs/data/augmented/final/ClinTox_final_augmented.yml	Adds final augmented ClinTox dataset config.
configs/data/augmented/final/chebi50_final_augmented.yml	Adds final augmented ChEBI50 dataset config.
configs/data/augmented/final/BBBP_final_augmented.yml	Adds final augmented BBBP dataset config.
configs/data/augmented/final/BACE_final_augmented.yml	Adds final augmented BACE dataset config.
chebai_graph/preprocessing/datasets/molecule_net_classification.py	Adds MoleculeNet dataset wrappers using augmented graph readers + property merging.
chebai_graph/preprocessing/datasets/chebi.py	Refactors ChEBI dataset code to use the new shared dataset base and augmentation mixins.
chebai_graph/preprocessing/datasets/base.py	Introduces shared property-processing and “as-per-node-type” feature merging implementations.
chebai_graph/preprocessing/datasets/augmentation_base.py	Adds augmentation mixins and node-type ablation variants.
chebai_graph/preprocessing/datasets/init.py	Exposes newly added dataset variants and MoleculeNet wrappers.
chebai_graph/preprocessing/bin/AtomFunctionalGroup/indices_one_hot.txt	Adds a new functional-group token entry.
chebai_graph/models/pooling.py	Adds AA/AMG pooling wrappers for augmented graphs.
chebai_graph/models/graph.py	Updates classification heads to stop accounting for molecule features (but still has remaining molecule_attr usage).
chebai_graph/models/gin_net.py	Removes legacy GINE implementation in favor of the new architecture module.
chebai_graph/models/dynamic_gni.py	Updates imports to new architectures base layout.
chebai_graph/models/base.py	Removes legacy model base definitions (replaced by architectures/base).
chebai_graph/models/augmented.py	Rewires augmented model classes to new pooling and architectures modules; adds GINE pooling variants.
chebai_graph/models/architectures/resgated.py	Adds the new ResGated architecture implementation.
chebai_graph/models/architectures/gine.py	Adds the new GINE architecture implementation.
chebai_graph/models/architectures/gat.py	Adds the new GAT architecture implementation.
chebai_graph/models/architectures/base.py	Adds the new shared model base/wrapper definitions (without molecule_attr concatenation).
chebai_graph/models/architectures/init.py	Introduces the architectures package.
chebai_graph/models/init.py	Updates public exports to include new architectures and pooling variants.

Suppressed comments (2)

chebai_graph/models/pooling.py:129

• AMGPool.forward overrides GraphNetWrapper.forward but does not move graph_data to self.device, which can lead to device-mismatch errors. Additionally, passing an explicit dim_size to scatter_add avoids shape mismatches if any node type is missing for a graph in the batch.

        graph_data = batch["features"][0]
        assert isinstance(graph_data, GraphData)

        is_graph_node = graph_data.is_graph_node.bool()
        is_atom_node = graph_data.is_atom_node.bool()

chebai_graph/models/graph.py:197

• ResGatedAugmentedGraphPred is still concatenating graph_data.molecule_attr into graph_vector later in forward(), but molecule-level features were removed from the pipeline (and the datasets produced by the updated preprocessing code no longer attach molecule_attr). This will raise at runtime and also makes the head input-size inconsistent with the configured 2 * hidden_length concatenation.

                torch.nn.Linear(
                    self.gnn.hidden_length + (i == 0) * self.gnn.hidden_length,
                    self.gnn.hidden_length,
                )
                for i in range(n_linear_layers - 1)

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.