🧪 REAP

Towards Effective Training-Free Chemical Reasoning with Explicit Atomic Priors

Features • Quick Start • Methodology • Citation

---

📖 Overview

REAP is a novel training-free framework that enhances Large Language Models (LLMs) with chemical reasoning capabilities. By bridging the gap between molecular graph structures and language models through atom-level tokenization priors, REAP enables LLMs to effectively predict molecular properties without extensive fine-tuning.

Traditional SMILES-based approaches often struggle to capture precise structural information. REAP addresses this by leveraging interpretable atom-level representations and a three-stage retrieval-augmented generation (RAG) pipeline.

Figure: Comparison between traditional SMILES representation and our REAP framework.

---

✨ Features

🔬 Three-Stage Reasoning Pipeline Task-Aware Feature Selection: Selects relevant features via LLM dialogue. AtomDisc-Based Retrieval: Retrieves similar molecules using discrete atom tokens. Structured Inference: Performs instance-specific LLM inference with structured outputs.	🧬 AtomDisc Backbone Graph Neural Network: Encodes molecular graphs using GIN. Vector Quantization: Maps embeddings to discrete, interpretable codebook entries. Prior Knowledge: Enriches tokens with chemical properties (charges, aromaticity, etc.).
🔍 Intelligent Retrieval Atom-Level Retrieval: Uses atom tokens for precise similarity matching. Instance Patching: Handles ambiguous tokens dynamically. Context Expansion: Incorporates two-hop neighbor information for better context.	🔌 Flexible & Scalable Multi-Backend Support: Compatible with OpenAI, DeepSeek, and Nuwa. Load Balancing: Multi-key support for high-throughput experiments. Easy Extension: Modular client architecture for adding new models.

---

⚡ Quick Start

Installation

# Clone the repository
git clone https://github.com/mingxuZhang2/ChemATP
cd ChemATP

# Create conda environment
conda env create -f environment.yml
conda activate reap

Configuration

Set up your API keys. You can use environment variables or a config file.

export DEEPSEEK_API_KEY=your_api_key_here
# OR
echo "your_api_key_here" > configs/deepseek_keys.txt

Running Predictions

Run the complete REAP pipeline on a dataset (e.g., BACE):

python three_stage_reap.py \
    --dataset bace \
    --model deepseek-chat \
    --kb_path knowledge_base/atom_knowledge_base.json \
    --checkpoint atomdisc/gnn_vq_decoder_model_trained.pth \
    --output_dir results/

See run_all_reap.sh for a batch script that runs the full pipeline across tasks and models.

The pretrained AtomDisc checkpoint (atomdisc/gnn_vq_decoder_model_trained.pth) and the atom-level knowledge base (knowledge_base/atom_knowledge_base.json) are bundled in this repository.

---

🔬 Methodology

Pipeline Architecture

AtomDisc Tokenization

AtomDisc transforms molecules into a language-compatible format:

1. GNN Encoding: Extracts diverse node embeddings.
2. Vector Quantization (VQ): Discretizes embeddings into tokens.
3. Knowledge Enrichment: Annotates tokens with priors like Gasteiger charges and H-bond properties.

---

📝 Citation

If you use REAP in your research, please cite our paper:

@inproceedings{zhang-etal-2026-reap,
    title = "{REAP}: Towards Effective Training-Free Chemical Reasoning with Explicit Atomic Priors",
    author = "Zhang, Mingxu  and
      Shen, Dazhong  and
      Zhang, Qi  and
      Sun, Ying",
    editor = "Liakata, Maria  and
      Moreira, Viviane P.  and
      Zhang, Jiajun  and
      Jurgens, David",
    booktitle = "Findings of the {A}ssociation for {C}omputational {L}inguistics: {ACL} 2026",
    month = jul,
    year = "2026",
    address = "San Diego, California, United States",
    publisher = "Association for Computational Linguistics",
    url = "https://aclanthology.org/2026.findings-acl.97/",
    doi = "10.18653/v1/2026.findings-acl.97",
    pages = "2037--2062"
}

---

📄 License

This project is licensed under the MIT License.