diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 0000000000000000000000000000000000000000..ba8da9c59420a22c9ff09e5d9a5f36ded8d0a4d2 --- /dev/null +++ b/.gitattributes @@ -0,0 +1,3 @@ +*.ckpt filter=lfs diff=lfs merge=lfs -text +*.png filter=lfs diff=lfs merge=lfs -text +*.pickle filter=lfs diff=lfs merge=lfs -text diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000000000000000000000000000000000000..2cfc844feee3f78e19755b0e41eb0f95317ac7b9 --- /dev/null +++ b/.gitignore @@ -0,0 +1,15 @@ +# Environment +env +.venv + +# Jupyter +.ipynb_checkpoints + +# PyCharm +.idea + +# Python bytecode +__pycache__/ +*.pyc +*.pyo +*.pyd diff --git a/LICENSE.OpenMDW-1.1 b/LICENSE.OpenMDW-1.1 new file mode 100644 index 0000000000000000000000000000000000000000..ec297ac5456384786644013ec196da33b916be97 --- /dev/null +++ b/LICENSE.OpenMDW-1.1 @@ -0,0 +1,49 @@ +OpenMDW License Agreement, version 1.1 (OpenMDW-1.1) + +By exercising rights granted to you under this agreement, you accept and agree +to its terms. + +As used in this agreement, "Model Materials" means the materials provided to +you under this agreement, consisting of: (1) one or more machine learning +models (including architecture and parameters); and (2) all related artifacts +(including associated data, documentation and software) that are provided to +you hereunder. + +Subject to your compliance with this agreement, permission is hereby granted, +free of charge, to deal in the Model Materials without restriction, including +under all copyright, patent, database, and trade secret rights included or +embodied therein. + +If you distribute any portion of the Model Materials, you shall retain in your +distribution (1) a copy of this agreement, and (2) all copyright notices and +other notices of origin included in the Model Materials that are applicable to +your distribution. + +If you file, maintain, or voluntarily participate in a lawsuit against any +person or entity asserting that the Model Materials directly or indirectly +infringe any patent or copyright, then all rights and grants made to you +hereunder are terminated, unless that lawsuit was in response to a +corresponding lawsuit first brought against you. + +This agreement does not impose any restrictions or obligations with respect to +any use, modification, or sharing of any outputs generated by using the Model +Materials. + +THE MODEL MATERIALS ARE 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, TITLE, NONINFRINGEMENT, ACCURACY, OR THE +ABSENCE OF LATENT OR OTHER DEFECTS OR ERRORS, WHETHER OR NOT DISCOVERABLE, ALL +TO THE GREATEST EXTENT PERMISSIBLE UNDER APPLICABLE LAW. + +YOU ARE SOLELY RESPONSIBLE FOR (1) CLEARING RIGHTS OF OTHER PERSONS THAT MAY +APPLY TO THE MODEL MATERIALS OR ANY USE THEREOF, INCLUDING WITHOUT LIMITATION +ANY PERSON'S COPYRIGHTS OR OTHER RIGHTS INCLUDED OR EMBODIED IN THE MODEL +MATERIALS; (2) OBTAINING ANY NECESSARY CONSENTS, PERMISSIONS OR OTHER RIGHTS +REQUIRED FOR ANY USE OF THE MODEL MATERIALS; OR (3) PERFORMING ANY DUE +DILIGENCE OR UNDERTAKING ANY OTHER INVESTIGATIONS INTO THE MODEL MATERIALS OR +ANYTHING INCORPORATED OR EMBODIED THEREIN. + +IN NO EVENT SHALL THE PROVIDERS OF THE MODEL MATERIALS 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 MODEL MATERIALS, THE +USE THEREOF OR OTHER DEALINGS THEREIN. diff --git a/README.md b/README.md new file mode 100644 index 0000000000000000000000000000000000000000..b761f62d949c7b4adb104fc7517439d197a386fa --- /dev/null +++ b/README.md @@ -0,0 +1,351 @@ +--- +language: + - en +tags: + - chemistry + - reaction-prediction + - retrosynthesis +license: openmdw-1.1 +--- + +# RXN-Sandbox + +A lightweight, local instance of RXN for running forward reaction and retrosynthesis predictions. + +> [!NOTE] +> This repository is cross-listed on **[GitHub](https://github.com/rxn4chemistry/rxn-sandbox)** and **[Hugging Face](https://huggingface.co/rxn4chemistry/rxn-sandbox)**. + +## Table of Contents +- [Overview](#overview) +- [Prerequisites](#prerequisites) +- [Quick Start](#quick-start) +- [Usage](#usage) + - [Jupyter Notebook](#jupyter-notebook) + - [MCP Integration with OpenWebUI](#mcp-integration-with-openwebui) + - [Python Scripts](#python-scripts) +- [Container Management](#container-management) + +## Overview + +This repository enables offline chemical reaction prediction using transformer models. The available tasks are RXN's most used prediction functionality: forward reaction, single-step retrosynthesis, and retrosynthesis tree generation. These tasks are performed locally via Jupyter notebook, via dedicated Python scripts, or via LLM (using OpenWeb UI and MCP). The transformer models were trained using 2025Q2 Pistachio data. + +| Forward Reaction | Retrosynthesis (Single Step) | Retrosynthesis (Tree) | +|:---:|:---:|:---:| +| ![Product Prediction](img/predict_product.svg "Product Prediction")
Predict products from reactants | ![Retrosynthesis Prediction](img/predict_retrosynthesis.svg "Retrosynthesis Prediction")
Predict reactants for one step | ![Retrosynthesis Tree Prediction](img/predict_retrosynthesis_tree.svg "Retrosynthesis Tree Prediction")
Generate multi-step routes | + + +## Prerequisites + +**Software**: + +- [Docker](https://docs.docker.com/get-docker/) and [Docker Compose](https://docs.docker.com/compose/install) OR +[Podman](https://podman.io/getting-started/installation) and [Podman Compose](https://podman-desktop.io/docs/compose/setting-up-compose) +- [Git LFS](https://git-lfs.com/) + +> [!NOTE] +> If using Podman, replace `docker` commands with `podman` throughout this guide. +> +> This repository uses Git LFS for `.ckpt` model files. Install Git LFS, then run `git lfs install` once on your machine before cloning or pulling the repository. + +**Hardware**: + +- **8GB RAM** (16GB recommended for retrosynthesis tree predictions) +- **10GB free disk space** for container images and models +- **Supported platforms:** macOS, Linux, Windows +- **Supported architectures:** Intel/AMD (x86_64) and ARM64 (Apple Silicon) +- **NVIDIA GPU** (optional) — required only when using `compose-cuda.yaml` + +## Quick Start + +Two Compose files are provided: + +| File | Description | +|------|-------------| +| [`compose.yaml`](compose.yaml) | Default — CPU-only inference | +| [`compose-cuda.yaml`](compose-cuda.yaml) | GPU-accelerated inference via NVIDIA CUDA (requires an NVIDIA GPU and the [NVIDIA Container Toolkit](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/install-guide.html)) | + +Clone this repository, and from the root directory run the following commands. Replace `compose.yaml` with `compose-cuda.yaml` in each command to enable GPU acceleration. + +#### 1. Build the Docker Images + +```bash +docker compose -f compose.yaml build +``` + +#### 2. Start the Services + +```bash +docker compose -f compose.yaml up -d +``` + +#### 3. Verify Running Containers + +You should see all containers running: + +![Containers](img/containers.png "Containers") + +#### 4. Access the Interface(s) + +- **Jupyter Notebook**: http://localhost:8888/ +- **OpenWeb UI**: http://localhost:3000/ + + + +## Usage + +
+ +Jupyter Notebook (click to expand) + +### Jupyter Notebook + +Access Jupyter at http://localhost:8888/ to use the interactive notebook environment. + +![Jupyter](img/jupyter.png "Jupyter") + +Use the provided [notebook.ipynb](http://localhost:8888/notebooks/notebook.ipynb) to explore examples and interact with the models. + +![Notebook](img/notebook.png "Notebook") + +### Initial Setup + +Run the **Celery setup and helper functions** section first to: +- Import required libraries +- Configure the Celery application +- Define helper functions for visualizing results + +### Product Prediction + +Two examples are provided for product prediction (batch and single reaction). Customize the reactants list: + +```python +# Set up a list of reactants to make predictions +reactants_list = ["CCI.O=Cc1ccc([N+](=O)[O-])c(O)c1"] +``` + +Configure prediction parameters: + +```python +# Setup task kwargs +kwargs = { + "topn": 3, # Number of results per reactant + "num_beams": 5, # Number of beams used for prediction. Must be >= topn +} +``` +After running the prediction, the results will be displayed in a table: + +![Product Prediction Results](img/forward_reactions.png "Product Prediction Results") + +### Retrosynthesis Prediction + +Retrosynthesis predictions process one product at a time. Set the target product: + +```python +# Choose product for retrosynthesis prediction +product = "C=CC(=C)C[Si](C)(C)C" +``` +Configure retrosynthesis-specific parameters: + +```python +# Setup task kwargs +kwargs = { + "topn": 10, # Number of results per reactant + "num_beams": 10, # Number of beams used for prediction. Must be >= topn + "fap": 0.6, # Forward likelihood acceptance probability (not length averaged) + "fld": 0.2, # Forward likelihood delta required between the top2 forward prediction results +} +``` + +Results are displayed in a similar table format: + +![Retrosynthesis Prediction Results](img/retro_reactions.png "Retrosynthesis Prediction Results") + +### Retrosynthesis Tree Prediction + +Start by selecting a target product SMILES: + +```python +# Choose product for retrosynthesis tree prediction +product = "C1C(C[Si](C)(C)C)=CCC2C(=O)OC(=O)C12" +``` + +Configure prediction and tree-specific parameters: + +```python +# Setup task kwargs +kwargs = { + "topn": 15, # Number of results per reactant + "num_beams": 15, # Number of beams used for prediction. Must be >= topn + "fap": 0.6, # Forward likelihood acceptance probability (not length averaged) + "fld": 0.2, # Forward likelihood delta required between the top2 forward prediction results + "max_depth": 4, # Max depth of the retrosynthesis tree + "beam_width": 6, # Max amount of nodes being expanded in each step +} + +``` + +> **⚠️ Performance Note**: Retrosynthesis tree predictions are computationally intensive and may take significant time to complete. + +### Result Visualization + +Two visualization options are available: + +**1. Text Representation** - Complete textual view of predicted routes and steps: + +![Retrosynthesis Tree Text Representation](img/retro_tree_text.png "Retrosynthesis Tree Text Representation") + +**2. Graph Representation** - Visual tree structure with molecule expansion paths. Use the selector to switch between different prediction results: + +![Retrosynthesis Tree Graph Representation](img/retro_tree_graph.png "Retrosynthesis Tree Graph Representation") + +
+ +
+LLM (click to expand) + +### MCP Integration with OpenWebUI + +Access OpenWeb UI at http://localhost:3000/ to interact with RXN models using natural language through AI assistants. + +![OpenWeb UI](img/openwebui.png "OpenWeb UI") + +#### Setup MCP Server Connection + +1. Navigate to **Settings → Integrations** +2. Under **Manage Tool Servers**, click the `+` icon +3. Set URL to `http://localhost:8000` +4. Click **Verify Connection** and **Save** + +OpenWeb UI New Integration + +#### Configure External Models (Optional) + +To use external AI models: +1. Go to **Admin Settings → Connections** +2. Add your API key for the desired model provider + +![OpenWeb UI Connections](img/openwebui_connections.png "OpenWeb UI Connections") + +#### Enable RXN Tools + +1. Select your preferred AI model +2. Click the **Integrations** button below the prompt input +3. Select **Tools** and toggle on `rxn-mcp-server` + +> **Note**: You must re-enable the tool when switching models. + +![OpenWeb UI Activate Tool](img/openwebui_tools.png "OpenWeb UI Activate Tool") + +#### Using Natural Language + +Interact with RXN functions using conversational prompts: + +![OpenWeb UI Chat Prompt](img/openwebui_prompt.png "OpenWeb UI Chat Prompt") + +
+ +
+ +Python Scripts (click to expand) + +### Running via Scripts + +If you prefer a command-line workflow, you can run predictions directly from the provided Python scripts without using Jupyter or OpenWebUI. This method may be useful for +running analyses on remote machines (which often lack GUIs). + +#### Open a Shell in the Worker Container + +The scripts are executed inside the `worker` container, where the models and Celery configuration are already available: + +```bash +docker exec -it rxn-worker-1 bash +``` + +#### Available Example Scripts + +- `python scripts/predict_product.py` — run forward reaction prediction examples +- `python scripts/predict_retrosynthesis.py` — run single-step retrosynthesis examples +- `python scripts/predict_retrosynthesis_tree.py` — run retrosynthesis tree examples +- `python scripts/run_notebook_examples.py` — run the same examples shown in the notebook in sequence (this is effectively a combination of the three prior scripts) + +#### Customize Inputs and Parameters + +Each script is intended to be edited before execution. The `scripts/` directory is mounted into the worker container, so local changes are immediately available without rebuilding the image. Update the input SMILES and prediction parameters directly in the file: + +- `reactants_list` for forward reaction prediction +- `product` for retrosynthesis and retrosynthesis tree prediction +- `topn`, `num_beams`, `fap`, `fld`, `max_depth`, and `beam_width` as needed + +Then run the script you want: + +```bash +python scripts/predict_product.py +``` + +#### Exit the Container + +When you are finished, leave the container shell with: + +```bash +exit +``` +
+ +## Container Management + +### Container Architecture + +The system consists of six containers: + +- **redis** - Results backend for Celery tasks +- **broker** - RabbitMQ message queue for task distribution +- **worker** - Celery worker running the transformer models +- **jupyter** - Interactive notebook environment +- **mcp** - Model Context Protocol server for LLM integration +- **openwebui** - Web interface for AI assistant interaction + +### Useful Commands + +```bash +# Stopping services +docker compose -f compose.yaml stop + +# Restarting services +docker compose -f compose.yaml restart + +# Viewing logs (all services) +docker compose -f compose.yaml logs -f + +# Viewing logs (specific service) +docker compose -f compose.yaml logs -f worker + +# Removing everything (containers, networks, and volumes) +docker compose -f compose.yaml down -v +``` + +> [!TIP] +> Replace `-f compose.yaml` with `-f compose-cuda.yaml` in any of the commands above to manage the GPU-accelerated stack instead. + +### Performance Notes + +- **Forward Predictions**: Fast (seconds) +- **Single-Step Retrosynthesis**: Moderate (seconds to minutes) +- **Tree Retrosynthesis**: Slow (minutes to hours depending on depth/width) + +**Optimization Tips**: + +- Start with smaller `topn` and `num_beams` values +- Limit `max_depth` to 3-4 for tree predictions +- Use `beam_width` of 5-10 for reasonable performance +- Allocate 16GB RAM for complex tree predictions + +**Benchmarking**: + +GPU acceleration provides significant speedup, especially for complex retrosynthesis tree predictions. These results are from the Python scripts in the `scripts/` folder. Analysis +times can vary substantially depending on the query molecule(s) and parameters. + +| Script | T4 GPU (AWS g4dn.xlarge) | M1 Mac CPU | GPU Speedup | +|--------|---------------------|------------|-------------| +| `predict_product.py` | **1.7 s** | 3.4 s | 2.0x | +| `predict_retrosynthesis.py` | **2.5 s** | 8.5 s | 3.4x | +| `predict_retrosynthesis_tree.py` | **87 s (1.5 min)** | 763 s (12.7 min) | 8.8x | diff --git a/compose-cuda.yaml b/compose-cuda.yaml new file mode 100644 index 0000000000000000000000000000000000000000..a50c95a9fb98e1935c29e9208c0295e47e1428c5 --- /dev/null +++ b/compose-cuda.yaml @@ -0,0 +1,71 @@ +name: rxn + +services: + + redis: + image: public.ecr.aws/docker/library/redis + command: redis-server --requirepass ubuntu + environment: + - REDIS_PASSWORD=ubuntu + + broker: + image: public.ecr.aws/docker/library/rabbitmq:3.13-management + environment: + RABBITMQ_DEFAULT_PASS: ubuntu + RABBITMQ_DEFAULT_USER: ubuntu + + worker: + build: + context: ./worker + volumes: + - ./models:/app/models + - ./vocab:/app/vocab + - ./scripts:/app/scripts + environment: + CELERY_QUEUE: product_prediction,retro_prediction + CELERY_BROKER_URL: amqp://ubuntu:ubuntu@broker:5672// + CELERY_RESULT_BACKEND: redis://:ubuntu@redis:6379/0 + depends_on: [redis, broker] + deploy: + resources: + reservations: + devices: + - driver: nvidia + count: all + capabilities: [gpu] + + jupyter: + build: + context: ./jupyter + working_dir: /rxn + volumes: + - ./jupyter:/rxn + environment: + CELERY_QUEUE: product_prediction,retro_prediction + CELERY_BROKER_URL: amqp://ubuntu:ubuntu@broker:5672// + CELERY_RESULT_BACKEND: redis://:ubuntu@redis:6379/0 + ports: + - "8888:8888" + depends_on: [redis, broker, worker] + + mcp: + build: + context: ./mcp + ports: + - "8000:8000" + environment: + CELERY_BROKER_URL: amqp://ubuntu:ubuntu@broker:5672// + CELERY_RESULT_BACKEND: redis://:ubuntu@redis:6379/0 + depends_on: [redis, broker, worker] + + openwebui: + image: ghcr.io/open-webui/open-webui + ports: + - "3000:8080" + environment: + WEBUI_AUTH: "False" + volumes: + - openwebui-data:/app/backend/data + +volumes: + openwebui-data: diff --git a/compose.yaml b/compose.yaml new file mode 100644 index 0000000000000000000000000000000000000000..5a10e28d8a0996d6fbdac47de47985c44b005d0b --- /dev/null +++ b/compose.yaml @@ -0,0 +1,64 @@ +name: rxn + +services: + + redis: + image: public.ecr.aws/docker/library/redis + command: redis-server --requirepass ubuntu + environment: + - REDIS_PASSWORD=ubuntu + + broker: + image: public.ecr.aws/docker/library/rabbitmq:3.13-management + environment: + RABBITMQ_DEFAULT_PASS: ubuntu + RABBITMQ_DEFAULT_USER: ubuntu + + worker: + build: + context: ./worker + volumes: + - ./models:/app/models + - ./vocab:/app/vocab + - ./scripts:/app/scripts + environment: + CELERY_QUEUE: product_prediction,retro_prediction + CELERY_BROKER_URL: amqp://ubuntu:ubuntu@broker:5672// + CELERY_RESULT_BACKEND: redis://:ubuntu@redis:6379/0 + depends_on: [redis, broker] + + jupyter: + build: + context: ./jupyter + working_dir: /rxn + volumes: + - ./jupyter:/rxn + environment: + CELERY_QUEUE: product_prediction,retro_prediction + CELERY_BROKER_URL: amqp://ubuntu:ubuntu@broker:5672// + CELERY_RESULT_BACKEND: redis://:ubuntu@redis:6379/0 + ports: + - "8888:8888" + depends_on: [redis, broker, worker] + + mcp: + build: + context: ./mcp + ports: + - "8000:8000" + environment: + CELERY_BROKER_URL: amqp://ubuntu:ubuntu@broker:5672// + CELERY_RESULT_BACKEND: redis://:ubuntu@redis:6379/0 + depends_on: [redis, broker, worker] + + openwebui: + image: ghcr.io/open-webui/open-webui + ports: + - "3000:8080" + environment: + WEBUI_AUTH: "False" + volumes: + - openwebui-data:/app/backend/data + +volumes: + openwebui-data: diff --git a/img/containers.png b/img/containers.png new file mode 100644 index 0000000000000000000000000000000000000000..9df77fd80e7f8bc84d5f22beeca0b017dddb6c4b --- /dev/null +++ b/img/containers.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:035462e638efa5558d8b6b2fe283cfa9c40a6e5c6a9e93db56426e09516a4970 +size 123364 diff --git a/img/forward_reactions.png b/img/forward_reactions.png new file mode 100644 index 0000000000000000000000000000000000000000..92863db8465d20d5b8b4e618dc34c3c2fdda9376 Binary files /dev/null and b/img/forward_reactions.png differ diff --git a/img/jupyter.png b/img/jupyter.png new file mode 100644 index 0000000000000000000000000000000000000000..8361ab892fabfe9d34c73f666fec2f05f6cdc0e5 Binary files /dev/null and b/img/jupyter.png differ diff --git a/img/notebook.png b/img/notebook.png new file mode 100644 index 0000000000000000000000000000000000000000..f9008c51c0dfea730d9c8d70a281c04971f11a1b --- /dev/null +++ b/img/notebook.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:cccecc90b3ac8f103e6347cf656ad7dfe036358d28a75607d82bbff85451138d +size 203214 diff --git a/img/openwebui.png b/img/openwebui.png new file mode 100644 index 0000000000000000000000000000000000000000..3cdee9a8ce223fc81276892d72dfe40432cdbb42 --- /dev/null +++ b/img/openwebui.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:041adbe138f5fd9ce266286ae1688352ee54ad2e6913e6df15dddfa77d6bc82b +size 178334 diff --git a/img/openwebui_connections.png b/img/openwebui_connections.png new file mode 100644 index 0000000000000000000000000000000000000000..427391c90a973bef9eb7541d60604ee27ee6870a --- /dev/null +++ b/img/openwebui_connections.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:40231f8d1ca79b14659431a160ca105c99cff3ce778b67dc2fe651d1571306ab +size 266659 diff --git a/img/openwebui_integrations.png b/img/openwebui_integrations.png new file mode 100644 index 0000000000000000000000000000000000000000..6236b46c488d1c49324845f9c7796a7f169c0e52 Binary files /dev/null and b/img/openwebui_integrations.png differ diff --git a/img/openwebui_prompt.png b/img/openwebui_prompt.png new file mode 100644 index 0000000000000000000000000000000000000000..837b6d51e0ea2ac01d68eb1df085e24f9db3e03d --- /dev/null +++ b/img/openwebui_prompt.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a0b4c74b0be255c4bbfa80959cea30e424fe42385f8389bf955d587db7be9008 +size 212143 diff --git a/img/openwebui_tools.png b/img/openwebui_tools.png new file mode 100644 index 0000000000000000000000000000000000000000..fc6c5b85262bad5231ad31b68c19240a48f29542 Binary files /dev/null and b/img/openwebui_tools.png differ diff --git a/img/predict_product.svg b/img/predict_product.svg new file mode 100644 index 0000000000000000000000000000000000000000..6ecb5de770b66e708cf1759665e780833ae3c2fa --- /dev/null +++ b/img/predict_product.svg @@ -0,0 +1,10 @@ + + + + + + + + + + diff --git a/img/predict_retrosynthesis.svg b/img/predict_retrosynthesis.svg new file mode 100644 index 0000000000000000000000000000000000000000..492b58230798d0ea9ede5d045556431dc8dae677 --- /dev/null +++ b/img/predict_retrosynthesis.svg @@ -0,0 +1,9 @@ + + + + + + + + + diff --git a/img/predict_retrosynthesis_tree.svg b/img/predict_retrosynthesis_tree.svg new file mode 100644 index 0000000000000000000000000000000000000000..300a46b156c971115a61c9578dd32951889b3cca --- /dev/null +++ b/img/predict_retrosynthesis_tree.svg @@ -0,0 +1,15 @@ + + + + + + + + + + + + + + + diff --git a/img/retro_reactions.png b/img/retro_reactions.png new file mode 100644 index 0000000000000000000000000000000000000000..5673b24dc9955113b61d9746d8be3bba72828359 Binary files /dev/null and b/img/retro_reactions.png differ diff --git a/img/retro_tree_graph.png b/img/retro_tree_graph.png new file mode 100644 index 0000000000000000000000000000000000000000..396e892abb56b1e3d2e0ba28f91c87f45bef0aa4 --- /dev/null +++ b/img/retro_tree_graph.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:6c13f19ec7b19f09eb51352b905432784937d6273e6e7c75e933a028b170c783 +size 208656 diff --git a/img/retro_tree_text.png b/img/retro_tree_text.png new file mode 100644 index 0000000000000000000000000000000000000000..601f163014e4978740b9c16b4c2b043fcb419c49 --- /dev/null +++ b/img/retro_tree_text.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9a3855fed5a8d1ac7ab8f35b54524bf8065ec2bb49004cac2cc4c480260d1d4a +size 166113 diff --git a/jupyter/Dockerfile b/jupyter/Dockerfile new file mode 100644 index 0000000000000000000000000000000000000000..461a3a9842cab442ffef3cfc9b8a937f72020d42 --- /dev/null +++ b/jupyter/Dockerfile @@ -0,0 +1,31 @@ +FROM public.ecr.aws/docker/library/python:3.13-slim + +# Prevents interactive prompts during apt install +ENV DEBIAN_FRONTEND=noninteractive + +# Install system dependencies required by RDKit and matplotlib +RUN apt-get update && apt-get install -y \ + build-essential \ + libxrender1 \ + libxext6 \ + libsm6 \ + libglib2.0-0 \ + libgl1 \ + ca-certificates \ + && rm -rf /var/lib/apt/lists/* + +WORKDIR /rxn + +# Copy and install Python dependencies +COPY requirements.txt . + +RUN pip install --no-cache-dir -r requirements.txt + +CMD ["jupyter", "notebook", \ + "--ip=0.0.0.0", \ + "--port=8888", \ + "--no-browser", \ + "--ServerApp.allow_root=True", \ + "--IdentityProvider.token=", \ + "--ServerApp.password=", \ + "--ServerApp.default_url=/notebook.ipynb"] \ No newline at end of file diff --git a/jupyter/notebook.ipynb b/jupyter/notebook.ipynb new file mode 100644 index 0000000000000000000000000000000000000000..87dde33e0e03528c3199943678bdd2d9b2e75b6d --- /dev/null +++ b/jupyter/notebook.ipynb @@ -0,0 +1,1567 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "6964ba26-6c99-4150-b2d2-0e2b55072993", + "metadata": {}, + "source": [ + "# RXN Sandbox" + ] + }, + { + "cell_type": "markdown", + "id": "290b4d5f-2315-469f-afcd-9bffde994f1b", + "metadata": {}, + "source": [ + "### Celery setup and helper functions" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "831d0e08-3b3d-4348-9de4-5befa7a081a6", + "metadata": {}, + "outputs": [], + "source": [ + "# All necessary imports\n", + "import time\n", + "import base64\n", + "import pandas as pd\n", + "import ipywidgets as widgets\n", + "from io import BytesIO\n", + "from rdkit import Chem\n", + "from rdkit.Chem import Draw\n", + "from celery import Celery\n", + "from celery.result import AsyncResult\n", + "from typing import Optional, Dict, Any\n", + "from IPython.display import display, HTML, clear_output, Markdown\n", + "\n", + "\n", + "# Initialize Celery\n", + "celery_app = Celery()\n", + "#print(\"Broker:\", celery_app.conf.broker_url)\n", + "#print(\"Backend:\", celery_app.conf.result_backend)\n", + "\n", + "\n", + "# Wait for the celery result\n", + "def wait_for_result(app: Celery, task_id: str, timeout: float = 120.0, poll: float = 0.5) -> Dict[str, Any]:\n", + " \"\"\"\n", + " Poll for a result with a timeout. If task updates state with meta (e.g., PROGRESS),\n", + " we surface that along the way.\n", + " \"\"\"\n", + " res = AsyncResult(task_id, app=app)\n", + " t0 = time.time()\n", + " last_state = None\n", + "\n", + " while True:\n", + " state = res.state\n", + " if state != last_state:\n", + " print(f\"State: {state} | Info: {res.info}\")\n", + " last_state = state\n", + "\n", + " if res.ready():\n", + " # could be SUCCESS or FAILURE; .get() will raise on FAILURE\n", + " return res.get(propagate=False) # returns exception object if failed\n", + "\n", + " if time.time() - t0 > timeout:\n", + " raise TimeoutError(f\"Task {task_id} did not finish in {timeout} seconds.\")\n", + " time.sleep(poll)\n", + "\n", + "\n", + "# Build results dataframe\n", + "def build_results_dataframe(records):\n", + " \"\"\"\n", + " Convert Celery result(s) into a DataFrame.\n", + " Only keeps: smiles, predicted_smiles, confidence\n", + " \"\"\"\n", + " if records is None:\n", + " return pd.DataFrame()\n", + " if isinstance(records, dict):\n", + " records = [records]\n", + " if not isinstance(records, list):\n", + " return pd.DataFrame()\n", + "\n", + " df = pd.DataFrame(records)\n", + " \n", + " if {\"reactants\", \"product\"}.issubset(df.columns):\n", + " df.insert(\n", + " 0, # put as first column (or change to another position)\n", + " \"Predicted Reaction\",\n", + " df[\"reactants\"].fillna(\"\").astype(str) + \" → \" + df[\"product\"].fillna(\"\").astype(str)\n", + " )\n", + " # Optionally remove the original columns:\n", + " df = df.drop(columns=[\"reactants\", \"product\"])\n", + "\n", + " cols = []\n", + " for c in [\"Predicted Reaction\", \"confidence\"]:\n", + " if c in df.columns:\n", + " cols.append(c)\n", + "\n", + " return df[cols]\n", + "\n", + "\n", + "# Show results table\n", + "def style_results_table(df, caption=\"Reaction Predictions\", align=\"left\"):\n", + " \"\"\"\n", + " Style the table:\n", + " - Bigger, bold caption\n", + " - Confidence mapped red→yellow→green (0→1), fixed range vmin=0, vmax=1\n", + " - Removes index\n", + " \"\"\"\n", + "\n", + " clear_output(wait=True)\n", + " \n", + " if df.empty:\n", + " display(HTML(\"

No results to display.

\"))\n", + " return\n", + "\n", + " df = df.copy()\n", + "\n", + " # Build the Styler\n", + " styler = df.style\n", + "\n", + " # Confidence formatting + red→yellow→green gradient (low→high)\n", + " # Lock range to [0, 1] so colors are consistent even for partial data.\n", + " if \"confidence\" in df.columns:\n", + " try:\n", + " styler = styler.format({\"confidence\": \"{:.3f}\"}, escape=None)\n", + " except TypeError:\n", + " styler = styler.format({\"confidence\": \"{:.3f}\"}, escape=False)\n", + " styler = styler.background_gradient(\n", + " subset=[\"confidence\"],\n", + " cmap=\"RdYlGn\", # red (low) → yellow → green (high)\n", + " vmin=0.0,\n", + " vmax=1.0,\n", + " )\n", + " styler = styler.set_properties(\n", + " subset=[\"confidence\"], **{\"font-weight\": \"500\",})\n", + "\n", + " # SMILES styling: bold + monospace + wrapping\n", + " if \"Predicted Reaction\" in df.columns:\n", + " styler = styler.set_properties(\n", + " subset=[\"Predicted Reaction\"], \n", + " **{\n", + " \"background-color\": \"#fafafa\", \n", + " \"text-align\": align,\n", + " \"font-weight\": \"600\",\n", + " \"font-family\": \"ui-monospace, SFMono-Regular, Menlo, Monaco, Consolas, 'Liberation Mono', 'Courier New', monospace\",\n", + " })\n", + "\n", + " # Table styles (caption, headers, borders)\n", + " styler = styler.set_table_styles(\n", + " [\n", + " {\n", + " \"selector\": \"caption\",\n", + " \"props\": [\n", + " (\"caption-side\", \"top\"),\n", + " (\"font-weight\", \"700\"),\n", + " (\"font-size\", \"1.2rem\"),\n", + " (\"margin-bottom\", \"10px\"),\n", + " ],\n", + " },\n", + " {\n", + " \"selector\": \"th.col_heading\",\n", + " \"props\": [\n", + " (\"font-weight\", \"700\"),\n", + " (\"padding\", \"8px\"),\n", + " (\"text-align\", \"center\"),\n", + " ],\n", + " },\n", + " {\n", + " \"selector\": \"td\",\n", + " \"props\": [\n", + " (\"padding\", \"8px\"),\n", + " (\"vertical-align\", \"center\"),\n", + " ],\n", + " },\n", + " {\n", + " \"selector\": \"table\",\n", + " \"props\": [\n", + " (\"border-collapse\", \"separate\"),\n", + " (\"border-spacing\", \"0px\"),\n", + " (\"width\", \"100%\"),\n", + " ],\n", + " },\n", + " ]\n", + " )\n", + "\n", + " # Render with inline HTML (formatters already handled escaping logic)\n", + " try:\n", + " styler = styler.hide(axis=\"index\")\n", + " except Exception:\n", + " styler = styler.hide_index()\n", + "\n", + " # Use modern API first; fall back for older pandas\n", + " try:\n", + " html = styler.set_caption(caption).to_html() # escape behavior controlled in format()\n", + " except TypeError:\n", + " html = styler.set_caption(caption).to_html(escape=False)\n", + "\n", + " display(HTML(html))\n", + "\n", + "\n", + "# Display a retrosynthesis tree textually\n", + "def display_route(route, idx):\n", + " display(Markdown(\n", + " f\"## Route {idx + 1}\\n\"\n", + " f\"- **Score:** `{route['score']:.3f}`\\n\"\n", + " f\"- **Steps:** `{route['steps']}`\"\n", + " ))\n", + "\n", + " open_nodes = route.get(\"open_nodes\", [])\n", + " if open_nodes:\n", + " display(Markdown(f\"**Remaining molecules:** `{', '.join(open_nodes)}`\"))\n", + "\n", + " for i, arc in enumerate(route.get(\"arcs\", []), start=1):\n", + " product = arc[\"product\"]\n", + " reactants = \".\".join(arc[\"reactants\"])\n", + "\n", + " display(Markdown(\n", + " f\"#### Step {i}\\n\"\n", + " f\"**Product:** `{product}` \\n\"\n", + " f\"**Reactants:** `{reactants}` \\n\"\n", + " f\"- Forward likelihood: `{arc['forward_likelihood']:.3f}` \\n\"\n", + " f\"- Arc score: `{arc['arc_score']:.3f}` \\n\"\n", + " f\"- Retro confidence: `{arc.get('confidence', 'n/a')}`\"\n", + " ))\n", + "\n", + "\n", + "\n", + "\n", + "# Convert a molecule to PNG for faster representation\n", + "def mol_png_base64(smiles, size=(250, 200)):\n", + " mol = Chem.MolFromSmiles(smiles)\n", + " if mol is None:\n", + " return \"\"\n", + "\n", + " img = Draw.MolToImage(mol, size=size)\n", + " buf = BytesIO()\n", + " img.save(buf, format=\"PNG\")\n", + " return base64.b64encode(buf.getvalue()).decode(\"utf-8\")\n", + "\n", + "\n", + "# Draw a row of reactants that make a retrosynthesis tree step\n", + "def draw_reactant_row_png(reactants, expanded_smiles=None):\n", + " imgs = [mol_png_base64(smi) for smi in reactants]\n", + "\n", + " html = \"\"\"\n", + "
\n", + " \"\"\"\n", + "\n", + " for smi, img in zip(reactants, imgs):\n", + " is_expanded = (smi == expanded_smiles)\n", + "\n", + " border = \"2px solid #1f77b4\" if is_expanded else \"1px solid #ccc\"\n", + " label = \"
expanded ↓
\" if is_expanded else \"\"\n", + "\n", + " html += f\"\"\"\n", + "
\n", + " \n", + "
\n", + " {smi}\n", + "
\n", + " {label}\n", + "
\n", + " \"\"\"\n", + "\n", + " html += \"
\"\n", + " display(HTML(html))\n", + "\n", + "\n", + "# Draw the full retrosynthesis tree\n", + "def draw_route_png(route):\n", + " display(Markdown(\"## Target molecule\"))\n", + " draw_reactant_row_png([route[\"arcs\"][0][\"product\"]])\n", + "\n", + " for i, arc in enumerate(route[\"arcs\"]):\n", + " display(Markdown(f\"## Step {i + 1} reactants\"))\n", + "\n", + " # Determine which reactant is further decomposed\n", + " expanded = None\n", + " if i + 1 < len(route[\"arcs\"]):\n", + " expanded = route[\"arcs\"][i + 1][\"product\"]\n", + "\n", + " draw_reactant_row_png(\n", + " arc[\"reactants\"],\n", + " expanded_smiles=expanded\n", + " )\n", + "\n", + "\n", + "# Builds a selector to display the full retrosynthesis tree for the selected route\n", + "def tree_route_selector(results):\n", + " options = [\n", + " (f\"Route {i+1} | score={r['score']:.3f}\", i)\n", + " for i, r in enumerate(results)\n", + " ]\n", + " \n", + " def show_route(idx):\n", + " draw_route_png(results[idx])\n", + " \n", + " dropdown = widgets.Dropdown(\n", + " options=options,\n", + " description=\"Route\",\n", + " )\n", + " \n", + " ui = widgets.interactive(show_route, idx=dropdown)\n", + " display(ui)" + ] + }, + { + "cell_type": "markdown", + "id": "66ebc007-7a61-4f2f-badb-ede0efcfbc47", + "metadata": {}, + "source": [ + "### Product prediction for a batch of reactions" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "9708b32c-41d5-4796-a4a7-16e7a740a979", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Task sent. Assigned task_id: 31bccfad-38b8-4cc5-965f-08552ed376f6\n", + "State: PENDING | Info: None\n", + "State: STARTED | Info: None\n", + "State: SUCCESS | Info: {'result': [{'reactants': 'CCI.O=Cc1ccc([N+](=O)[O-])c(O)c1', 'product': 'CCOc1cc(C=O)ccc1[N+](=O)[O-]', 'confidence': 0.9987357632971271, 'smiles': 'CCI.O=Cc1ccc([N+](=O)[O-])c(O)c1>>CCOc1cc(C=O)ccc1[N+](=O)[O-]'}, {'reactants': 'CCOc1cc(O)c(C=O)cc1OCC.OCCCBr', 'product': 'CCOc1cc(C=O)c(OCCCO)cc1OCC', 'confidence': 0.992653825419336, 'smiles': 'CCOc1cc(O)c(C=O)cc1OCC.OCCCBr>>CCOc1cc(C=O)c(OCCCO)cc1OCC'}, {'reactants': 'C=CCc1cc(OCc2ccccc2)ccc1O.CCBr', 'product': 'C=CCc1cc(OCc2ccccc2)ccc1OCC', 'confidence': 0.9966198460612675, 'smiles': 'C=CCc1cc(OCc2ccccc2)ccc1O.CCBr>>C=CCc1cc(OCc2ccccc2)ccc1OCC'}], 'time': 3.18107008934021}\n" + ] + } + ], + "source": [ + "# Set up a list of reactants to make predictions\n", + "reactants_list = [\"CCI.O=Cc1ccc([N+](=O)[O-])c(O)c1\", \"CCOc1cc(O)c(C=O)cc1OCC.OCCCBr\", \"C=CCc1cc(OCc2ccccc2)ccc1O.CCBr\"]\n", + "\n", + "# Setup task kwargs\n", + "kwargs = {\n", + " \"topn\": 1, # Number of results per reactant\n", + " \"num_beams\": 3, # Number of beams used for prediction. Must be >= topn\n", + " \"device\": None, # Device used for predicting, either \"cuda\" or \"cpu\", None defaults to cuda if available\n", + " \"ckpt_forward\": \"Pistachio2025Q2-Forward\", # Default forward model\n", + " \"vocab\": \"Pistachio2025Q2\", # Vocab for default forward model\n", + " # \"ckpt_forward_path\": \"models/forward/Pistachio2025Q2-Forward.ckpt\", # Can be used instead of ckpt_forward\n", + " # \"vocab_path\": \"vocab/Pistachio2025Q2.txt\", # Can be used instead of vocab\n", + "}\n", + "\n", + "# Send the product_prediction task with the reaction list and kwargs\n", + "task = celery_app.send_task(\n", + " \"tasks.product_prediction\",\n", + " [reactants_list],\n", + " kwargs=kwargs,\n", + " queue=\"product_prediction\",\n", + ")\n", + "print(\"Task sent. Assigned task_id: {}\".format(task.id))\n", + "\n", + "# Use the task id to get the result. Increase timeout if needed.\n", + "response = wait_for_result(celery_app, task.id, timeout=180)" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "df6007a1-1e6d-4b71-9991-8d6a21febdf0", + "metadata": { + "scrolled": true + }, + "outputs": [ + { + "data": { + "text/html": [ + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
Product Prediction Results
Predicted Reactionconfidence
CCI.O=Cc1ccc([N+](=O)[O-])c(O)c1 → CCOc1cc(C=O)ccc1[N+](=O)[O-]0.999
CCOc1cc(O)c(C=O)cc1OCC.OCCCBr → CCOc1cc(C=O)c(OCCCO)cc1OCC0.993
C=CCc1cc(OCc2ccccc2)ccc1O.CCBr → C=CCc1cc(OCc2ccccc2)ccc1OCC0.997
\n" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + } + ], + "source": [ + "# Build a dataframe with the result in the response\n", + "df = build_results_dataframe(response[\"result\"])\n", + "\n", + "# Create and show the data in a table\n", + "style_results_table(df, caption=\"Product Prediction Results\", align=\"center\")" + ] + }, + { + "cell_type": "markdown", + "id": "71f84c47-07d2-405f-b824-86d95dc9ea2a", + "metadata": {}, + "source": [ + "### Product prediction for a single reaction - top 3 results" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "8c8fb6b9-2a2b-42b6-9fb8-e5356377a7c4", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Task sent. Assigned task_id: 10bdf9f8-e49f-4c5f-8f11-5d03f3f75d83\n", + "State: PENDING | Info: None\n", + "State: STARTED | Info: None\n", + "State: SUCCESS | Info: {'result': [{'reactants': 'CCI.O=Cc1ccc([N+](=O)[O-])c(O)c1', 'product': 'CCOc1cc(C=O)ccc1[N+](=O)[O-]', 'confidence': 0.9987357632971271, 'smiles': 'CCI.O=Cc1ccc([N+](=O)[O-])c(O)c1>>CCOc1cc(C=O)ccc1[N+](=O)[O-]'}, {'reactants': 'CCI.O=Cc1ccc([N+](=O)[O-])c(O)c1', 'product': 'CCOc1cc(C=O)ccc1[N+](=O)[O-].CCOc1cc(C=O)ccc1[N+](=O)[O-]', 'confidence': 0.8466943868002893, 'smiles': 'CCI.O=Cc1ccc([N+](=O)[O-])c(O)c1>>CCOc1cc(C=O)ccc1[N+](=O)[O-].CCOc1cc(C=O)ccc1[N+](=O)[O-]'}, {'reactants': 'CCI.O=Cc1ccc([N+](=O)[O-])c(O)c1', 'product': 'CCOc1cc(C=O)ccc1[N+](=O)[O-].CCOc1cc(CO)ccc1[N+](=O)[O-]', 'confidence': 0.7820610951050019, 'smiles': 'CCI.O=Cc1ccc([N+](=O)[O-])c(O)c1>>CCOc1cc(C=O)ccc1[N+](=O)[O-].CCOc1cc(CO)ccc1[N+](=O)[O-]'}], 'time': 3.020906448364258}\n" + ] + } + ], + "source": [ + "# Set up a list of reactants to make predictions\n", + "reactants_list = [\"CCI.O=Cc1ccc([N+](=O)[O-])c(O)c1\"]\n", + "\n", + "# Setup task kwargs\n", + "kwargs = {\n", + " \"topn\": 3, # Number of results per reactant\n", + " \"num_beams\": 5, # Number of beams used for prediction. Must be >= topn\n", + " \"device\": None, # Device used for predicting, either \"cuda\" or \"cpu\", None defaults to cuda if available\n", + " \"ckpt_forward\": \"Pistachio2025Q2-Forward\", # Default forward model\n", + " \"vocab\": \"Pistachio2025Q2\", # Vocab for default forward model\n", + " # \"ckpt_forward_path\": \"models/forward/Pistachio2025Q2-Forward.ckpt\", # Can be used instead of ckpt_forward\n", + " # \"vocab_path\": \"vocab/Pistachio2025Q2.txt\", # Can be used instead of vocab\n", + "}\n", + "\n", + "# Send the product_prediction task with the reaction list and kwargs\n", + "task = celery_app.send_task(\n", + " \"tasks.product_prediction\",\n", + " [reactants_list],\n", + " kwargs=kwargs,\n", + " queue=\"product_prediction\",\n", + ")\n", + "print(\"Task sent. Assigned task_id: {}\".format(task.id))\n", + "\n", + "# Use the task id to get the result. Increase timeout if needed.\n", + "response = wait_for_result(celery_app, task.id, timeout=180)" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "2872dd07-110a-4cdb-a905-588e31bdae49", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
Product Prediction Results
Predicted Reactionconfidence
CCI.O=Cc1ccc([N+](=O)[O-])c(O)c1 → CCOc1cc(C=O)ccc1[N+](=O)[O-]0.999
CCI.O=Cc1ccc([N+](=O)[O-])c(O)c1 → CCOc1cc(C=O)ccc1[N+](=O)[O-].CCOc1cc(C=O)ccc1[N+](=O)[O-]0.847
CCI.O=Cc1ccc([N+](=O)[O-])c(O)c1 → CCOc1cc(C=O)ccc1[N+](=O)[O-].CCOc1cc(CO)ccc1[N+](=O)[O-]0.782
\n" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + } + ], + "source": [ + "# Build a dataframe with the result in the response\n", + "df = build_results_dataframe(response[\"result\"])\n", + "\n", + "# Create and show the data in a table\n", + "style_results_table(df, caption=\"Product Prediction Results\")" + ] + }, + { + "cell_type": "markdown", + "id": "ee465c64-f56c-4295-8f29-5adf23b494c1", + "metadata": {}, + "source": [ + "### Retrosynthesis prediction" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "id": "c46f70ed-2b19-43cb-b5d9-ae846cdc99c7", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Task sent. Assigned task_id: 2ac795d1-67f2-42ab-b6e8-db9d0f04e37d\n", + "State: PENDING | Info: None\n", + "State: STARTED | Info: None\n", + "State: SUCCESS | Info: {'result': [{'reactants': 'C1CCOC1.C=CC(=C)Cl.C[Si](C)(C)[CH2][Mg][Cl]', 'product': 'C=CC(=C)C[Si](C)(C)C', 'confidence': 0.8633497322540727, 'smiles': 'C1CCOC1.C=CC(=C)Cl.C[Si](C)(C)[CH2][Mg][Cl]>>C=CC(=C)C[Si](C)(C)C'}, {'reactants': 'C1CCOC1.C=CC(=C)Cl.C[Si](C)(C)[CH2][Mg][Cl].N#N', 'product': 'C=CC(=C)C[Si](C)(C)C', 'confidence': 0.853555662500899, 'smiles': 'C1CCOC1.C=CC(=C)Cl.C[Si](C)(C)[CH2][Mg][Cl].N#N>>C=CC(=C)C[Si](C)(C)C'}, {'reactants': 'C1CCOC1.C=CC(=C)Cl.CCOCC.C[Si](C)(C)[CH2][Mg][Cl]', 'product': 'C=CC(=C)C[Si](C)(C)C', 'confidence': 0.8437746350333076, 'smiles': 'C1CCOC1.C=CC(=C)Cl.CCOCC.C[Si](C)(C)[CH2][Mg][Cl]>>C=CC(=C)C[Si](C)(C)C'}, {'reactants': 'C1CCOC1.C=CC(=C)Cl.C[Si](C)(C)[CH2][Mg][Cl].O.[Cl-].[NH4+]', 'product': 'C=CC(=C)C[Si](C)(C)C', 'confidence': 0.8341045215463807, 'smiles': 'C1CCOC1.C=CC(=C)Cl.C[Si](C)(C)[CH2][Mg][Cl].O.[Cl-].[NH4+]>>C=CC(=C)C[Si](C)(C)C'}, {'reactants': 'C1CCOC1.C=CC(=C)Cl.CCOCC.C[Si](C)(C)[CH2][Mg][Cl].N#N', 'product': 'C=CC(=C)C[Si](C)(C)C', 'confidence': 0.833862536714611, 'smiles': 'C1CCOC1.C=CC(=C)Cl.CCOCC.C[Si](C)(C)[CH2][Mg][Cl].N#N>>C=CC(=C)C[Si](C)(C)C'}, {'reactants': 'C1CCOC1.C=CC(=C)Cl.CCOCC.C[Si](C)(C)[CH2][Mg][Cl].O.[Cl-].[NH4+]', 'product': 'C=CC(=C)C[Si](C)(C)C', 'confidence': 0.8328274628653165, 'smiles': 'C1CCOC1.C=CC(=C)Cl.CCOCC.C[Si](C)(C)[CH2][Mg][Cl].O.[Cl-].[NH4+]>>C=CC(=C)C[Si](C)(C)C'}, {'reactants': 'C1CCOC1.C=CC(=C)Cl.CCOCC.C[Si](C)(C)[CH2][Mg][Cl].Cl.O', 'product': 'C=CC(=C)C[Si](C)(C)C', 'confidence': 0.8303334989016697, 'smiles': 'C1CCOC1.C=CC(=C)Cl.CCOCC.C[Si](C)(C)[CH2][Mg][Cl].Cl.O>>C=CC(=C)C[Si](C)(C)C'}, {'reactants': 'C1CCOC1.C=CC(=C)Cl.CCOCC.C[Si](C)(C)[CH2][Mg][Cl].O.[Cl-].[Mg].[NH4+]', 'product': 'C=CC(=C)C[Si](C)(C)C', 'confidence': 0.8301065237221663, 'smiles': 'C1CCOC1.C=CC(=C)Cl.CCOCC.C[Si](C)(C)[CH2][Mg][Cl].O.[Cl-].[Mg].[NH4+]>>C=CC(=C)C[Si](C)(C)C'}], 'time': 14.295113801956177}\n" + ] + } + ], + "source": [ + "# Choose product for retrosynthesis prediction\n", + "product = \"C=CC(=C)C[Si](C)(C)C\"\n", + "\n", + "# Setup task kwargs\n", + "kwargs = {\n", + " \"topn\": 15, # Number of results per reactant\n", + " \"num_beams\": 15, # Number of beams used for prediction. Must be >= topn\n", + " \"fap\": 0.6, # Forward likelihood acceptance probability (not length averaged)\n", + " \"fld\": 0.2, # Forward likelihood delta required between the top2 forward prediction results\n", + " \"device\": None, # Device used for predicting, either \"cuda\" or \"cpu\", None defaults to cuda if available\n", + " \"ckpt_forward\": \"Pistachio2025Q2-Forward\", # Default forward model\n", + " \"ckpt_retro\": \"Pistachio2025Q2-Retro\", # Default retrosynthesis model\n", + " \"vocab\": \"Pistachio2025Q2\", # Vocab for default forward and retrosynthesis models\n", + " # \"ckpt_forward_path\": \"models/forward/Pistachio2025Q2-Forward.ckpt\", # Can be used instead of ckpt_forward\n", + " # \"ckpt_retro_path\": \"models/retrosynthesis/Pistachio2025Q2-Retro.ckpt\", # Can be used instead of ckpt_retro\n", + " # \"vocab_path\": \"vocab/Pistachio2025Q2.txt\", # Can be used instead of vocab\n", + "}\n", + "\n", + "# Send the retro_prediction task with the product and kwargs\n", + "task = celery_app.send_task(\n", + " \"tasks.retro_prediction\",\n", + " [product],\n", + " kwargs=kwargs,\n", + " queue=\"retro_prediction\",\n", + ")\n", + "print(\"Task sent. Assigned task_id: {}\".format(task.id))\n", + "\n", + "# Use the task id to get the result. Increase timeout if needed.\n", + "response = wait_for_result(celery_app, task.id, timeout=300)" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "id": "f7601ac3-2a27-43eb-956a-209c8dc3d4cc", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
Retrosynthesis Prediction Results
Predicted Reactionconfidence
C1CCOC1.C=CC(=C)Cl.C[Si](C)(C)[CH2][Mg][Cl] → C=CC(=C)C[Si](C)(C)C0.863
C1CCOC1.C=CC(=C)Cl.C[Si](C)(C)[CH2][Mg][Cl].N#N → C=CC(=C)C[Si](C)(C)C0.854
C1CCOC1.C=CC(=C)Cl.CCOCC.C[Si](C)(C)[CH2][Mg][Cl] → C=CC(=C)C[Si](C)(C)C0.844
C1CCOC1.C=CC(=C)Cl.C[Si](C)(C)[CH2][Mg][Cl].O.[Cl-].[NH4+] → C=CC(=C)C[Si](C)(C)C0.834
C1CCOC1.C=CC(=C)Cl.CCOCC.C[Si](C)(C)[CH2][Mg][Cl].N#N → C=CC(=C)C[Si](C)(C)C0.834
C1CCOC1.C=CC(=C)Cl.CCOCC.C[Si](C)(C)[CH2][Mg][Cl].O.[Cl-].[NH4+] → C=CC(=C)C[Si](C)(C)C0.833
C1CCOC1.C=CC(=C)Cl.CCOCC.C[Si](C)(C)[CH2][Mg][Cl].Cl.O → C=CC(=C)C[Si](C)(C)C0.830
C1CCOC1.C=CC(=C)Cl.CCOCC.C[Si](C)(C)[CH2][Mg][Cl].O.[Cl-].[Mg].[NH4+] → C=CC(=C)C[Si](C)(C)C0.830
\n" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + } + ], + "source": [ + "# Build a dataframe with the result in the response\n", + "df = build_results_dataframe(response[\"result\"])\n", + "\n", + "# Create and show the data in a table\n", + "style_results_table(df, caption=\"Retrosynthesis Prediction Results\", align=\"right\")" + ] + }, + { + "cell_type": "markdown", + "id": "98ba77e3-fcf3-488a-b51c-54ce2dcfe7e0", + "metadata": {}, + "source": [ + "### Retro tree prediction" + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "id": "4cdba192-ec0c-494d-b356-b48cbdeb5cb8", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Task sent. Assigned task_id: 26ea63e0-2527-4baf-9905-673f5d98442e\n", + "State: PENDING | Info: None\n", + "State: STARTED | Info: None\n", + "State: SUCCESS | Info: {'result': [{'open_nodes': ['ClCCl', '[Cl][Ru]([Cl])(=[CH]c1ccccc1)([P](C1CCCCC1)(C1CCCCC1)C1CCCCC1)[P](C1CCCCC1)(C1CCCCC1)C1CCCCC1', 'C=CCC1C(=O)OC(=O)C1CC(=C)C[Si](C)(C)C'], 'expanded': ['C[Si](C)(C)CC1=CCC2C(=O)OC(=O)C2C1'], 'score': 0.7056475777714836, 'steps': 1, 'arcs': [{'product': 'C[Si](C)(C)CC1=CCC2C(=O)OC(=O)C2C1', 'reactants': ['C=CCC1C(=O)OC(=O)C1CC(=C)C[Si](C)(C)C', 'ClCCl', '[Cl][Ru]([Cl])(=[CH]c1ccccc1)([P](C1CCCCC1)(C1CCCCC1)C1CCCCC1)[P](C1CCCCC1)(C1CCCCC1)C1CCCCC1'], 'forward_likelihood': 0.8842267990112305, 'arc_score': 0.7056475777704836, 'confidence': 0.9439050974158504, 'metadata': {'reactants': 'C=CCC1C(=O)OC(=O)C1CC(=C)C[Si](C)(C)C.ClCCl.[Cl][Ru]([Cl])(=[CH]c1ccccc1)([P](C1CCCCC1)(C1CCCCC1)C1CCCCC1)[P](C1CCCCC1)(C1CCCCC1)C1CCCCC1', 'product': 'C[Si](C)(C)CC1=CCC2C(=O)OC(=O)C2C1', 'confidence': 0.9439050974158504, 'smiles': 'C=CCC1C(=O)OC(=O)C1CC(=C)C[Si](C)(C)C.ClCCl.[Cl][Ru]([Cl])(=[CH]c1ccccc1)([P](C1CCCCC1)(C1CCCCC1)C1CCCCC1)[P](C1CCCCC1)(C1CCCCC1)C1CCCCC1>>C[Si](C)(C)CC1=CCC2C(=O)OC(=O)C2C1'}}]}, {'open_nodes': ['ClCCl', 'Cc1cc(C)c(N2CCN(c3c(C)cc(C)cc3C)[C]2=[Ru]([Cl])([Cl])(=[CH]c2ccccc2)[P](C2CCCCC2)(C2CCCCC2)C2CCCCC2)c(C)c1', 'C=CCC1C(=O)OC(=O)C1CC(=C)C[Si](C)(C)C'], 'expanded': ['C[Si](C)(C)CC1=CCC2C(=O)OC(=O)C2C1'], 'score': 0.37149402762328515, 'steps': 1, 'arcs': [{'product': 'C[Si](C)(C)CC1=CCC2C(=O)OC(=O)C2C1', 'reactants': ['C=CCC1C(=O)OC(=O)C1CC(=C)C[Si](C)(C)C', 'Cc1cc(C)c(N2CCN(c3c(C)cc(C)cc3C)[C]2=[Ru]([Cl])([Cl])(=[CH]c2ccccc2)[P](C2CCCCC2)(C2CCCCC2)C2CCCCC2)c(C)c1', 'ClCCl'], 'forward_likelihood': 0.8087736368179321, 'arc_score': 0.3714940276222852, 'confidence': 0.9503329457084253, 'metadata': {'reactants': 'C=CCC1C(=O)OC(=O)C1CC(=C)C[Si](C)(C)C.Cc1cc(C)c(N2CCN(c3c(C)cc(C)cc3C)[C]2=[Ru]([Cl])([Cl])(=[CH]c2ccccc2)[P](C2CCCCC2)(C2CCCCC2)C2CCCCC2)c(C)c1.ClCCl', 'product': 'C[Si](C)(C)CC1=CCC2C(=O)OC(=O)C2C1', 'confidence': 0.9503329457084253, 'smiles': 'C=CCC1C(=O)OC(=O)C1CC(=C)C[Si](C)(C)C.Cc1cc(C)c(N2CCN(c3c(C)cc(C)cc3C)[C]2=[Ru]([Cl])([Cl])(=[CH]c2ccccc2)[P](C2CCCCC2)(C2CCCCC2)C2CCCCC2)c(C)c1.ClCCl>>C[Si](C)(C)CC1=CCC2C(=O)OC(=O)C2C1'}}]}, {'open_nodes': ['ClCCl', 'C=CCC1C(=O)OC(=O)C1CC(=C[Si](C)(C)C)C[Si](C)(C)C', 'Cc1cc(C)c(N2CCN(c3c(C)cc(C)cc3C)[C]2=[Ru]([Cl])([Cl])(=[CH]c2ccccc2)[P](C2CCCCC2)(C2CCCCC2)C2CCCCC2)c(C)c1'], 'expanded': ['C[Si](C)(C)CC1=CCC2C(=O)OC(=O)C2C1'], 'score': 0.33907517755787414, 'steps': 1, 'arcs': [{'product': 'C[Si](C)(C)CC1=CCC2C(=O)OC(=O)C2C1', 'reactants': ['C=CCC1C(=O)OC(=O)C1CC(=C[Si](C)(C)C)C[Si](C)(C)C', 'Cc1cc(C)c(N2CCN(c3c(C)cc(C)cc3C)[C]2=[Ru]([Cl])([Cl])(=[CH]c2ccccc2)[P](C2CCCCC2)(C2CCCCC2)C2CCCCC2)c(C)c1', 'ClCCl'], 'forward_likelihood': 0.8071097731590271, 'arc_score': 0.3390751775568742, 'confidence': 0.9455194125006082, 'metadata': {'reactants': 'C=CCC1C(=O)OC(=O)C1CC(=C[Si](C)(C)C)C[Si](C)(C)C.Cc1cc(C)c(N2CCN(c3c(C)cc(C)cc3C)[C]2=[Ru]([Cl])([Cl])(=[CH]c2ccccc2)[P](C2CCCCC2)(C2CCCCC2)C2CCCCC2)c(C)c1.ClCCl', 'product': 'C[Si](C)(C)CC1=CCC2C(=O)OC(=O)C2C1', 'confidence': 0.9455194125006082, 'smiles': 'C=CCC1C(=O)OC(=O)C1CC(=C[Si](C)(C)C)C[Si](C)(C)C.Cc1cc(C)c(N2CCN(c3c(C)cc(C)cc3C)[C]2=[Ru]([Cl])([Cl])(=[CH]c2ccccc2)[P](C2CCCCC2)(C2CCCCC2)C2CCCCC2)c(C)c1.ClCCl>>C[Si](C)(C)CC1=CCC2C(=O)OC(=O)C2C1'}}]}, {'open_nodes': ['C1CCOC1', 'O=C1C=CC(=O)O1', 'C=CC(=C)O', 'ClP(Cl)(Cl)(Cl)Cl', 'Cc1ccccc1', 'C[Si](C)(C)[CH2][Mg][Cl]'], 'expanded': ['C=CC(=C)C[Si](C)(C)C', 'C=CC(=C)Cl', 'C[Si](C)(C)CC1=CCC2C(=O)OC(=O)C2C1'], 'score': 0.24857583642032657, 'steps': 3, 'arcs': [{'product': 'C[Si](C)(C)CC1=CCC2C(=O)OC(=O)C2C1', 'reactants': ['C=CC(=C)C[Si](C)(C)C', 'Cc1ccccc1', 'O=C1C=CC(=O)O1'], 'forward_likelihood': 0.39315515756607056, 'arc_score': 0.43673660563264405, 'confidence': 0.9271861522313083, 'metadata': {'reactants': 'C=CC(=C)C[Si](C)(C)C.Cc1ccccc1.O=C1C=CC(=O)O1', 'product': 'C[Si](C)(C)CC1=CCC2C(=O)OC(=O)C2C1', 'confidence': 0.9271861522313083, 'smiles': 'C=CC(=C)C[Si](C)(C)C.Cc1ccccc1.O=C1C=CC(=O)O1>>C[Si](C)(C)CC1=CCC2C(=O)OC(=O)C2C1'}}, {'product': 'C=CC(=C)C[Si](C)(C)C', 'reactants': ['C1CCOC1', 'C=CC(=C)Cl', 'C[Si](C)(C)[CH2][Mg][Cl]'], 'forward_likelihood': 0.9396551847457886, 'arc_score': 1.0504772488213368, 'confidence': 0.8633497322540727, 'metadata': {'reactants': 'C1CCOC1.C=CC(=C)Cl.C[Si](C)(C)[CH2][Mg][Cl]', 'product': 'C=CC(=C)C[Si](C)(C)C', 'confidence': 0.8633497322540727, 'smiles': 'C1CCOC1.C=CC(=C)Cl.C[Si](C)(C)[CH2][Mg][Cl]>>C=CC(=C)C[Si](C)(C)C'}}, {'product': 'C=CC(=C)Cl', 'reactants': ['C=CC(=C)O', 'ClP(Cl)(Cl)(Cl)Cl'], 'forward_likelihood': 0.5561506152153015, 'arc_score': 0.541817045936755, 'confidence': 0.8283158501723433, 'metadata': {'reactants': 'C=CC(=C)O.ClP(Cl)(Cl)(Cl)Cl', 'product': 'C=CC(=C)Cl', 'confidence': 0.8283158501723433, 'smiles': 'C=CC(=C)O.ClP(Cl)(Cl)(Cl)Cl>>C=CC(=C)Cl'}}]}, {'open_nodes': ['O=C1C=CC(=O)O1', '[NH4+]', 'C=CC(=C)O', 'ClP(Cl)(Cl)(Cl)Cl', '[Cl-]', 'Cc1ccccc1', 'C[Si](C)(C)[CH2][Mg][Cl]', 'C1CCOC1', 'O'], 'expanded': ['C=CC(=C)C[Si](C)(C)C', 'C=CC(=C)Cl', 'C[Si](C)(C)CC1=CCC2C(=O)OC(=O)C2C1'], 'score': 0.2424893468373566, 'steps': 3, 'arcs': [{'product': 'C[Si](C)(C)CC1=CCC2C(=O)OC(=O)C2C1', 'reactants': ['C=CC(=C)C[Si](C)(C)C', 'Cc1ccccc1', 'O=C1C=CC(=O)O1'], 'forward_likelihood': 0.39315515756607056, 'arc_score': 0.43673660563264405, 'confidence': 0.9271861522313083, 'metadata': {'reactants': 'C=CC(=C)C[Si](C)(C)C.Cc1ccccc1.O=C1C=CC(=O)O1', 'product': 'C[Si](C)(C)CC1=CCC2C(=O)OC(=O)C2C1', 'confidence': 0.9271861522313083, 'smiles': 'C=CC(=C)C[Si](C)(C)C.Cc1ccccc1.O=C1C=CC(=O)O1>>C[Si](C)(C)CC1=CCC2C(=O)OC(=O)C2C1'}}, {'product': 'C=CC(=C)C[Si](C)(C)C', 'reactants': ['C1CCOC1', 'C=CC(=C)Cl', 'C[Si](C)(C)[CH2][Mg][Cl]', 'O', '[Cl-]', '[NH4+]'], 'forward_likelihood': 0.9166916608810425, 'arc_score': 1.02475584756135, 'confidence': 0.8341045215463807, 'metadata': {'reactants': 'C1CCOC1.C=CC(=C)Cl.C[Si](C)(C)[CH2][Mg][Cl].O.[Cl-].[NH4+]', 'product': 'C=CC(=C)C[Si](C)(C)C', 'confidence': 0.8341045215463807, 'smiles': 'C1CCOC1.C=CC(=C)Cl.C[Si](C)(C)[CH2][Mg][Cl].O.[Cl-].[NH4+]>>C=CC(=C)C[Si](C)(C)C'}}, {'product': 'C=CC(=C)Cl', 'reactants': ['C=CC(=C)O', 'ClP(Cl)(Cl)(Cl)Cl'], 'forward_likelihood': 0.5561506152153015, 'arc_score': 0.541817045936755, 'confidence': 0.8283158501723433, 'metadata': {'reactants': 'C=CC(=C)O.ClP(Cl)(Cl)(Cl)Cl', 'product': 'C=CC(=C)Cl', 'confidence': 0.8283158501723433, 'smiles': 'C=CC(=C)O.ClP(Cl)(Cl)(Cl)Cl>>C=CC(=C)Cl'}}]}, {'open_nodes': ['O=C1C=CC(=O)O1', 'C=CC(=C)O', 'CCOCC', 'ClP(Cl)(Cl)(Cl)Cl', 'Cc1ccccc1', 'C[Si](C)(C)[CH2][Mg][Cl]', 'C1CCOC1', 'Cl', 'O'], 'expanded': ['C=CC(=C)C[Si](C)(C)C', 'C=CC(=C)Cl', 'C[Si](C)(C)CC1=CCC2C(=O)OC(=O)C2C1'], 'score': 0.23326155626917652, 'steps': 3, 'arcs': [{'product': 'C[Si](C)(C)CC1=CCC2C(=O)OC(=O)C2C1', 'reactants': ['C=CC(=C)C[Si](C)(C)C', 'Cc1ccccc1', 'O=C1C=CC(=O)O1'], 'forward_likelihood': 0.39315515756607056, 'arc_score': 0.43673660563264405, 'confidence': 0.9271861522313083, 'metadata': {'reactants': 'C=CC(=C)C[Si](C)(C)C.Cc1ccccc1.O=C1C=CC(=O)O1', 'product': 'C[Si](C)(C)CC1=CCC2C(=O)OC(=O)C2C1', 'confidence': 0.9271861522313083, 'smiles': 'C=CC(=C)C[Si](C)(C)C.Cc1ccccc1.O=C1C=CC(=O)O1>>C[Si](C)(C)CC1=CCC2C(=O)OC(=O)C2C1'}}, {'product': 'C=CC(=C)C[Si](C)(C)C', 'reactants': ['C1CCOC1', 'C=CC(=C)Cl', 'CCOCC', 'C[Si](C)(C)[CH2][Mg][Cl]', 'Cl', 'O'], 'forward_likelihood': 0.8817673325538635, 'arc_score': 0.9857593618676268, 'confidence': 0.8303334989016697, 'metadata': {'reactants': 'C1CCOC1.C=CC(=C)Cl.CCOCC.C[Si](C)(C)[CH2][Mg][Cl].Cl.O', 'product': 'C=CC(=C)C[Si](C)(C)C', 'confidence': 0.8303334989016697, 'smiles': 'C1CCOC1.C=CC(=C)Cl.CCOCC.C[Si](C)(C)[CH2][Mg][Cl].Cl.O>>C=CC(=C)C[Si](C)(C)C'}}, {'product': 'C=CC(=C)Cl', 'reactants': ['C=CC(=C)O', 'ClP(Cl)(Cl)(Cl)Cl'], 'forward_likelihood': 0.5561506152153015, 'arc_score': 0.541817045936755, 'confidence': 0.8283158501723433, 'metadata': {'reactants': 'C=CC(=C)O.ClP(Cl)(Cl)(Cl)Cl', 'product': 'C=CC(=C)Cl', 'confidence': 0.8283158501723433, 'smiles': 'C=CC(=C)O.ClP(Cl)(Cl)(Cl)Cl>>C=CC(=C)Cl'}}]}, {'open_nodes': ['O=C1C=CC(=O)O1', '[NH4+]', 'C=CC(=C)O', 'CCOCC', '[Cl-]', 'ClP(Cl)(Cl)(Cl)Cl', 'Cc1ccccc1', 'C[Si](C)(C)[CH2][Mg][Cl]', 'C1CCOC1', 'O', '[Mg]'], 'expanded': ['C=CC(=C)C[Si](C)(C)C', 'C=CC(=C)Cl', 'C[Si](C)(C)CC1=CCC2C(=O)OC(=O)C2C1'], 'score': 0.21788071070873768, 'steps': 3, 'arcs': [{'product': 'C[Si](C)(C)CC1=CCC2C(=O)OC(=O)C2C1', 'reactants': ['C=CC(=C)C[Si](C)(C)C', 'Cc1ccccc1', 'O=C1C=CC(=O)O1'], 'forward_likelihood': 0.39315515756607056, 'arc_score': 0.43673660563264405, 'confidence': 0.9271861522313083, 'metadata': {'reactants': 'C=CC(=C)C[Si](C)(C)C.Cc1ccccc1.O=C1C=CC(=O)O1', 'product': 'C[Si](C)(C)CC1=CCC2C(=O)OC(=O)C2C1', 'confidence': 0.9271861522313083, 'smiles': 'C=CC(=C)C[Si](C)(C)C.Cc1ccccc1.O=C1C=CC(=O)O1>>C[Si](C)(C)CC1=CCC2C(=O)OC(=O)C2C1'}}, {'product': 'C=CC(=C)C[Si](C)(C)C', 'reactants': ['C1CCOC1', 'C=CC(=C)Cl', 'CCOCC', 'C[Si](C)(C)[CH2][Mg][Cl]', 'O', '[Cl-]', '[Mg]', '[NH4+]'], 'forward_likelihood': 0.8617748022079468, 'arc_score': 0.9207601706285793, 'confidence': 0.8301065237221663, 'metadata': {'reactants': 'C1CCOC1.C=CC(=C)Cl.CCOCC.C[Si](C)(C)[CH2][Mg][Cl].O.[Cl-].[Mg].[NH4+]', 'product': 'C=CC(=C)C[Si](C)(C)C', 'confidence': 0.8301065237221663, 'smiles': 'C1CCOC1.C=CC(=C)Cl.CCOCC.C[Si](C)(C)[CH2][Mg][Cl].O.[Cl-].[Mg].[NH4+]>>C=CC(=C)C[Si](C)(C)C'}}, {'product': 'C=CC(=C)Cl', 'reactants': ['C=CC(=C)O', 'ClP(Cl)(Cl)(Cl)Cl'], 'forward_likelihood': 0.5561506152153015, 'arc_score': 0.541817045936755, 'confidence': 0.8283158501723433, 'metadata': {'reactants': 'C=CC(=C)O.ClP(Cl)(Cl)(Cl)Cl', 'product': 'C=CC(=C)Cl', 'confidence': 0.8283158501723433, 'smiles': 'C=CC(=C)O.ClP(Cl)(Cl)(Cl)Cl>>C=CC(=C)Cl'}}]}, {'open_nodes': ['C1CCOC1', 'O=C1C=CC(=O)O1', 'C=CC(=C)O', 'O=S(Cl)Cl', 'Cc1ccccc1', 'C[Si](C)(C)[CH2][Mg][Cl]'], 'expanded': ['C=CC(=C)C[Si](C)(C)C', 'C=CC(=C)Cl', 'C[Si](C)(C)CC1=CCC2C(=O)OC(=O)C2C1'], 'score': 0.18048843471208253, 'steps': 3, 'arcs': [{'product': 'C[Si](C)(C)CC1=CCC2C(=O)OC(=O)C2C1', 'reactants': ['C=CC(=C)C[Si](C)(C)C', 'Cc1ccccc1', 'O=C1C=CC(=O)O1'], 'forward_likelihood': 0.39315515756607056, 'arc_score': 0.43673660563264405, 'confidence': 0.9271861522313083, 'metadata': {'reactants': 'C=CC(=C)C[Si](C)(C)C.Cc1ccccc1.O=C1C=CC(=O)O1', 'product': 'C[Si](C)(C)CC1=CCC2C(=O)OC(=O)C2C1', 'confidence': 0.9271861522313083, 'smiles': 'C=CC(=C)C[Si](C)(C)C.Cc1ccccc1.O=C1C=CC(=O)O1>>C[Si](C)(C)CC1=CCC2C(=O)OC(=O)C2C1'}}, {'product': 'C=CC(=C)C[Si](C)(C)C', 'reactants': ['C1CCOC1', 'C=CC(=C)Cl', 'C[Si](C)(C)[CH2][Mg][Cl]'], 'forward_likelihood': 0.9396551847457886, 'arc_score': 1.0504772488213368, 'confidence': 0.8633497322540727, 'metadata': {'reactants': 'C1CCOC1.C=CC(=C)Cl.C[Si](C)(C)[CH2][Mg][Cl]', 'product': 'C=CC(=C)C[Si](C)(C)C', 'confidence': 0.8633497322540727, 'smiles': 'C1CCOC1.C=CC(=C)Cl.C[Si](C)(C)[CH2][Mg][Cl]>>C=CC(=C)C[Si](C)(C)C'}}, {'product': 'C=CC(=C)Cl', 'reactants': ['C=CC(=C)O', 'O=S(Cl)Cl'], 'forward_likelihood': 0.403808057308197, 'arc_score': 0.39340795118967853, 'confidence': 0.8269444351225532, 'metadata': {'reactants': 'C=CC(=C)O.O=S(Cl)Cl', 'product': 'C=CC(=C)Cl', 'confidence': 0.8269444351225532, 'smiles': 'C=CC(=C)O.O=S(Cl)Cl>>C=CC(=C)Cl'}}]}, {'open_nodes': ['O=C1OC(=O)C2CC=CCC12', 'Cc1cc(C)c(N2CCN(c3c(C)cc(C)cc3C)[C]2=[Ru]([Cl])([Cl])=[CH]c2ccccc2)c(C)c1', 'ClCCl', 'C=CC(=C)C[Si](C)(C)C', 'C1CCC(P(C2CCCCC2)C2CCCCC2)CC1', 'N#N'], 'expanded': ['C[Si](C)(C)CC1=CCC2C(=O)OC(=O)C2C1'], 'score': 0.1796396903569043, 'steps': 1, 'arcs': [{'product': 'C[Si](C)(C)CC1=CCC2C(=O)OC(=O)C2C1', 'reactants': ['C1CCC(P(C2CCCCC2)C2CCCCC2)CC1', 'C=CC(=C)C[Si](C)(C)C', 'Cc1cc(C)c(N2CCN(c3c(C)cc(C)cc3C)[C]2=[Ru]([Cl])([Cl])=[CH]c2ccccc2)c(C)c1', 'ClCCl', 'N#N', 'O=C1OC(=O)C2CC=CCC12'], 'forward_likelihood': 0.395216703414917, 'arc_score': 0.1796396903559043, 'confidence': 0.9290369156214755, 'metadata': {'reactants': 'C1CCC(P(C2CCCCC2)C2CCCCC2)CC1.C=CC(=C)C[Si](C)(C)C.Cc1cc(C)c(N2CCN(c3c(C)cc(C)cc3C)[C]2=[Ru]([Cl])([Cl])=[CH]c2ccccc2)c(C)c1.ClCCl.N#N.O=C1OC(=O)C2CC=CCC12', 'product': 'C[Si](C)(C)CC1=CCC2C(=O)OC(=O)C2C1', 'confidence': 0.9290369156214755, 'smiles': 'C1CCC(P(C2CCCCC2)C2CCCCC2)CC1.C=CC(=C)C[Si](C)(C)C.Cc1cc(C)c(N2CCN(c3c(C)cc(C)cc3C)[C]2=[Ru]([Cl])([Cl])=[CH]c2ccccc2)c(C)c1.ClCCl.N#N.O=C1OC(=O)C2CC=CCC12>>C[Si](C)(C)CC1=CCC2C(=O)OC(=O)C2C1'}}]}, {'open_nodes': ['O=C1C=CC(=O)O1', '[NH4+]', 'C=CC(=C)O', '[Cl-]', 'O=S(Cl)Cl', 'Cc1ccccc1', 'C[Si](C)(C)[CH2][Mg][Cl]', 'C1CCOC1', 'O'], 'expanded': ['C=CC(=C)C[Si](C)(C)C', 'C=CC(=C)Cl', 'C[Si](C)(C)CC1=CCC2C(=O)OC(=O)C2C1'], 'score': 0.17606909535254775, 'steps': 3, 'arcs': [{'product': 'C[Si](C)(C)CC1=CCC2C(=O)OC(=O)C2C1', 'reactants': ['C=CC(=C)C[Si](C)(C)C', 'Cc1ccccc1', 'O=C1C=CC(=O)O1'], 'forward_likelihood': 0.39315515756607056, 'arc_score': 0.43673660563264405, 'confidence': 0.9271861522313083, 'metadata': {'reactants': 'C=CC(=C)C[Si](C)(C)C.Cc1ccccc1.O=C1C=CC(=O)O1', 'product': 'C[Si](C)(C)CC1=CCC2C(=O)OC(=O)C2C1', 'confidence': 0.9271861522313083, 'smiles': 'C=CC(=C)C[Si](C)(C)C.Cc1ccccc1.O=C1C=CC(=O)O1>>C[Si](C)(C)CC1=CCC2C(=O)OC(=O)C2C1'}}, {'product': 'C=CC(=C)C[Si](C)(C)C', 'reactants': ['C1CCOC1', 'C=CC(=C)Cl', 'C[Si](C)(C)[CH2][Mg][Cl]', 'O', '[Cl-]', '[NH4+]'], 'forward_likelihood': 0.9166916608810425, 'arc_score': 1.02475584756135, 'confidence': 0.8341045215463807, 'metadata': {'reactants': 'C1CCOC1.C=CC(=C)Cl.C[Si](C)(C)[CH2][Mg][Cl].O.[Cl-].[NH4+]', 'product': 'C=CC(=C)C[Si](C)(C)C', 'confidence': 0.8341045215463807, 'smiles': 'C1CCOC1.C=CC(=C)Cl.C[Si](C)(C)[CH2][Mg][Cl].O.[Cl-].[NH4+]>>C=CC(=C)C[Si](C)(C)C'}}, {'product': 'C=CC(=C)Cl', 'reactants': ['C=CC(=C)O', 'O=S(Cl)Cl'], 'forward_likelihood': 0.403808057308197, 'arc_score': 0.39340795118967853, 'confidence': 0.8269444351225532, 'metadata': {'reactants': 'C=CC(=C)O.O=S(Cl)Cl', 'product': 'C=CC(=C)Cl', 'confidence': 0.8269444351225532, 'smiles': 'C=CC(=C)O.O=S(Cl)Cl>>C=CC(=C)Cl'}}]}], 'time': 1306.3053419589996}\n" + ] + } + ], + "source": [ + "# Choose product for retrosynthesis tree prediction\n", + "product = \"C1C(C[Si](C)(C)C)=CCC2C(=O)OC(=O)C12\"\n", + "#product = \"Cc1cc2scnc2cc1N\"\n", + "#product = \"CC(C)(C(=O)O)C1C=CC=C(C2CC2)C1=O\"\n", + "#product = \"Nc1ccc2scnc2c1Br\"\n", + "\n", + "# Setup task kwargs\n", + "kwargs = {\n", + " \"topn\": 15, # Number of results per reactant\n", + " \"num_beams\": 15, # Number of beams used for prediction. Must be >= topn\n", + " \"fap\": 0.6, # Forward likelihood acceptance probability (not length averaged)\n", + " \"fld\": 0.2, # Forward likelihood delta required between the top2 forward prediction results\n", + " \"max_depth\": 4, # Max depth of the retrosynthesis tree\n", + " \"beam_width\": 6, # Max amount of nodes being expanded in each step\n", + " \"device\": None, # Device used for predicting, either \"cuda\" or \"cpu\", None defaults to cuda if available\n", + " \"ckpt_forward\": \"Pistachio2025Q2-Forward\", # Default forward model\n", + " \"ckpt_retro\": \"Pistachio2025Q2-Retro\", # Default retrosynthesis model\n", + " \"vocab\": \"Pistachio2025Q2\", # Vocab for default forward and retrosynthesis models\n", + " # \"ckpt_forward_path\": \"models/forward/Pistachio2025Q2-Forward.ckpt\", # Can be used instead of ckpt_forward\n", + " # \"ckpt_retro_path\": \"models/retrosynthesis/Pistachio2025Q2-Retro.ckpt\", # Can be used instead of ckpt_retro\n", + " # \"vocab_path\": \"vocab/Pistachio2025Q2.txt\", # Can be used instead of vocab\n", + "}\n", + "\n", + "# Send the retro_prediction_tree task with the product and kwargs\n", + "task = celery_app.send_task(\n", + " \"tasks.retro_prediction_tree\",\n", + " [product],\n", + " kwargs=kwargs,\n", + " queue=\"retro_prediction\",\n", + ")\n", + "print(\"Task sent. Assigned task_id: {}\".format(task.id))\n", + "\n", + "# Use the task id to get the result. Increase timeout if needed.\n", + "response = wait_for_result(celery_app, task.id, timeout=3000)" + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "id": "ee06bf27-b7ac-4d07-bff1-2cf99322b80d", + "metadata": {}, + "outputs": [ + { + "data": { + "text/markdown": [ + "## Route 1\n", + "- **Score:** `0.706`\n", + "- **Steps:** `1`" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/markdown": [ + "**Remaining molecules:** `ClCCl, [Cl][Ru]([Cl])(=[CH]c1ccccc1)([P](C1CCCCC1)(C1CCCCC1)C1CCCCC1)[P](C1CCCCC1)(C1CCCCC1)C1CCCCC1, C=CCC1C(=O)OC(=O)C1CC(=C)C[Si](C)(C)C`" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/markdown": [ + "#### Step 1\n", + "**Product:** `C[Si](C)(C)CC1=CCC2C(=O)OC(=O)C2C1` \n", + "**Reactants:** `C=CCC1C(=O)OC(=O)C1CC(=C)C[Si](C)(C)C.ClCCl.[Cl][Ru]([Cl])(=[CH]c1ccccc1)([P](C1CCCCC1)(C1CCCCC1)C1CCCCC1)[P](C1CCCCC1)(C1CCCCC1)C1CCCCC1` \n", + "- Forward likelihood: `0.884` \n", + "- Arc score: `0.706` \n", + "- Retro confidence: `0.9439050974158504`" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/markdown": [ + "## Route 2\n", + "- **Score:** `0.371`\n", + "- **Steps:** `1`" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/markdown": [ + "**Remaining molecules:** `ClCCl, Cc1cc(C)c(N2CCN(c3c(C)cc(C)cc3C)[C]2=[Ru]([Cl])([Cl])(=[CH]c2ccccc2)[P](C2CCCCC2)(C2CCCCC2)C2CCCCC2)c(C)c1, C=CCC1C(=O)OC(=O)C1CC(=C)C[Si](C)(C)C`" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/markdown": [ + "#### Step 1\n", + "**Product:** `C[Si](C)(C)CC1=CCC2C(=O)OC(=O)C2C1` \n", + "**Reactants:** `C=CCC1C(=O)OC(=O)C1CC(=C)C[Si](C)(C)C.Cc1cc(C)c(N2CCN(c3c(C)cc(C)cc3C)[C]2=[Ru]([Cl])([Cl])(=[CH]c2ccccc2)[P](C2CCCCC2)(C2CCCCC2)C2CCCCC2)c(C)c1.ClCCl` \n", + "- Forward likelihood: `0.809` \n", + "- Arc score: `0.371` \n", + "- Retro confidence: `0.9503329457084253`" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/markdown": [ + "## Route 3\n", + "- **Score:** `0.339`\n", + "- **Steps:** `1`" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/markdown": [ + "**Remaining molecules:** `ClCCl, C=CCC1C(=O)OC(=O)C1CC(=C[Si](C)(C)C)C[Si](C)(C)C, Cc1cc(C)c(N2CCN(c3c(C)cc(C)cc3C)[C]2=[Ru]([Cl])([Cl])(=[CH]c2ccccc2)[P](C2CCCCC2)(C2CCCCC2)C2CCCCC2)c(C)c1`" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/markdown": [ + "#### Step 1\n", + "**Product:** `C[Si](C)(C)CC1=CCC2C(=O)OC(=O)C2C1` \n", + "**Reactants:** `C=CCC1C(=O)OC(=O)C1CC(=C[Si](C)(C)C)C[Si](C)(C)C.Cc1cc(C)c(N2CCN(c3c(C)cc(C)cc3C)[C]2=[Ru]([Cl])([Cl])(=[CH]c2ccccc2)[P](C2CCCCC2)(C2CCCCC2)C2CCCCC2)c(C)c1.ClCCl` \n", + "- Forward likelihood: `0.807` \n", + "- Arc score: `0.339` \n", + "- Retro confidence: `0.9455194125006082`" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/markdown": [ + "## Route 4\n", + "- **Score:** `0.249`\n", + "- **Steps:** `3`" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/markdown": [ + "**Remaining molecules:** `C1CCOC1, O=C1C=CC(=O)O1, C=CC(=C)O, ClP(Cl)(Cl)(Cl)Cl, Cc1ccccc1, C[Si](C)(C)[CH2][Mg][Cl]`" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/markdown": [ + "#### Step 1\n", + "**Product:** `C[Si](C)(C)CC1=CCC2C(=O)OC(=O)C2C1` \n", + "**Reactants:** `C=CC(=C)C[Si](C)(C)C.Cc1ccccc1.O=C1C=CC(=O)O1` \n", + "- Forward likelihood: `0.393` \n", + "- Arc score: `0.437` \n", + "- Retro confidence: `0.9271861522313083`" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/markdown": [ + "#### Step 2\n", + "**Product:** `C=CC(=C)C[Si](C)(C)C` \n", + "**Reactants:** `C1CCOC1.C=CC(=C)Cl.C[Si](C)(C)[CH2][Mg][Cl]` \n", + "- Forward likelihood: `0.940` \n", + "- Arc score: `1.050` \n", + "- Retro confidence: `0.8633497322540727`" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/markdown": [ + "#### Step 3\n", + "**Product:** `C=CC(=C)Cl` \n", + "**Reactants:** `C=CC(=C)O.ClP(Cl)(Cl)(Cl)Cl` \n", + "- Forward likelihood: `0.556` \n", + "- Arc score: `0.542` \n", + "- Retro confidence: `0.8283158501723433`" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/markdown": [ + "## Route 5\n", + "- **Score:** `0.242`\n", + "- **Steps:** `3`" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/markdown": [ + "**Remaining molecules:** `O=C1C=CC(=O)O1, [NH4+], C=CC(=C)O, ClP(Cl)(Cl)(Cl)Cl, [Cl-], Cc1ccccc1, C[Si](C)(C)[CH2][Mg][Cl], C1CCOC1, O`" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/markdown": [ + "#### Step 1\n", + "**Product:** `C[Si](C)(C)CC1=CCC2C(=O)OC(=O)C2C1` \n", + "**Reactants:** `C=CC(=C)C[Si](C)(C)C.Cc1ccccc1.O=C1C=CC(=O)O1` \n", + "- Forward likelihood: `0.393` \n", + "- Arc score: `0.437` \n", + "- Retro confidence: `0.9271861522313083`" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/markdown": [ + "#### Step 2\n", + "**Product:** `C=CC(=C)C[Si](C)(C)C` \n", + "**Reactants:** `C1CCOC1.C=CC(=C)Cl.C[Si](C)(C)[CH2][Mg][Cl].O.[Cl-].[NH4+]` \n", + "- Forward likelihood: `0.917` \n", + "- Arc score: `1.025` \n", + "- Retro confidence: `0.8341045215463807`" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/markdown": [ + "#### Step 3\n", + "**Product:** `C=CC(=C)Cl` \n", + "**Reactants:** `C=CC(=C)O.ClP(Cl)(Cl)(Cl)Cl` \n", + "- Forward likelihood: `0.556` \n", + "- Arc score: `0.542` \n", + "- Retro confidence: `0.8283158501723433`" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/markdown": [ + "## Route 6\n", + "- **Score:** `0.233`\n", + "- **Steps:** `3`" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/markdown": [ + "**Remaining molecules:** `O=C1C=CC(=O)O1, C=CC(=C)O, CCOCC, ClP(Cl)(Cl)(Cl)Cl, Cc1ccccc1, C[Si](C)(C)[CH2][Mg][Cl], C1CCOC1, Cl, O`" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/markdown": [ + "#### Step 1\n", + "**Product:** `C[Si](C)(C)CC1=CCC2C(=O)OC(=O)C2C1` \n", + "**Reactants:** `C=CC(=C)C[Si](C)(C)C.Cc1ccccc1.O=C1C=CC(=O)O1` \n", + "- Forward likelihood: `0.393` \n", + "- Arc score: `0.437` \n", + "- Retro confidence: `0.9271861522313083`" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/markdown": [ + "#### Step 2\n", + "**Product:** `C=CC(=C)C[Si](C)(C)C` \n", + "**Reactants:** `C1CCOC1.C=CC(=C)Cl.CCOCC.C[Si](C)(C)[CH2][Mg][Cl].Cl.O` \n", + "- Forward likelihood: `0.882` \n", + "- Arc score: `0.986` \n", + "- Retro confidence: `0.8303334989016697`" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/markdown": [ + "#### Step 3\n", + "**Product:** `C=CC(=C)Cl` \n", + "**Reactants:** `C=CC(=C)O.ClP(Cl)(Cl)(Cl)Cl` \n", + "- Forward likelihood: `0.556` \n", + "- Arc score: `0.542` \n", + "- Retro confidence: `0.8283158501723433`" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/markdown": [ + "## Route 7\n", + "- **Score:** `0.218`\n", + "- **Steps:** `3`" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/markdown": [ + "**Remaining molecules:** `O=C1C=CC(=O)O1, [NH4+], C=CC(=C)O, CCOCC, [Cl-], ClP(Cl)(Cl)(Cl)Cl, Cc1ccccc1, C[Si](C)(C)[CH2][Mg][Cl], C1CCOC1, O, [Mg]`" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/markdown": [ + "#### Step 1\n", + "**Product:** `C[Si](C)(C)CC1=CCC2C(=O)OC(=O)C2C1` \n", + "**Reactants:** `C=CC(=C)C[Si](C)(C)C.Cc1ccccc1.O=C1C=CC(=O)O1` \n", + "- Forward likelihood: `0.393` \n", + "- Arc score: `0.437` \n", + "- Retro confidence: `0.9271861522313083`" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/markdown": [ + "#### Step 2\n", + "**Product:** `C=CC(=C)C[Si](C)(C)C` \n", + "**Reactants:** `C1CCOC1.C=CC(=C)Cl.CCOCC.C[Si](C)(C)[CH2][Mg][Cl].O.[Cl-].[Mg].[NH4+]` \n", + "- Forward likelihood: `0.862` \n", + "- Arc score: `0.921` \n", + "- Retro confidence: `0.8301065237221663`" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/markdown": [ + "#### Step 3\n", + "**Product:** `C=CC(=C)Cl` \n", + "**Reactants:** `C=CC(=C)O.ClP(Cl)(Cl)(Cl)Cl` \n", + "- Forward likelihood: `0.556` \n", + "- Arc score: `0.542` \n", + "- Retro confidence: `0.8283158501723433`" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/markdown": [ + "## Route 8\n", + "- **Score:** `0.180`\n", + "- **Steps:** `3`" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/markdown": [ + "**Remaining molecules:** `C1CCOC1, O=C1C=CC(=O)O1, C=CC(=C)O, O=S(Cl)Cl, Cc1ccccc1, C[Si](C)(C)[CH2][Mg][Cl]`" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/markdown": [ + "#### Step 1\n", + "**Product:** `C[Si](C)(C)CC1=CCC2C(=O)OC(=O)C2C1` \n", + "**Reactants:** `C=CC(=C)C[Si](C)(C)C.Cc1ccccc1.O=C1C=CC(=O)O1` \n", + "- Forward likelihood: `0.393` \n", + "- Arc score: `0.437` \n", + "- Retro confidence: `0.9271861522313083`" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/markdown": [ + "#### Step 2\n", + "**Product:** `C=CC(=C)C[Si](C)(C)C` \n", + "**Reactants:** `C1CCOC1.C=CC(=C)Cl.C[Si](C)(C)[CH2][Mg][Cl]` \n", + "- Forward likelihood: `0.940` \n", + "- Arc score: `1.050` \n", + "- Retro confidence: `0.8633497322540727`" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/markdown": [ + "#### Step 3\n", + "**Product:** `C=CC(=C)Cl` \n", + "**Reactants:** `C=CC(=C)O.O=S(Cl)Cl` \n", + "- Forward likelihood: `0.404` \n", + "- Arc score: `0.393` \n", + "- Retro confidence: `0.8269444351225532`" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/markdown": [ + "## Route 9\n", + "- **Score:** `0.180`\n", + "- **Steps:** `1`" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/markdown": [ + "**Remaining molecules:** `O=C1OC(=O)C2CC=CCC12, Cc1cc(C)c(N2CCN(c3c(C)cc(C)cc3C)[C]2=[Ru]([Cl])([Cl])=[CH]c2ccccc2)c(C)c1, ClCCl, C=CC(=C)C[Si](C)(C)C, C1CCC(P(C2CCCCC2)C2CCCCC2)CC1, N#N`" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/markdown": [ + "#### Step 1\n", + "**Product:** `C[Si](C)(C)CC1=CCC2C(=O)OC(=O)C2C1` \n", + "**Reactants:** `C1CCC(P(C2CCCCC2)C2CCCCC2)CC1.C=CC(=C)C[Si](C)(C)C.Cc1cc(C)c(N2CCN(c3c(C)cc(C)cc3C)[C]2=[Ru]([Cl])([Cl])=[CH]c2ccccc2)c(C)c1.ClCCl.N#N.O=C1OC(=O)C2CC=CCC12` \n", + "- Forward likelihood: `0.395` \n", + "- Arc score: `0.180` \n", + "- Retro confidence: `0.9290369156214755`" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/markdown": [ + "## Route 10\n", + "- **Score:** `0.176`\n", + "- **Steps:** `3`" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/markdown": [ + "**Remaining molecules:** `O=C1C=CC(=O)O1, [NH4+], C=CC(=C)O, [Cl-], O=S(Cl)Cl, Cc1ccccc1, C[Si](C)(C)[CH2][Mg][Cl], C1CCOC1, O`" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/markdown": [ + "#### Step 1\n", + "**Product:** `C[Si](C)(C)CC1=CCC2C(=O)OC(=O)C2C1` \n", + "**Reactants:** `C=CC(=C)C[Si](C)(C)C.Cc1ccccc1.O=C1C=CC(=O)O1` \n", + "- Forward likelihood: `0.393` \n", + "- Arc score: `0.437` \n", + "- Retro confidence: `0.9271861522313083`" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/markdown": [ + "#### Step 2\n", + "**Product:** `C=CC(=C)C[Si](C)(C)C` \n", + "**Reactants:** `C1CCOC1.C=CC(=C)Cl.C[Si](C)(C)[CH2][Mg][Cl].O.[Cl-].[NH4+]` \n", + "- Forward likelihood: `0.917` \n", + "- Arc score: `1.025` \n", + "- Retro confidence: `0.8341045215463807`" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/markdown": [ + "#### Step 3\n", + "**Product:** `C=CC(=C)Cl` \n", + "**Reactants:** `C=CC(=C)O.O=S(Cl)Cl` \n", + "- Forward likelihood: `0.404` \n", + "- Arc score: `0.393` \n", + "- Retro confidence: `0.8269444351225532`" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + } + ], + "source": [ + "# Textual retrosynthesis tree representation\n", + "for i, route in enumerate(response[\"result\"]):\n", + " display_route(route, i)" + ] + }, + { + "cell_type": "code", + "execution_count": 13, + "id": "8491bcb4-b096-4277-9244-8b59b4499fe2", + "metadata": {}, + "outputs": [ + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "37b74ae49d7748b7acbc0e82ad6b285d", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "interactive(children=(Dropdown(description='Route', options=(('Route 1 | score=0.706', 0), ('Route 2 | score=0…" + ] + }, + "metadata": {}, + "output_type": "display_data" + } + ], + "source": [ + "# Graphical retrosynthesis tree representation with route selection\n", + "tree_route_selector(response[\"result\"])" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5f7af0df-b7a9-4e93-ac40-27fc33d0a103", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.13.12" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/jupyter/requirements.txt b/jupyter/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..699b8257d6b8e5e60e32d9c1f8db362b157a868b --- /dev/null +++ b/jupyter/requirements.txt @@ -0,0 +1,7 @@ +celery[redis] +ipykernel +ipywidgets +matplotlib +notebook +pandas +rdkit diff --git a/mcp/Dockerfile b/mcp/Dockerfile new file mode 100644 index 0000000000000000000000000000000000000000..d1db2d1ecb6c1b51e52718d2f4ff5568ae162eba --- /dev/null +++ b/mcp/Dockerfile @@ -0,0 +1,11 @@ +FROM public.ecr.aws/docker/library/python:3.13-slim + +WORKDIR /app + +COPY requirements.txt . +RUN pip install --no-cache-dir -r requirements.txt + +COPY server.py start.sh . +RUN chmod +x start.sh + +CMD ["./start.sh"] diff --git a/mcp/requirements.txt b/mcp/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..2d729fe6cd26df87006d78fd12c8ef194cce70dd --- /dev/null +++ b/mcp/requirements.txt @@ -0,0 +1,4 @@ +celery +mcp +mcpo +redis diff --git a/mcp/server.py b/mcp/server.py new file mode 100644 index 0000000000000000000000000000000000000000..e89f3601a2825bd2eb592a3775b419c3bcfce392 --- /dev/null +++ b/mcp/server.py @@ -0,0 +1,198 @@ +import os +import asyncio +from typing import Optional +from mcp.server.fastmcp import FastMCP +from celery import Celery +from celery.exceptions import TimeoutError + + +CELERY_BROKER_URL = os.getenv("CELERY_BROKER_URL", "amqp://ubuntu:ubuntu@broker:5672//") +CELERY_RESULT_BACKEND = os.getenv("CELERY_RESULT_BACKEND", "redis://:ubuntu@redis:6379/0") + +app = Celery("mcp-client", broker=CELERY_BROKER_URL, backend=CELERY_RESULT_BACKEND) + +mcp = FastMCP("rxn-mcp-server", json_response=False) + + +async def wait_for_celery_result(task, timeout: int = 300): + def _get(): + return task.get(timeout=timeout) + + return await asyncio.to_thread(_get) + + +@mcp.tool() +async def product_prediction( + reactants_list: list[str], + topn: int = 3, + num_beams: int = 3, + device: Optional[str] = None, + ckpt_forward = "Pistachio2025Q2-Forward", + vocab = "Pistachio2025Q2", +): + """ + Gets the product prediction for a batch of reactants SMILES. + + Args: + reactants_list (list[str]): List of reactants SMILES. + topn (int): number of predictions per reactant (Defaults to 3). + num_beams (int): beams used for prediction (num_beams >= topn) (Defaults to 3). + device (Optional str): device used for predicting, either "cuda" or "cpu", None defaults to cuda if available. + ckpt_forward (Optional str): name of the forward model without extension. Should exist inside the "models/forward" directory and be a .ckpt file. + vocab (Optional str): name of the vocab file without extension. Should exist inside the "vocab" directory and be a .txt file. + """ + try: + kwargs = { + "topn": topn, + "num_beams": num_beams, + "device": device, + "ckpt_forward": ckpt_forward, + "vocab": vocab, + } + task = app.send_task( + "tasks.product_prediction", + [reactants_list], + kwargs=kwargs, + queue="product_prediction", + ) + result = await wait_for_celery_result(task, timeout=600) + result["status"] = "success" + return result + except TimeoutError: + return { + "status": "timeout", + "message": "Computation exceeded time limit", + } + except Exception as e: + return { + "status": "error", + "message": str(e), + } + + +@mcp.tool() +async def retro_prediction( + product: str, + topn: int = 15, + num_beams: int = 15, + fap: float = 0.6, + fld: float = 0.2, + device: Optional[str] = None, + ckpt_forward = "Pistachio2025Q2-Forward", + ckpt_retro = "Pistachio2025Q2-Retro", + vocab = "Pistachio2025Q2", +): + """ + Gets the retrosynthesis prediction for a product SMILES. + + Args: + product (str): List of reactants SMILES. + topn (int): number of retrosynthesis predictions (Defaults to 15). + num_beams (int): beams used for prediction (num_beams >= topn) (Defaults to 15). + fap (float): forward acceptance level (Defaults to 0.6). + fld (float): forward likelihood delta (Defaults to 0.2). + device (Optional str): device used for predicting, either "cuda" or "cpu", None defaults to cuda if available. + ckpt_forward (Optional str): name of the forward model without extension. Should exist inside the "models/forward" directory and be a .ckpt file. + ckpt_retro (Optional str): name of the retrosynthesis model without extension. Should exist inside the "models/retrosynthesis" directory and be a .ckpt file. + vocab (Optional str): name of the vocab file without extension. Should exist inside the "vocab" directory and be a .txt file. + """ + try: + kwargs = { + "topn": topn, + "num_beams": num_beams, + "fap": fap, + "fld": fld, + "device": device, + "ckpt_forward": ckpt_forward, + "ckpt_retro": ckpt_retro, + "vocab": vocab, + } + task = app.send_task( + "tasks.retro_prediction", + [product], + kwargs=kwargs, + queue="retro_prediction", + ) + result = await wait_for_celery_result(task, timeout=600) + result["status"] = "success" + return result + except TimeoutError: + return { + "status": "timeout", + "message": "Computation exceeded time limit", + } + except Exception as e: + return { + "status": "error", + "message": str(e), + } + + +@mcp.tool() +async def retro_tree_prediction( + product: str, + topn: int = 15, + num_beams: int = 15, + fap: float = 0.6, + fld: float = 0.2, + max_depth: int = 3, + beam_width: int = 5, + device: Optional[str] = None, + ckpt_forward = "Pistachio2025Q2-Forward", + ckpt_retro = "Pistachio2025Q2-Retro", + vocab = "Pistachio2025Q2", +): + """ + Gets the full retrosynthesis tree for a product SMILES. + + Args: + product (str): List of reactants SMILES. + topn (int): number of retrosynthesis predictions (Defaults to 15). + num_beams (int): beams used for prediction (num_beams >= topn) (Defaults to 15). + fap (float): forward acceptance level (Defaults to 0.6). + fld (float): forward likelihood delta (Defaults to 0.2). + max_depth (int): maximum number of retrosynthesis expansions in the tree (Defaults to 3). + beam_width (int): number of molecules chosen for expansion in each step of the tree (Defaults to 5). + device (Optional str): device used for predicting, either "cuda" or "cpu", None defaults to cuda if available. + ckpt_forward (Optional str): name of the forward model without extension. Should exist inside the "models/forward" directory and be a .ckpt file. + ckpt_retro (Optional str): name of the retrosynthesis model without extension. Should exist inside the "models/retrosynthesis" directory and be a .ckpt file. + vocab (Optional str): name of the vocab file without extension. Should exist inside the "vocab" directory and be a .txt file. + """ + try: + kwargs = { + "topn": topn, + "num_beams": num_beams, + "fap": fap, + "fld": fld, + "max_depth": max_depth, + "beam_width": beam_width, + "device": device, + "ckpt_forward": ckpt_forward, + "ckpt_retro": ckpt_retro, + "vocab": vocab, + } + task = app.send_task( + "tasks.retro_prediction_tree", + [product], + kwargs=kwargs, + queue="retro_prediction", + ) + result = await wait_for_celery_result(task, 1500) + result["status"] = "success" + return result + except TimeoutError: + return { + "status": "timeout", + "message": "Computation exceeded time limit", + } + except Exception as e: + return { + "status": "error", + "message": str(e), + } + + +if __name__ == "__main__": + mcp.settings.host = "0.0.0.0" + mcp.settings.port = 8001 + asyncio.run(mcp.run_streamable_http_async()) diff --git a/mcp/start.sh b/mcp/start.sh new file mode 100644 index 0000000000000000000000000000000000000000..d94b97f7a29d29931432cdb784f6d392482a447d --- /dev/null +++ b/mcp/start.sh @@ -0,0 +1,12 @@ +#!/bin/sh +set -e + +echo "Starting FastMCP on port 8001" +python server.py & + +echo "Starting MCPO on port 8000" +exec mcpo \ + --host 0.0.0.0 \ + --port 8000 \ + --server-type streamable_http \ + -- http://127.0.0.1:8001/mcp diff --git a/model-training/README.md b/model-training/README.md new file mode 100644 index 0000000000000000000000000000000000000000..f8ed2749d770556da4d67b4d5cb257d1895ac632 --- /dev/null +++ b/model-training/README.md @@ -0,0 +1,186 @@ +# Model Training + +This directory contains everything needed to retrain the forward prediction and retrosynthesis models. + +- **`pistachio/`** — Scripts to download and preprocess Pistachio reaction data into the JSONL format required for training. See [`pistachio/README.md`](pistachio/README.md) for the full data pipeline. +- **`training/`** — The LightningCLI entry point (`cli_main.py`) and YAML configuration files for training, testing, and running predictions. +- **`tools/`** — Standalone utility scripts for data analysis, model diagnostics, and running training/testing/prediction directly without LightningCLI. + +The model source code lives in [`worker/transformers_model/`](../worker/transformers_model/). The key files are: +- [`model.py`](../worker/transformers_model/model.py) — `LitVanillaTransformer`, the Lightning module that wraps the model. +- [`smiles_datamodule.py`](../worker/transformers_model/smiles_datamodule.py) — `LitSmilesDataset`, the Lightning data module. +- [`configuration.py`](../worker/transformers_model/configuration.py) — `VanillaTransformerConfig`, the model's default hyperparameter configuration. + +--- + +## Setup + +### Copy the model package + +`cli_main.py` imports `transformers_model` as a local package. Python has no way to resolve this import across the repository boundary to `worker/transformers_model/` without installing it, so the folder must be copied into `model-training/` before running any command: + +```bash +cp -r ../worker/transformers_model model-training/transformers_model +``` + +This only needs to be done once, or again whenever the model source in `worker/` changes. + +### Install dependencies + +#### With uv (recommended) + +```bash +# Install uv if not already available +curl -LsSf https://astral.sh/uv/install.sh | sh + +# Create and activate a virtual environment +uv venv --python 3.11 +source .venv/bin/activate + +# Install dependencies +uv pip install -r requirements.txt +``` + +#### Without uv + +```bash +python3.11 -m venv .venv +source .venv/bin/activate +pip install -r requirements.txt +``` + +> On Windows replace `source .venv/bin/activate` with `.venv\Scripts\activate`. + +--- + +## Running LightningCLI commands + +All commands are run from the `model-training/` directory using `cli_main.py` together with one of the provided YAML configuration files. + +### Training + +```bash +# Forward model — with uv +uv run python training/cli_main.py fit --config=training/train.yaml + +# Forward model — without uv (venv already activated) +python training/cli_main.py fit --config=training/train.yaml + +# Retrosynthesis model +python training/cli_main.py fit --config=training/train_retro.yaml +``` + +### Testing + +```bash +# Forward model +python training/cli_main.py test --config=training/test.yaml + +# Retrosynthesis model +python training/cli_main.py test --config=training/test_retro.yaml +``` + +### Prediction + +```bash +python training/cli_main.py predict --config=training/predict.yaml +``` + +--- + +## Configuration + +Default values are defined in [`VanillaTransformerConfig`](../worker/transformers_model/configuration.py). Additional defaults are set directly in [`LitVanillaTransformer`](../worker/transformers_model/model.py) and [`LitSmilesDataset`](../worker/transformers_model/smiles_datamodule.py). Note that some defaults from `VanillaTransformerConfig` are overridden inside `LitVanillaTransformer`, so always verify which value is actually used at runtime. + +Lightning saves the configuration used for each run to `hparams.yaml` inside the corresponding `lightning_logs//` folder. + +### Passing options inline + +Any option from a YAML file can also be passed directly on the command line: + +**Model flags** — forwarded to `LitVanillaTransformer.__init__`: +``` +--model.vocab_path=vocab/vocab.txt +--model.learning_rate=0.0002 +--model.task=forward +``` + +**Data flags** — forwarded to `LitSmilesDataset.__init__`: +``` +--data.vocab_path=vocab/vocab.txt +--data.train_path=data/pistachio/data.forward/train.jsonl +--data.validation_path=data/pistachio/data.forward/validation.jsonl +--data.num_dataloader_workers=16 +--data.batch_size=256 +``` + +**Trainer flags** — forwarded to the Lightning `Trainer`: +``` +--trainer.max_epochs=10 +--trainer.accelerator=gpu +--trainer.devices=1 +--trainer.precision=32 +--trainer.enable_progress_bar=false +``` + +**Global flags**: +``` +--seed_everything=42 +--config=training/train.yaml +``` + +### Example YAML with callbacks + +More complex options, such as callbacks, are easier to specify via a YAML file: + +```yaml +model: + vocab_path: vocab/vocab.txt + learning_rate: 0.0002 +data: + vocab_path: vocab/vocab.txt + train_path: data/pistachio/data.forward/train.jsonl + validation_path: data/pistachio/data.forward/validation.jsonl + num_dataloader_workers: 16 + batch_size: 256 +trainer: + max_epochs: 10 + accelerator: gpu + enable_progress_bar: false + devices: 1 + precision: 32 + callbacks: + - class_path: lightning.pytorch.callbacks.ModelCheckpoint + init_args: + monitor: val_accuracy + mode: max + save_top_k: 3 + filename: "{epoch:02d}-{val_accuracy:.4f}" +``` + +--- + +## Requirements + +`requirements.txt` covers all dependencies for both model training and the Pistachio data pipeline: + +| Package | Used by | +|---------|---------| +| `lightning` | LightningCLI, model and data modules | +| `pandas` | Pistachio data pipeline (steps 7 and 8) | +| `python-Levenshtein` | Edit-distance metrics during validation and testing | +| `rdkit` | SMILES validation and Tanimoto similarity during training | +| `rxn-chem-utils` | SMILES tokenization regex pattern | +| `torch` | Model training | +| `transformers` | Tokenizer, scheduler, model base classes | + + +--- + +## Improving training throughput + +- **`num_dataloader_workers`** — Number of CPU workers for data loading. 16 has shown the best results in practice. +- **`batch_size`** — Samples per training step. Increasing from the default 32 to 256 (or higher) significantly reduces wall time per epoch. +- **`accumulate_grad_batches`** — Simulates a larger effective batch size without increasing GPU memory usage. +- **`precision`** — `32` is the stable default. `16-mixed` can speed up training but has shown instability (NaN loss values) in practice. +- **`profiler`** — Set `trainer.profiler: simple` to identify time bottlenecks across training steps. diff --git a/model-training/pistachio/README.md b/model-training/pistachio/README.md new file mode 100644 index 0000000000000000000000000000000000000000..e3856a12ae4658348175a9b90a38fb3b02080b73 --- /dev/null +++ b/model-training/pistachio/README.md @@ -0,0 +1,133 @@ +# Pistachio Data Pipeline + +This directory contains a sequence of scripts that download and process [Pistachio](https://www.nextmovesoftware.com/pistachio.html) reaction data into the JSONL format expected by the forward and retrosynthesis model training pipelines. + +## Prerequisites + +- **Pistachio license** — Access to the Pistachio release archives requires a license from [NextMove Software](https://www.nextmovesoftware.com/). Obtain your credentials before running step 1. +- **`rxn-reaction-preprocessing`** — Provides the `rxn-data-pipeline` CLI used in steps 4 and 6. Install it from [rxn4chemistry/rxn-reaction-preprocessing](https://github.com/rxn4chemistry/rxn-reaction-preprocessing/). +- **Python dependencies** — `pandas` is required for steps 7 and 8. + +> **Working directory**: unless noted otherwise, all scripts assume they are run from the root of the repository, and data is written under `data/pistachio/`. + +--- + +## Step 1 — Download the Pistachio archive + +Edit `step1_download.sh` to set the desired release version and replace `` and `` with your NextMove Software credentials, then run: + +```bash +bash model-training/pistachio/step1_download.sh +``` + +This downloads `pistachio.tar.gz` for the configured release into the current directory. The `-C -` flag enables resuming an interrupted download. + +--- + +## Step 2 — Extract reaction SMILES + +```bash +python model-training/pistachio/step2_get_reactions.py +``` + +Iterates over every record in `pistachio.tar.gz` and writes one reaction SMILES per line to `reactions.txt`. + +The input and output paths can be overridden if needed: + +```bash +python model-training/pistachio/step2_get_reactions.py \ + --input path/to/pistachio.tar.gz \ + --output path/to/reactions.txt +``` + +Both a `.tar.gz` archive and an already-extracted directory of JSON files are accepted as input. + +--- + +## Step 3 — Remove extended SMILES annotations + +```bash +python model-training/pistachio/step3_clean_reactions.py +``` + +Some reaction SMILES contain extended SMILES notation blocks such as `|&1:4,24,f:2.4|` or `|f:1.2|` that are not handled correctly by downstream tools. This script strips everything after the first `|` character on each line and writes the result to `reactions_cleaned.txt`. + +--- + +## Step 4 — Standardize reactions (IMPORT + STANDARDIZE) + +```bash +rxn-data-pipeline \ + --config-dir model-training/pistachio \ + --config-name step4_standardize_reactions_config.yaml +``` + +Runs the **IMPORT** and **STANDARDIZE** stages of the `rxn-data-pipeline` on `reactions_cleaned.txt`. After completion, the pipeline directory (`data/pistachio/pipeline/`) will contain: + +| File | Description | +|------|-------------| +| `data.imported.csv` | Raw reactions imported into the pipeline CSV format | +| `data.standardized.csv` | Reactions after SMILES standardization | + +> The full pipeline will fail at this point because a subset of standardized reactions contain coordinate-bond arrow annotations such as `->[Cu+2]<-`. Step 5 removes these before the pipeline is resumed. + +--- + +## Step 5 — Remove arrow annotations + +```bash +bash model-training/pistachio/step5_remove_arrow_annotations.sh +``` + +Filters `data.standardized.csv` to remove any rows whose reaction SMILES contain arrow-bond annotations (`->` or `<-`). The removed rows are saved to `arrow.forward.csv` and `arrow.backward.csv` for inspection. The cleaned output is written to `arrow.removed.csv`, which is the input for the next step. + +--- + +## Step 6 — Preprocess and split reactions (PREPROCESS + SPLIT) + +```bash +rxn-data-pipeline \ + --config-dir model-training/pistachio \ + --config-name step6_preprocess_reactions_config.yaml +``` + +Resumes the pipeline from `arrow.removed.csv`, running the **PREPROCESS** and **SPLIT** stages. After completion, the pipeline directory will contain: + +| File | Description | +|------|-------------| +| `data.processed.train.csv` | Training split | +| `data.processed.test.csv` | Test split | +| `data.processed.validation.csv` | Validation split | + +The split ratio is set to 5 % for test and validation (see `step6_preprocess_reactions_config.yaml`). + +--- + +## Step 7 — Convert splits to JSONL + +```bash +python model-training/pistachio/step7_convert_to_jsonl.py +``` + +Reads the three CSV splits and writes them as JSONL files in the format expected by the model trainer (`{"source": "...", "target": "..."}`). Two sets of output files are produced — one for each task direction: + +| Directory | Task | source | target | +|-----------|------|--------|--------| +| `data/pistachio/data.forward/` | Forward prediction | reactants | products | +| `data/pistachio/data.retro/` | Retrosynthesis | products | reactants | + +Each directory contains `train.jsonl`, `test.jsonl`, and `validation.jsonl`. + +--- + +## Step 8 — Generate model vocabulary + +The vocabulary can be built with either the forward or retro data, from its training and validation splits, since the data is the same with only the source and target being switched. + +```bash +python model-training/pistachio/step8_generate_vocab.py \ + --data-dir data/pistachio/data.forward \ + --output data/pistachio/vocab.txt +``` + +The script tokenizes all SMILES strings in `train.jsonl` and `validation.jsonl`, counts every token (bracketed atoms, two-character elements such as `Cl`/`Br`, `%nn` ring-closure labels, and individual characters), and writes a `vocab.txt` file with reserved BERT tokens prepended and all remaining tokens ordered by frequency. diff --git a/model-training/pistachio/step1_download.sh b/model-training/pistachio/step1_download.sh new file mode 100644 index 0000000000000000000000000000000000000000..a5365100d961e2ce706014d63b68ed046a30e079 --- /dev/null +++ b/model-training/pistachio/step1_download.sh @@ -0,0 +1,5 @@ +#!/bin/bash + +BASE_URL=https://www.nextmovesoftware.com/downloads/pistachio/releases/ +RELEASE=2025Q2 +curl -L -C - --user : $BASE_URL/$RELEASE/data/pistachio.tar.gz -o pistachio.tar.gz diff --git a/model-training/pistachio/step2_get_reactions.py b/model-training/pistachio/step2_get_reactions.py new file mode 100644 index 0000000000000000000000000000000000000000..181c2615136c2629f2112c94384e8ec131d109be --- /dev/null +++ b/model-training/pistachio/step2_get_reactions.py @@ -0,0 +1,107 @@ +""" +Extract reaction SMILES from a Pistachio archive. + +This script is self-contained: all logic that was previously provided by the +rxn_pistachio package has been inlined here. No external dependencies beyond +the Python standard library are required. + +Usage: + python step2_get_reactions.py # defaults below + python step2_get_reactions.py --input my.tar.gz --output reactions.txt +""" + +import argparse +import json +import logging +import tarfile +from pathlib import Path +from typing import Any, Dict, Iterator, Union + +# --------------------------------------------------------------------------- +# Pistachio archive navigation (inlined from rxn_pistachio) +# --------------------------------------------------------------------------- + +def _is_json(path: Union[str, Path]) -> bool: + path_str = str(path) + return path_str.endswith(".json") or path_str.endswith(".JSON") + + +def _iter_reaction_dicts(archive_path: Path) -> Iterator[Dict[str, Any]]: + """ + Iterate over every reaction JSON object inside a Pistachio .tar.gz file + or a directory tree of JSON files. + """ + if archive_path.is_dir(): + for json_file in (p for p in archive_path.rglob("*") if _is_json(p)): + print(f'Reading file "{json_file}"') + with open(json_file, "rt") as f: + for line in f: + line = line.strip() + if line: + yield json.loads(line) + else: + with tarfile.open(archive_path, "r:gz") as tar: + for member in tar.getmembers(): + if not _is_json(member.name): + continue + extracted = tar.extractfile(member) + if extracted is None: + continue + for raw_line in extracted: + line = raw_line.strip() + if line: + yield json.loads(line) + + +def _reaction_smiles(record: Dict[str, Any]) -> str: + """ + Return the reaction SMILES string from a Pistachio record dictionary. + Mirrors Record.reaction_smiles from rxn_pistachio. + """ + data = record.get("data", {}) + if "reactionSmiles" in data: + return data["reactionSmiles"] + return data.get("smiles", "") + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + +def extract_reactions(input_path: str, output_path: str) -> None: + archive = Path(input_path) + if not archive.exists(): + raise FileNotFoundError(f"Input not found: {archive}") + + written = 0 + skipped = 0 + + with open(output_path, "w") as out: + for record in _iter_reaction_dicts(archive): + smiles = _reaction_smiles(record) + if smiles: + out.write(smiles + "\n") + written += 1 + else: + skipped += 1 + + print(f"Done. Written: {written}, skipped (no SMILES): {skipped}") + print(f"Output: {output_path}") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Extract reaction SMILES from a Pistachio archive." + ) + parser.add_argument( + "--input", + default="pistachio.tar.gz", + help="Path to pistachio.tar.gz or extracted directory (default: pistachio.tar.gz)", + ) + parser.add_argument( + "--output", + default="reactions.txt", + help="Output file path (default: reactions.txt)", + ) + args = parser.parse_args() + extract_reactions(args.input, args.output) diff --git a/model-training/pistachio/step3_clean_reactions.py b/model-training/pistachio/step3_clean_reactions.py new file mode 100644 index 0000000000000000000000000000000000000000..7565af757be35c89ed030d2fa97741c491a2c5f4 --- /dev/null +++ b/model-training/pistachio/step3_clean_reactions.py @@ -0,0 +1,10 @@ + +input_file = 'reactions.txt' +output_file = 'reactions_cleaned.txt' + +with open(input_file, 'r') as infile, open(output_file, 'w') as outfile: + for line in infile: + cleaned_line = line.split('|')[0].strip() + outfile.write(cleaned_line + '\n') + +print(f"Cleaned SMILES reactions saved to '{output_file}'") diff --git a/model-training/pistachio/step4_standardize_reactions_config.yaml b/model-training/pistachio/step4_standardize_reactions_config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..a85da698a81a25a796a18d10e57153ba7d54f2b4 --- /dev/null +++ b/model-training/pistachio/step4_standardize_reactions_config.yaml @@ -0,0 +1,31 @@ +data: + path: data/pistachio/reactions_cleaned.txt + name: data + proc_dir: data/pistachio/pipeline +common: + sequence: + - IMPORT + - STANDARDIZE + fragment_bond: TILDE + reaction_column_name: rxn + keep_intermediate_columns: false +rxn_import: + input_file: data/pistachio/reactions_cleaned.txt + output_csv: data/pistachio/pipeline/data.imported.csv + data_format: TXT + input_csv_column_name: rxn + reaction_column_name: rxn + fragment_bond: TILDE + remove_atom_mapping: true + column_for_light: null + column_for_heat: null + keep_original_rxn_column: false +standardize: + input_file_path: data/pistachio/pipeline/data.imported.csv + annotation_file_paths: [] + discard_unannotated_metals: false + output_file_path: data/pistachio/pipeline/data.standardized.csv + fragment_bond: TILDE + reaction_column_name: rxn + remove_stereo_if_not_defined_in_precursors: false + keep_intermediate_columns: false diff --git a/model-training/pistachio/step5_remove_arrow_annotations.sh b/model-training/pistachio/step5_remove_arrow_annotations.sh new file mode 100644 index 0000000000000000000000000000000000000000..429a04c9a5ca3271270048aad969163d0274f267 --- /dev/null +++ b/model-training/pistachio/step5_remove_arrow_annotations.sh @@ -0,0 +1,19 @@ +#!/bin/bash + +BASE_PATH=data/pistachio/pipeline +FILEPATH=$BASE_PATH/data.standardized.csv + +# Extract the header +head -n 1 $FILEPATH > $BASE_PATH/arrow.backward.csv +head -n 1 $FILEPATH > $BASE_PATH/arrow.forward.csv +head -n 1 $FILEPATH > $BASE_PATH/arrow.removed.csv +head -n 1 $FILEPATH > $BASE_PATH/tmp.csv + +# Append matching lines (excluding header) to each file +tail -n +2 $FILEPATH | grep '<' >> $BASE_PATH/arrow.backward.csv +tail -n +2 $FILEPATH | grep -v '<' >> $BASE_PATH/tmp.csv +tail -n +2 $BASE_PATH/tmp.csv | grep "\\->" >> $BASE_PATH/arrow.forward.csv +tail -n +2 $BASE_PATH/tmp.csv | grep -v "\\->" >> $BASE_PATH/arrow.removed.csv + +# Remove auxiliary csv +rm $BASE_PATH/tmp.csv diff --git a/model-training/pistachio/step6_preprocess_reactions_config.yaml b/model-training/pistachio/step6_preprocess_reactions_config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..baa32ce74aa5404882d30338fe785281330bc6cb --- /dev/null +++ b/model-training/pistachio/step6_preprocess_reactions_config.yaml @@ -0,0 +1,35 @@ +data: + path: data/pistachio/reactions_cleaned.txt + name: data + proc_dir: data/pistachio/pipeline +common: + sequence: + - PREPROCESS + - SPLIT + fragment_bond: TILDE + reaction_column_name: rxn + keep_intermediate_columns: false +preprocess: + input_file_path: data/pistachio/pipeline/arrow.removed.csv + output_file_path: data/pistachio/pipeline/data.processed.csv + min_reactants: 1 + max_reactants: 10 + max_reactants_tokens: 254 + min_agents: 0 + max_agents: 0 + max_agents_tokens: 0 + min_products: 1 + max_products: 10 + max_products_tokens: 254 + max_absolute_formal_charge: 2 + fragment_bond: TILDE + reaction_column_name: rxn + keep_intermediate_columns: false +split: + input_file_path: data/pistachio/pipeline/data.processed.csv + output_directory: data/pistachio/pipeline + split_ratio: 0.05 + reaction_column_name: rxn + index_column: products + hash_seed: 42 + shuffle_seed: 42 diff --git a/model-training/pistachio/step7_convert_to_jsonl.py b/model-training/pistachio/step7_convert_to_jsonl.py new file mode 100644 index 0000000000000000000000000000000000000000..dfb4036043f9403facf122ebfb681573a985cb24 --- /dev/null +++ b/model-training/pistachio/step7_convert_to_jsonl.py @@ -0,0 +1,46 @@ +import pandas as pd +import os +import json + +# Define base paths +base_path = "data/pistachio" +data_path = os.path.join(base_path, "pipeline") + +# Define input CSV files +csv_files = { + "train": os.path.join(data_path, "data.processed.train.csv"), + "test": os.path.join(data_path, "data.processed.test.csv"), + "validation": os.path.join(data_path, "data.processed.validation.csv") +} + +# Define output directories +forward_dir = os.path.join(base_path, "data.forward") +retro_dir = os.path.join(base_path, "data.retro") +os.makedirs(forward_dir, exist_ok=True) +os.makedirs(retro_dir, exist_ok=True) + +# Process each dataset +for split, csv_file in csv_files.items(): + # Read CSV file + df = pd.read_csv(csv_file, header=0, names=["reactants", "products"], sep=">>", engine="python") + + # Escape backslashes + df["reactants"] = df["reactants"].astype(str) + df["products"] = df["products"].astype(str) + + # Write forward JSONL + forward_file = os.path.join(forward_dir, f"{split}.jsonl") + with open(forward_file, "w") as f_out: + for _, row in df.iterrows(): + json.dump({"source": row["reactants"], "target": row["products"]}, f_out) + f_out.write("\n") + + # Write retro JSONL + retro_file = os.path.join(retro_dir, f"{split}.jsonl") + with open(retro_file, "w") as f_out: + for _, row in df.iterrows(): + json.dump({"source": row["products"], "target": row["reactants"]}, f_out) + f_out.write("\n") + +print("All JSONL files created successfully.") + diff --git a/model-training/pistachio/step8_generate_vocab.py b/model-training/pistachio/step8_generate_vocab.py new file mode 100644 index 0000000000000000000000000000000000000000..0b25c8da755108a1582b25fb16d86b232843ae7d --- /dev/null +++ b/model-training/pistachio/step8_generate_vocab.py @@ -0,0 +1,83 @@ +import json +import os +import re +from collections import Counter + +# Reserved tokens +reserved_tokens = [ + "[PAD]", "[unused1]", "[unused2]", "[unused3]", "[unused4]", "[unused5]", + "[unused6]", "[unused7]", "[unused8]", "[unused9]", "[unused10]", "[UNK]", + "[CLS]", "[SEP]", "[MASK]" +] + +def count_characters(file_path): + char_counter = Counter() + with open(file_path, 'r') as file: + for line in file: + data = json.loads(line) + source = data.get('source', '') + target = data.get('target', '') + combined = source + target + + # Extract and count bracketed tokens + bracket_tokens = re.findall(r'\[.*?\]', combined) + char_counter.update(bracket_tokens) + + # Remove bracketed tokens + combined = re.sub(r'\[.*?\]', '', combined) + + # Extract and count composite % tokens (e.g., %10, %11) + percent_tokens = re.findall(r'%\d{2}', combined) + char_counter.update(percent_tokens) + + # Remove composite % tokens + combined = re.sub(r'%\d{2}', '', combined) + + # Tokenize remaining string + tokens = [] + i = 0 + while i < len(combined): + if combined[i:i+2] in ('Cl', 'Br'): + tokens.append(combined[i:i+2]) + i += 2 + else: + tokens.append(combined[i]) + i += 1 + + char_counter.update(tokens) + + return char_counter + +def generate_vocab(train_file, val_file, output_file): + train_counter = count_characters(train_file) + val_counter = count_characters(val_file) + total_counter = train_counter + val_counter + sorted_tokens = [token for token, _ in total_counter.most_common()] + vocab = reserved_tokens + sorted_tokens + with open(output_file, 'w') as file: + for token in vocab: + file.write(token + '\n') + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser( + description="Generate a vocabulary file from forward or retro JSONL splits." + ) + parser.add_argument( + "--data-dir", + required=True, + help="Directory containing train.jsonl and validation.jsonl (e.g. data/pistachio/data.forward.256)", + ) + parser.add_argument( + "--output", + default="vocab.txt", + help="Output vocabulary file path (default: vocab.txt)", + ) + args = parser.parse_args() + + train_file = os.path.join(args.data_dir, "train.jsonl") + val_file = os.path.join(args.data_dir, "validation.jsonl") + generate_vocab(train_file, val_file, args.output) + print(f"Vocabulary written to '{args.output}'") diff --git a/model-training/requirements.txt b/model-training/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..11a5aac104442dfbd6210295e128d2ce55d7c5a0 --- /dev/null +++ b/model-training/requirements.txt @@ -0,0 +1,7 @@ +lightning +pandas +python-Levenshtein +rdkit +rxn-chem-utils +torch +transformers diff --git a/model-training/tools/analyze_explicit_hydrogens.py b/model-training/tools/analyze_explicit_hydrogens.py new file mode 100644 index 0000000000000000000000000000000000000000..5dd92e831f2821e014be0fd4841ff69ed12c97cc --- /dev/null +++ b/model-training/tools/analyze_explicit_hydrogens.py @@ -0,0 +1,122 @@ +#!/usr/bin/env python3 +""" +Script to analyze reactions with explicit hydrogens and check for organometallic elements. +Usage: python check_explicit_hydrogens.py +""" + +import sys +import re +from collections import Counter + +# Common organometallic and metalloid elements +METALS = { + 'Li', 'Na', 'K', 'Mg', 'Ca', 'Sc', 'Ti', 'V', 'Cr', 'Mn', 'Fe', 'Co', 'Ni', + 'Cu', 'Zn', 'Al', 'Ga', 'Ge', 'Sn', 'Pb', 'B', 'Si', 'P', 'Pd', 'Pt', 'Rh', + 'Ru', 'Os', 'Ir', 'Au', 'Ag', 'Cd', 'Hg', 'Zr', 'Mo', 'W', 'Re', 'Bi' +} + +def has_explicit_hydrogens(smiles): + """Check if SMILES contains explicit hydrogens like [CH], [CH2], [CH3]""" + return bool(re.search(r'\[CH\d?\]', smiles)) + +def extract_elements(smiles): + """Extract all elements in square brackets from SMILES""" + # Match elements in square brackets, including charges and other notation + pattern = r'\[([A-Z][a-z]?)' + elements = re.findall(pattern, smiles) + return set(elements) + +def main(): + if len(sys.argv) != 2: + print("Usage: python check_explicit_hydrogens.py ") + sys.exit(1) + + input_file = sys.argv[1] + + try: + with open(input_file, 'r') as f: + lines = f.readlines() + except FileNotFoundError: + print(f"Error: File {input_file} not found") + sys.exit(1) + + # Skip header if present + if lines and lines[0].strip().lower() in ['rxn', 'reaction', 'smiles']: + lines = lines[1:] + + total_reactions = len(lines) + reactions_with_explicit_h = [] + reactions_with_explicit_h_and_metals = [] + reactions_with_explicit_h_no_metals = [] + metal_counter = Counter() + + print("Analyzing reactions...") + print("=" * 60) + + for line in lines: + line = line.strip() + if not line: + continue + + if has_explicit_hydrogens(line): + reactions_with_explicit_h.append(line) + elements = extract_elements(line) + metals_found = elements & METALS + + if metals_found: + reactions_with_explicit_h_and_metals.append(line) + for metal in metals_found: + metal_counter[metal] += 1 + else: + reactions_with_explicit_h_no_metals.append(line) + + # Print statistics + print(f"\nTotal reactions: {total_reactions}") + print(f"Reactions with explicit hydrogens: {len(reactions_with_explicit_h)}") + print(f" - With organometallic elements: {len(reactions_with_explicit_h_and_metals)}") + print(f" - Without organometallic elements: {len(reactions_with_explicit_h_no_metals)}") + + if reactions_with_explicit_h: + percentage = (len(reactions_with_explicit_h_and_metals) / len(reactions_with_explicit_h)) * 100 + print(f"\nPercentage with metals: {percentage:.2f}%") + + # Show metal breakdown + if metal_counter: + print("\nBreakdown by metal type (for reactions with explicit H):") + print("=" * 60) + for metal, count in metal_counter.most_common(): + print(f" {metal}: {count} reactions") + + # Show examples without metals + if reactions_with_explicit_h_no_metals: + print(f"\nFirst 10 examples with explicit H but NO metals:") + print("=" * 60) + for i, rxn in enumerate(reactions_with_explicit_h_no_metals[:10], 1): + # Truncate long reactions for display + display_rxn = rxn if len(rxn) <= 100 else rxn[:97] + "..." + print(f"{i}. {display_rxn}") + + # Additional analysis: check for specific patterns + print("\n\nAdditional Analysis:") + print("=" * 60) + + # Check for radicals (indicated by explicit H on single atoms) + radical_pattern = r'\[CH?\](?![0-9])' # [C] or [CH] not followed by a digit + radicals = [rxn for rxn in reactions_with_explicit_h_no_metals if re.search(radical_pattern, rxn)] + print(f"Potential radicals (no metals): {len(radicals)}") + + # Check for charged species + charged_pattern = r'\[CH\d?[+-]\]' + charged = [rxn for rxn in reactions_with_explicit_h_no_metals if re.search(charged_pattern, rxn)] + print(f"Charged species with explicit H (no metals): {len(charged)}") + + if radicals and len(radicals) <= 5: + print("\nExamples of potential radicals:") + for rxn in radicals[:5]: + display_rxn = rxn if len(rxn) <= 100 else rxn[:97] + "..." + print(f" {display_rxn}") + +if __name__ == "__main__": + main() + +# Made with Bob diff --git a/model-training/tools/analyze_tokens.py b/model-training/tools/analyze_tokens.py new file mode 100644 index 0000000000000000000000000000000000000000..28371adb0157e40015282b846693a3732882fc41 --- /dev/null +++ b/model-training/tools/analyze_tokens.py @@ -0,0 +1,106 @@ +""" +Analyse token lengths across the train/test/validation JSONL splits. + +Reports the maximum token count for source and target sequences and how +many samples exceed the configured max_length threshold. + +Usage: + python analyze_tokens.py --data-dir data/pistachio/data.forward --vocab vocab/vocab.txt + python analyze_tokens.py --data-dir data/pistachio/data.forward --vocab vocab/vocab.txt --max-length 512 +""" + +import argparse +import json +import os +import sys + +# Allow running from any directory inside model-training/ +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) + +from transformers_model.smiles_tokenizer import SmilesTokenizer + +SPLITS = ["train", "test", "validation"] +BATCH_SIZE = 1000 + + +def analyse(data_dir: str, vocab_path: str, max_length: int) -> None: + tokenizer = SmilesTokenizer(vocab_path) + + source_max_tokens = 0 + target_max_tokens = 0 + source_over_limit = 0 + target_over_limit = 0 + total_samples = 0 + + for split in SPLITS: + file_path = os.path.join(data_dir, f"{split}.jsonl") + if not os.path.exists(file_path): + print(f" [skip] {file_path} not found") + continue + + print(f" Processing {file_path} ...") + with open(file_path, "r", encoding="utf-8") as f: + batch_sources = [] + batch_targets = [] + + def process_batch(sources, targets): + nonlocal source_max_tokens, target_max_tokens + nonlocal source_over_limit, target_over_limit + + source_encodings = tokenizer(sources, add_special_tokens=True, truncation=False) + target_encodings = tokenizer(targets, add_special_tokens=True, truncation=False) + + for s_ids, t_ids in zip(source_encodings["input_ids"], target_encodings["input_ids"]): + s_len = len(s_ids) + t_len = len(t_ids) + source_max_tokens = max(source_max_tokens, s_len) + target_max_tokens = max(target_max_tokens, t_len) + if s_len > max_length: + source_over_limit += 1 + if t_len > max_length: + target_over_limit += 1 + + for line in f: + data = json.loads(line) + batch_sources.append(data["source"]) + batch_targets.append(data["target"]) + total_samples += 1 + + if len(batch_sources) == BATCH_SIZE: + process_batch(batch_sources, batch_targets) + batch_sources = [] + batch_targets = [] + + if batch_sources: + process_batch(batch_sources, batch_targets) + + print() + print(f"Max length threshold : {max_length}") + print(f"Max tokens in source : {source_max_tokens}") + print(f"Max tokens in target : {target_max_tokens}") + print(f"Source over limit : {source_over_limit} / {total_samples} ({(source_over_limit / total_samples) * 100:.2f}%)" if total_samples else "No samples processed.") + print(f"Target over limit : {target_over_limit} / {total_samples} ({(target_over_limit / total_samples) * 100:.2f}%)" if total_samples else "") + print(f"Total samples : {total_samples}") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Analyse token lengths in JSONL splits.") + parser.add_argument( + "--data-dir", + required=True, + help="Directory containing train.jsonl, test.jsonl, validation.jsonl " + "(e.g. data/pistachio/data.forward)", + ) + parser.add_argument( + "--vocab", + required=True, + help="Path to vocab.txt (e.g. data/pistachio/vocab.txt)", + ) + parser.add_argument( + "--max-length", + type=int, + default=256, + help="Token length threshold to count over-limit samples (default: 256)", + ) + args = parser.parse_args() + analyse(args.data_dir, args.vocab, args.max_length) diff --git a/model-training/tools/check_mol_from_smiles.py b/model-training/tools/check_mol_from_smiles.py new file mode 100644 index 0000000000000000000000000000000000000000..bf9fdd92b236870174b814d15eccc6a2f4904847 --- /dev/null +++ b/model-training/tools/check_mol_from_smiles.py @@ -0,0 +1,10 @@ +from rdkit import Chem, DataStructs +from rdkit.Chem.rdFingerprintGenerator import GetMorganGenerator + + +smiles = "O=S(=O)([O-])c1ccccc1~O=S(=O)([O-])c1ccccc1~O=S(=O)([O-])c1ccccc1~O=S(=O)([O-])c1ccccc1~O=S(=O)([O-])c1ccccc1~[Fe+2]~[Fe+2]~[Fe+2]~[Fe+2]~[Fe+2]~[Fe+2]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~[Fe+3]~" + +if __name__ == "__main__": + print(f"Smiles: {smiles}") + mol1 = Chem.MolFromSmiles(smiles) + print(f"Molecule: {mol1}") \ No newline at end of file diff --git a/model-training/tools/check_tanimoto.py b/model-training/tools/check_tanimoto.py new file mode 100644 index 0000000000000000000000000000000000000000..ba08e2985b4bbffd6eced95113ac175f330b57aa --- /dev/null +++ b/model-training/tools/check_tanimoto.py @@ -0,0 +1,46 @@ +from rdkit import Chem, DataStructs +from rdkit.Chem.rdFingerprintGenerator import GetMorganGenerator + + +predicted_texts = [ + "", +] + +target_texts = [ + "", +] + + +if __name__ == "__main__": + similarities = [] + invalid_pairs = 0 + generator = GetMorganGenerator(radius=2, fpSize=2048) + for pred, tgt in zip(predicted_texts, target_texts): + + # Get molecule from prediction + print(f"Predicted smiles (pred): {pred}") + mol1 = Chem.MolFromSmiles(pred) + print(f"Predicted molecule (mol1): {mol1}") + + # Get molecule from ground truth + print(f"Target smiles (tgt): {tgt}") + mol2 = Chem.MolFromSmiles(tgt) + print(f"Target molecule (mol2): {mol2}") + + # Increment invalid pairs on failure + if mol1 is None or mol2 is None: + invalid_pairs += 1 + print(f"Invalid pair! Total: {invalid_pairs}") + continue + + + fp1 = generator.GetFingerprint(mol1) + print(f"Predicted fingerprint (fp1): {fp1}") + fp2 = generator.GetFingerprint(mol2) + print(f"Target fingerprint (fp2): {fp2}") + sim = DataStructs.TanimotoSimilarity(fp1, fp2) + print(f"Similarity (sim): {sim}") + similarities.append(sim) + + avg_tanimoto = sum(similarities) / len(similarities) if len(similarities) > 0 else 0 + print(f"Average Tanimoto (avg_tanimoto): {avg_tanimoto}") diff --git a/model-training/tools/run_fit.py b/model-training/tools/run_fit.py new file mode 100644 index 0000000000000000000000000000000000000000..56eb9925819fc73f838709ec53a29c4c2237756d --- /dev/null +++ b/model-training/tools/run_fit.py @@ -0,0 +1,50 @@ +""" +Train the model directly with Lightning, bypassing LightningCLI. + +Useful for quick experiments where you want to set options in code rather +than via a YAML config. For production training use cli_main.py with the +YAML configs in training/. + +Usage (run from model-training/): + python debug/run_fit.py +""" + +import os +import sys + +# Allow running from any directory inside model-training/ +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) + +import lightning as L + +from transformers_model.smiles_datamodule import LitSmilesDataset +from transformers_model.model import LitVanillaTransformer + + +if __name__ == "__main__": + L.seed_everything(42) + + smiles_dataset = LitSmilesDataset( + vocab_path="vocab/vocab.txt", + train_path="data/pistachio/data.forward/train.jsonl", + validation_path="data/pistachio/data.forward/validation.jsonl", + num_dataloader_workers=16, + batch_size=256, + max_length=256, + ) + + model = LitVanillaTransformer( + vocab_path="vocab/vocab.txt", + task="forward", + max_length=256, + ) + + trainer = L.Trainer( + max_epochs=10, + accelerator="gpu", + devices=1, + precision=32, + ) + + trainer.fit(model, datamodule=smiles_dataset) + print("Training complete.") diff --git a/model-training/tools/run_lr_finder.py b/model-training/tools/run_lr_finder.py new file mode 100644 index 0000000000000000000000000000000000000000..bb2887128ef3689d107b0e672e5a03b6491b29e5 --- /dev/null +++ b/model-training/tools/run_lr_finder.py @@ -0,0 +1,98 @@ +""" +Runs the Lightning learning rate finder and saves a plot and CSV of the results. + +Output files are written to the same directory as this script: + - lr_finder_plot.png + - lr_finder_results.csv + +Usage (run from model-training/): + python debug/run_lr_finder.py + python debug/run_lr_finder.py --data-dir data/pistachio/data.forward --vocab vocab/vocab.txt +""" + +import argparse +import csv +import os +import sys +import torch + +# Allow running from any directory inside model-training/ +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) + +import lightning as L +import matplotlib.pyplot as plt +from lightning.pytorch.tuner import Tuner + +from transformers_model.smiles_datamodule import LitSmilesDataset +from transformers_model.model import LitVanillaTransformer + +HERE = os.path.dirname(os.path.abspath(__file__)) + + +def run(data_dir: str, vocab_path: str) -> None: + L.seed_everything(42) + torch.set_float32_matmul_precision("high") + + smiles_dataset = LitSmilesDataset( + vocab_path=vocab_path, + train_path=os.path.join(data_dir, "train.jsonl"), + validation_path=os.path.join(data_dir, "validation.jsonl"), + num_dataloader_workers=16, + batch_size=1024, + max_length=256, + ) + + model = LitVanillaTransformer( + vocab_path=vocab_path, + task="forward", + max_length=256, + ignore_nan=True, + ) + + trainer = L.Trainer( + max_epochs=60, + accelerator="gpu", + devices=1, + precision="16-mixed", + deterministic=True, + gradient_clip_val=1.0, + gradient_clip_algorithm="norm", + ) + + tuner = Tuner(trainer) + lr_finder = tuner.lr_find(model, datamodule=smiles_dataset) + + # Plot + fig = lr_finder.plot(suggest=True) + plt.grid(True, which="both", linestyle="--", linewidth=0.5) + plot_path = os.path.join(HERE, "lr_finder_plot.png") + fig.savefig(plot_path) + print(f"Plot saved to {plot_path}") + plt.show() + + # Save results as CSV + csv_path = os.path.join(HERE, "lr_finder_results.csv") + with open(csv_path, "w", newline="") as f: + writer = csv.writer(f) + writer.writerow(["learning_rate", "loss"]) + writer.writerows(zip(lr_finder.results["lr"], lr_finder.results["loss"])) + print(f"Results saved to {csv_path}") + + print(f"Suggested learning rate: {lr_finder.suggestion()}") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Run the Lightning LR finder.") + parser.add_argument( + "--data-dir", + default="data/pistachio/data.forward", + help="Directory containing train.jsonl and validation.jsonl " + "(default: data/pistachio/data.forward)", + ) + parser.add_argument( + "--vocab", + default="vocab/vocab.txt", + help="Path to vocab.txt (default: vocab/vocab.txt)", + ) + args = parser.parse_args() + run(args.data_dir, args.vocab) diff --git a/model-training/tools/run_predict.py b/model-training/tools/run_predict.py new file mode 100644 index 0000000000000000000000000000000000000000..3151ec0b38a7cd79515e6258a7e9b21e90eb6199 --- /dev/null +++ b/model-training/tools/run_predict.py @@ -0,0 +1,48 @@ +""" +Run predictions with the model directly via Lightning, bypassing LightningCLI. + +Useful for quick experiments where you want to set options in code rather +than via a YAML config. For production inference use cli_main.py with +training/predict.yaml. + +Usage (run from model-training/): + python debug/run_predict.py +""" + +import os +import sys + +# Allow running from any directory inside model-training/ +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) + +import lightning as L + +from transformers_model.smiles_datamodule import LitSmilesDataset +from transformers_model.model import LitVanillaTransformer + + +if __name__ == "__main__": + L.seed_everything(42) + + smiles_dataset = LitSmilesDataset( + vocab_path="vocab/vocab.txt", + predict_path="data/input.jsonl", + num_dataloader_workers=16, + batch_size=256, + max_length=256, + ) + + # Replace with the actual checkpoint path + model = LitVanillaTransformer.load_from_checkpoint( + checkpoint_path="lightning_logs/forward/checkpoints/epoch=59-val_accuracy=0.6476.ckpt", + vocab_path="vocab/vocab.txt", + ) + + trainer = L.Trainer( + accelerator="gpu", + devices=1, + precision=32, + ) + + trainer.predict(model, datamodule=smiles_dataset) + print("Prediction complete.") diff --git a/model-training/tools/run_test.py b/model-training/tools/run_test.py new file mode 100644 index 0000000000000000000000000000000000000000..8c46b7462dc1fe931c35e2da85594333ba37dab6 --- /dev/null +++ b/model-training/tools/run_test.py @@ -0,0 +1,48 @@ +""" +Test the model directly with Lightning, bypassing LightningCLI. + +Useful for quick experiments where you want to set options in code rather +than via a YAML config. For production testing use cli_main.py with the +YAML configs in training/. + +Usage (run from model-training/): + python debug/run_test.py +""" + +import os +import sys + +# Allow running from any directory inside model-training/ +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) + +import lightning as L + +from transformers_model.smiles_datamodule import LitSmilesDataset +from transformers_model.model import LitVanillaTransformer + + +if __name__ == "__main__": + L.seed_everything(42) + + smiles_dataset = LitSmilesDataset( + vocab_path="vocab/vocab.txt", + test_path="data/pistachio/data.forward/test.jsonl", + num_dataloader_workers=16, + batch_size=256, + max_length=256, + ) + + # Replace with the actual checkpoint path + model = LitVanillaTransformer.load_from_checkpoint( + checkpoint_path="lightning_logs/forward/checkpoints/epoch=59-val_accuracy=0.6476.ckpt", + vocab_path="vocab/vocab.txt", + ) + + trainer = L.Trainer( + accelerator="gpu", + devices=1, + precision=32, + ) + + trainer.test(model, datamodule=smiles_dataset) + print("Testing complete.") diff --git a/model-training/training/cli_main.py b/model-training/training/cli_main.py new file mode 100644 index 0000000000000000000000000000000000000000..c52b1edce640ba40af1ba616860b6163be6224ba --- /dev/null +++ b/model-training/training/cli_main.py @@ -0,0 +1,20 @@ +import torch + +from lightning.pytorch.cli import LightningCLI +from transformers_model.model import LitVanillaTransformer +from transformers_model.smiles_datamodule import LitSmilesDataset + + +def main(): + + # Enable Tensor Core optimization + torch.set_float32_matmul_precision('high') + + LightningCLI( + model_class=LitVanillaTransformer, + datamodule_class=LitSmilesDataset, + ) + + +if __name__ == "__main__": + main() diff --git a/model-training/training/predict.yaml b/model-training/training/predict.yaml new file mode 100644 index 0000000000000000000000000000000000000000..3013a1394367952f0987a9acec6623f9a8a24ec6 --- /dev/null +++ b/model-training/training/predict.yaml @@ -0,0 +1,16 @@ +model: + vocab_path: vocab/vocab.txt + task: "forward" + predictions_path: "predictions.jsonl" + num_beams: 3 + topn: 3 + max_length: 256 # for generating predictions +data: + vocab_path: vocab/vocab.txt + predict_path: data/input.jsonl + truncation: true + max_length: 256 # for tokenizing input data + num_dataloader_workers: 16 + batch_size: 512 +seed_everything: 42 +ckpt_path: lightning_logs/forward/checkpoints/epoch=59-val_accuracy=0.6476.ckpt \ No newline at end of file diff --git a/model-training/training/test.yaml b/model-training/training/test.yaml new file mode 100644 index 0000000000000000000000000000000000000000..ef087631bc22090cae597dfc2ee92a1cc3e86dbf --- /dev/null +++ b/model-training/training/test.yaml @@ -0,0 +1,35 @@ +model: + vocab_path: vocab/vocab.txt + task: "forward" + learning_rate: 0.001584893192461114 + warmup_ratio: 0.1 + # attention_mask: -1e4 + max_length: 256 # for generating predictions +data: + vocab_path: vocab/vocab.txt + test_path: data/pistachio/data.forward/test.jsonl + truncation: true + max_length: 256 # for tokenizing input data + num_dataloader_workers: 16 + batch_size: 512 +trainer: + accumulate_grad_batches: 2 + max_epochs: 60 + # profiler: simple + accelerator: gpu + enable_progress_bar: false + devices: 1 + precision: 32 + deterministic: true + # detect_anomaly: true + gradient_clip_val: 1.0 + gradient_clip_algorithm: "norm" + callbacks: + - class_path: lightning.pytorch.callbacks.ModelCheckpoint + init_args: + monitor: val_accuracy + mode: max + save_top_k: 10 + filename: "{epoch:02d}-{val_accuracy:.4f}" +seed_everything: 42 +ckpt_path: lightning_logs/forward/checkpoints/epoch=59-val_accuracy=0.6476.ckpt \ No newline at end of file diff --git a/model-training/training/test_retro.yaml b/model-training/training/test_retro.yaml new file mode 100644 index 0000000000000000000000000000000000000000..de5450bc3d354a6b392d982a7fabe194950e7c10 --- /dev/null +++ b/model-training/training/test_retro.yaml @@ -0,0 +1,38 @@ +model: + vocab_path: vocab/vocab.txt + task: "retro" + learning_rate: 0.001584893192461114 + warmup_ratio: 0.1 + # attention_mask: -1e4 + max_length: 256 # for generating predictions + num_beams: 5 + round_trip_top_k: 5 + forward_model_path: lightning_logs/forward/checkpoints/epoch=59-val_accuracy=0.6476.ckpt +data: + vocab_path: vocab/vocab.txt + test_path: data/pistachio/data.retro/test.jsonl + truncation: true + max_length: 256 # for tokenizing input data + num_dataloader_workers: 16 + batch_size: 512 +trainer: + accumulate_grad_batches: 2 + max_epochs: 60 + # profiler: simple + accelerator: gpu + enable_progress_bar: false + devices: 1 + precision: 32 + deterministic: true + # detect_anomaly: true + gradient_clip_val: 1.0 + gradient_clip_algorithm: "norm" + callbacks: + - class_path: lightning.pytorch.callbacks.ModelCheckpoint + init_args: + monitor: val_accuracy + mode: max + save_top_k: 10 + filename: "{epoch:02d}-{val_accuracy:.4f}" +seed_everything: 42 +ckpt_path: lightning_logs/retro/checkpoints/epoch=59-val_accuracy=0.0766.ckpt \ No newline at end of file diff --git a/model-training/training/train.yaml b/model-training/training/train.yaml new file mode 100644 index 0000000000000000000000000000000000000000..49b21fe07350ef0929a75db9440bd49f7d5b393f --- /dev/null +++ b/model-training/training/train.yaml @@ -0,0 +1,35 @@ +model: + vocab_path: vocab/vocab.txt + task: "forward" + learning_rate: 0.001584893192461114 + warmup_ratio: 0.1 + # attention_mask: -1e4 + max_length: 256 # for generating predictions +data: + vocab_path: vocab/vocab.txt + train_path: data/pistachio/data.forward/train.jsonl + validation_path: data/pistachio/data.forward/validation.jsonl + truncation: true + max_length: 256 # for tokenizing input data + num_dataloader_workers: 16 + batch_size: 512 +trainer: + accumulate_grad_batches: 2 + max_epochs: 60 + # profiler: simple + accelerator: gpu + enable_progress_bar: false + devices: 1 + precision: 32 + deterministic: true + # detect_anomaly: true + gradient_clip_val: 1.0 + gradient_clip_algorithm: "norm" + callbacks: + - class_path: lightning.pytorch.callbacks.ModelCheckpoint + init_args: + monitor: val_accuracy + mode: max + save_top_k: 10 + filename: "{epoch:02d}-{val_accuracy:.4f}" +seed_everything: 42 \ No newline at end of file diff --git a/model-training/training/train_retro.yaml b/model-training/training/train_retro.yaml new file mode 100644 index 0000000000000000000000000000000000000000..59a1465994edb26dd11c0704bf7c2d2d0bd880ac --- /dev/null +++ b/model-training/training/train_retro.yaml @@ -0,0 +1,32 @@ +model: + vocab_path: vocab/vocab.txt + task: "retro" + learning_rate: 0.001584893192461114 + warmup_ratio: 0.1 + max_length: 256 # for generating predictions +data: + vocab_path: vocab/vocab.txt + train_path: data/pistachio/data.retro/train.jsonl + validation_path: data/pistachio/data.retro/validation.jsonl + truncation: true + max_length: 256 # for tokenizing input data + num_dataloader_workers: 16 + batch_size: 512 +trainer: + accumulate_grad_batches: 2 + max_epochs: 60 + accelerator: gpu + enable_progress_bar: false + devices: 1 + precision: 32 + deterministic: true + gradient_clip_val: 1.0 + gradient_clip_algorithm: "norm" + callbacks: + - class_path: lightning.pytorch.callbacks.ModelCheckpoint + init_args: + monitor: val_accuracy + mode: max + save_top_k: 10 + filename: "{epoch:02d}-{val_accuracy:.4f}" +seed_everything: 42 \ No newline at end of file diff --git a/models/forward/Pistachio2025Q2-Forward.ckpt b/models/forward/Pistachio2025Q2-Forward.ckpt new file mode 100644 index 0000000000000000000000000000000000000000..dfb0a7e920baf0ee1f2b23f44924821af5a29d24 --- /dev/null +++ b/models/forward/Pistachio2025Q2-Forward.ckpt @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:def72338009f6f82c4d64a55ed3c488cf3029e72f0a10c308d322e59257a10fe +size 155080231 diff --git a/models/retrosynthesis/Pistachio2025Q2-Retro.ckpt b/models/retrosynthesis/Pistachio2025Q2-Retro.ckpt new file mode 100644 index 0000000000000000000000000000000000000000..1e2c4a245fba581deebdeca326cf5d8adc484bf6 --- /dev/null +++ b/models/retrosynthesis/Pistachio2025Q2-Retro.ckpt @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:13d7fc03a9e4637ed19e18868de84dc86cdfc7764681604157d5a9796c30ba9c +size 155080231 diff --git a/scripts/predict_product.py b/scripts/predict_product.py new file mode 100644 index 0000000000000000000000000000000000000000..ca53c0b471e00102d1e5660ddfd483c9fefeaa3a --- /dev/null +++ b/scripts/predict_product.py @@ -0,0 +1,39 @@ +""" +Update this script and run it inside the worker container to run your product predictions via script +""" +print("Setting up product prediction...") + +# All necessary imports +from celery import Celery +from utils import wait_for_result + +# Initialize Celery +celery_app = Celery() +print("Broker:", celery_app.conf.broker_url) +print("Backend:", celery_app.conf.result_backend) + +# Set up a list of reactants to make predictions +reactants_list = ["CCI.O=Cc1ccc([N+](=O)[O-])c(O)c1", "CCOc1cc(O)c(C=O)cc1OCC.OCCCBr", "C=CCc1cc(OCc2ccccc2)ccc1O.CCBr"] + +# Setup task kwargs +kwargs = { + "topn": 1, # Number of results per reactant + "num_beams": 3, # Number of beams used for prediction. Must be >= topn + "device": None, # Device used for predicting, either "cuda" or "cpu", None defaults to cuda if available + "ckpt_forward": "Pistachio2025Q2-Forward", # Default forward model + "vocab": "Pistachio2025Q2", # Vocab for default forward model + # "ckpt_forward_path": "models/forward/Pistachio2025Q2-Forward.ckpt", # Can be used instead of ckpt_forward + # "vocab_path": "vocab/Pistachio2025Q2.txt", # Can be used instead of vocab +} + +# Send the product_prediction task with the reaction list and kwargs +task = celery_app.send_task( + "tasks.product_prediction", + [reactants_list], + kwargs=kwargs, + queue="product_prediction", +) +print("Task sent. Assigned task_id: {}".format(task.id)) + +# Use the task id to get the result. Increase timeout if needed. +wait_for_result(celery_app, task.id, timeout=180) diff --git a/scripts/predict_retrosynthesis.py b/scripts/predict_retrosynthesis.py new file mode 100644 index 0000000000000000000000000000000000000000..8d59021b450727d39d90bdb7c64a5c48b0a8fe01 --- /dev/null +++ b/scripts/predict_retrosynthesis.py @@ -0,0 +1,43 @@ +""" +Update this script and run it inside the worker container to run your retrosynthesis predictions via script +""" +print("Setting up retrosynthesis prediction...") + +# All necessary imports +from celery import Celery +from utils import wait_for_result + +# Initialize Celery +celery_app = Celery() +print("Broker:", celery_app.conf.broker_url) +print("Backend:", celery_app.conf.result_backend) + +# Choose product for retrosynthesis prediction +product = "C=CC(=C)C[Si](C)(C)C" + +# Setup task kwargs +kwargs = { + "topn": 15, # Number of results per reactant + "num_beams": 15, # Number of beams used for prediction. Must be >= topn + "fap": 0.6, # Forward likelihood acceptance probability (not length averaged) + "fld": 0.2, # Forward likelihood delta required between the top2 forward prediction results + "device": None, # Device used for predicting, either "cuda" or "cpu", None defaults to cuda if available + "ckpt_forward": "Pistachio2025Q2-Forward", # Default forward model + "ckpt_retro": "Pistachio2025Q2-Retro", # Default retrosynthesis model + "vocab": "Pistachio2025Q2", # Vocab for default forward and retrosynthesis models + # "ckpt_forward_path": "models/forward/Pistachio2025Q2-Forward.ckpt", # Can be used instead of ckpt_forward + # "ckpt_retro_path": "models/retrosynthesis/Pistachio2025Q2-Retro.ckpt", # Can be used instead of ckpt_retro + # "vocab_path": "vocab/Pistachio2025Q2.txt", # Can be used instead of vocab +} + +# Send the retro_prediction task with the product and kwargs +task = celery_app.send_task( + "tasks.retro_prediction", + [product], + kwargs=kwargs, + queue="retro_prediction", +) +print("Task sent. Assigned task_id: {}".format(task.id)) + +# Use the task id to get the result. Increase timeout if needed. +wait_for_result(celery_app, task.id, timeout=300) diff --git a/scripts/predict_retrosynthesis_tree.py b/scripts/predict_retrosynthesis_tree.py new file mode 100644 index 0000000000000000000000000000000000000000..002cc07a1546c0758214eec2867682ef2b03cea1 --- /dev/null +++ b/scripts/predict_retrosynthesis_tree.py @@ -0,0 +1,48 @@ +""" +Update this script and run it inside the worker container to run your retrosynthesis tree predictions via script +""" +print("Setting up retrosynthesis tree prediction...") + +# All necessary imports +from celery import Celery +from utils import wait_for_result + +# Initialize Celery +celery_app = Celery() +print("Broker:", celery_app.conf.broker_url) +print("Backend:", celery_app.conf.result_backend) + +# Choose product for retrosynthesis tree prediction +product = "C1C(C[Si](C)(C)C)=CCC2C(=O)OC(=O)C12" +#product = "Cc1cc2scnc2cc1N" +#product = "CC(C)(C(=O)O)C1C=CC=C(C2CC2)C1=O" +#product = "Nc1ccc2scnc2c1Br" + +# Setup task kwargs +kwargs = { + "topn": 15, # Number of results per reactant + "num_beams": 15, # Number of beams used for prediction. Must be >= topn + "fap": 0.6, # Forward likelihood acceptance probability (not length averaged) + "fld": 0.2, # Forward likelihood delta required between the top2 forward prediction results + "max_depth": 4, # Max depth of the retrosynthesis tree + "beam_width": 6, # Max amount of nodes being expanded in each step + "device": None, # Device used for predicting, either "cuda" or "cpu", None defaults to cuda if available + "ckpt_forward": "Pistachio2025Q2-Forward", # Default forward model + "ckpt_retro": "Pistachio2025Q2-Retro", # Default retrosynthesis model + "vocab": "Pistachio2025Q2", # Vocab for default forward and retrosynthesis models + # "ckpt_forward_path": "models/forward/Pistachio2025Q2-Forward.ckpt", # Can be used instead of ckpt_forward + # "ckpt_retro_path": "models/retrosynthesis/Pistachio2025Q2-Retro.ckpt", # Can be used instead of ckpt_retro + # "vocab_path": "vocab/Pistachio2025Q2.txt", # Can be used instead of vocab +} + +# Send the retro_prediction_tree task with the product and kwargs +task = celery_app.send_task( + "tasks.retro_prediction_tree", + [product], + kwargs=kwargs, + queue="retro_prediction", +) +print("Task sent. Assigned task_id: {}".format(task.id)) + +# Use the task id to get the result. Increase timeout if needed. +wait_for_result(celery_app, task.id, timeout=3000) diff --git a/scripts/run_notebook_examples.py b/scripts/run_notebook_examples.py new file mode 100644 index 0000000000000000000000000000000000000000..8c07fba0f8eb617f9780579377f69f21979c53ed --- /dev/null +++ b/scripts/run_notebook_examples.py @@ -0,0 +1,144 @@ +print("Setup underway...") + +# All necessary imports +from celery import Celery +from utils import wait_for_result + +# Initialize Celery +celery_app = Celery() +print("Broker:", celery_app.conf.broker_url) +print("Backend:", celery_app.conf.result_backend) + + +print("") +print("============================") +print("1a. Product prediction - batch") + +# Set up a list of reactants to make predictions +reactants_list = ["CCI.O=Cc1ccc([N+](=O)[O-])c(O)c1", "CCOc1cc(O)c(C=O)cc1OCC.OCCCBr", "C=CCc1cc(OCc2ccccc2)ccc1O.CCBr"] + +# Setup task kwargs +kwargs = { + "topn": 1, # Number of results per reactant + "num_beams": 3, # Number of beams used for prediction. Must be >= topn + "device": None, # Device used for predicting, either "cuda" or "cpu", None defaults to cuda if available + "ckpt_forward": "Pistachio2025Q2-Forward", # Default forward model + "vocab": "Pistachio2025Q2", # Vocab for default forward model + # "ckpt_forward_path": "models/forward/Pistachio2025Q2-Forward.ckpt", # Can be used instead of ckpt_forward + # "vocab_path": "vocab/Pistachio2025Q2.txt", # Can be used instead of vocab +} + +# Send the product_prediction task with the reaction list and kwargs +task = celery_app.send_task( + "tasks.product_prediction", + [reactants_list], + kwargs=kwargs, + queue="product_prediction", +) +print("Task sent. Assigned task_id: {}".format(task.id)) + +# Use the task id to get the result. Increase timeout if needed. +wait_for_result(celery_app, task.id, timeout=180) + +print("") +print("============================") +print("1b. Product prediction - top 3") + +# Set up a list of reactants to make predictions +reactants_list = ["CCI.O=Cc1ccc([N+](=O)[O-])c(O)c1"] + +# Setup task kwargs +kwargs = { + "topn": 3, # Number of results per reactant + "num_beams": 5, # Number of beams used for prediction. Must be >= topn + "device": None, # Device used for predicting, either "cuda" or "cpu", None defaults to cuda if available + "ckpt_forward": "Pistachio2025Q2-Forward", # Default forward model + "vocab": "Pistachio2025Q2", # Vocab for default forward model + # "ckpt_forward_path": "models/forward/Pistachio2025Q2-Forward.ckpt", # Can be used instead of ckpt_forward + # "vocab_path": "vocab/Pistachio2025Q2.txt", # Can be used instead of vocab +} + +# Send the product_prediction task with the reaction list and kwargs +task = celery_app.send_task( + "tasks.product_prediction", + [reactants_list], + kwargs=kwargs, + queue="product_prediction", +) +print("Task sent. Assigned task_id: {}".format(task.id)) + +# Use the task id to get the result. Increase timeout if needed. +wait_for_result(celery_app, task.id, timeout=180) + +print("") +print("============================") +print("2. Retrosynthesis prediction") + +# Choose product for retrosynthesis prediction +product = "C=CC(=C)C[Si](C)(C)C" + +# Setup task kwargs +kwargs = { + "topn": 15, # Number of results per reactant + "num_beams": 15, # Number of beams used for prediction. Must be >= topn + "fap": 0.6, # Forward likelihood acceptance probability (not length averaged) + "fld": 0.2, # Forward likelihood delta required between the top2 forward prediction results + "device": None, # Device used for predicting, either "cuda" or "cpu", None defaults to cuda if available + "ckpt_forward": "Pistachio2025Q2-Forward", # Default forward model + "ckpt_retro": "Pistachio2025Q2-Retro", # Default retrosynthesis model + "vocab": "Pistachio2025Q2", # Vocab for default forward and retrosynthesis models + # "ckpt_forward_path": "models/forward/Pistachio2025Q2-Forward.ckpt", # Can be used instead of ckpt_forward + # "ckpt_retro_path": "models/retrosynthesis/Pistachio2025Q2-Retro.ckpt", # Can be used instead of ckpt_retro + # "vocab_path": "vocab/Pistachio2025Q2.txt", # Can be used instead of vocab +} + +# Send the retro_prediction task with the product and kwargs +task = celery_app.send_task( + "tasks.retro_prediction", + [product], + kwargs=kwargs, + queue="retro_prediction", +) +print("Task sent. Assigned task_id: {}".format(task.id)) + +# Use the task id to get the result. Increase timeout if needed. +wait_for_result(celery_app, task.id, timeout=300) + +print("") +print("============================") +print("3. Retro tree prediction") + +# Choose product for retrosynthesis tree prediction +product = "C1C(C[Si](C)(C)C)=CCC2C(=O)OC(=O)C12" +#product = "Cc1cc2scnc2cc1N" +#product = "CC(C)(C(=O)O)C1C=CC=C(C2CC2)C1=O" +#product = "Nc1ccc2scnc2c1Br" + +# Setup task kwargs +kwargs = { + "topn": 15, # Number of results per reactant + "num_beams": 15, # Number of beams used for prediction. Must be >= topn + "fap": 0.6, # Forward likelihood acceptance probability (not length averaged) + "fld": 0.2, # Forward likelihood delta required between the top2 forward prediction results + "max_depth": 4, # Max depth of the retrosynthesis tree + "beam_width": 6, # Max amount of nodes being expanded in each step + "device": None, # Device used for predicting, either "cuda" or "cpu", None defaults to cuda if available + "ckpt_forward": "Pistachio2025Q2-Forward", # Default forward model + "ckpt_retro": "Pistachio2025Q2-Retro", # Default retrosynthesis model + "vocab": "Pistachio2025Q2", # Vocab for default forward and retrosynthesis models + # "ckpt_forward_path": "models/forward/Pistachio2025Q2-Forward.ckpt", # Can be used instead of ckpt_forward + # "ckpt_retro_path": "models/retrosynthesis/Pistachio2025Q2-Retro.ckpt", # Can be used instead of ckpt_retro + # "vocab_path": "vocab/Pistachio2025Q2.txt", # Can be used instead of vocab +} + +# Send the retro_prediction_tree task with the product and kwargs +task = celery_app.send_task( + "tasks.retro_prediction_tree", + [product], + kwargs=kwargs, + queue="retro_prediction", +) +print("Task sent. Assigned task_id: {}".format(task.id)) + +# Use the task id to get the result. Increase timeout if needed. +wait_for_result(celery_app, task.id, timeout=3000) diff --git a/scripts/utils.py b/scripts/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..82b0e15beb950968951f162fcbbb74c49d9eaf47 --- /dev/null +++ b/scripts/utils.py @@ -0,0 +1,29 @@ +import time +from celery import Celery +from celery.result import AsyncResult +from typing import Dict, Any + + +# Wait for the celery result +def wait_for_result(app: Celery, task_id: str, timeout: float = 120.0, poll: float = 0.5) -> Dict[str, Any]: + """ + Poll for a result with a timeout. If task updates state with meta (e.g., PROGRESS), + we surface that along the way. + """ + res = AsyncResult(task_id, app=app) + t0 = time.time() + last_state = None + + while True: + state = res.state + if state != last_state: + print(f"State: {state} | Info: {res.info}") + last_state = state + + if res.ready(): + # could be SUCCESS or FAILURE; .get() will raise on FAILURE + return res.get(propagate=False) # returns exception object if failed + + if time.time() - t0 > timeout: + raise TimeoutError(f"Task {task_id} did not finish in {timeout} seconds.") + time.sleep(poll) \ No newline at end of file diff --git a/vocab/Pistachio2025Q2.txt b/vocab/Pistachio2025Q2.txt new file mode 100644 index 0000000000000000000000000000000000000000..def7c5845e502e7a541475efb4eade8cc16d822b --- /dev/null +++ b/vocab/Pistachio2025Q2.txt @@ -0,0 +1,624 @@ +[PAD] +[unused1] +[unused2] +[unused3] +[unused4] +[unused5] +[unused6] +[unused7] +[unused8] +[unused9] +[unused10] +[UNK] +[CLS] +[SEP] +[MASK] +c +C +( +) +O +1 +2 += +N +. +n +3 +F +Cl +4 +- +[C@H] +[C@@H] +S +# +Br +[O] +5 +[nH] +/ +[N+] +s +[K] +[O-] +[Na] +P +o +[Si] +B +[Cl] +[Pd] +6 +I +[C@] +[P] +[C@@] +[OH] +[Cs] +[H] +[Li] +\ +[CH2] +[C] +[N-] +[Na+] +[P-] +7 +[NaH] +[NH4+] +[Cl-] +[I] +[Fe] +[Cu] +[Mg] +[BH4-] +[CH3] +[Br] +[N] +[n+] +8 +[BH-] +[Zn] +[Sn] +[Li+] +[B-] +[Al] +[c] +[c-] +[BH3-] +[CH] +[Fe+2] +[F-] +[P+] +[Br-] +[OH-] +[O+] +[AlH4-] +[I-] +[S@] +[Pd+2] +[F] +[C-] +[Ni] +[SiH] +[S@@] +9 +[S] +[Ti] +[Ag] +[K+] +[AlH] +[Mn] +[NH3+] +[Pt] +[nH+] +[Cr] +[I+3] +[PH] +[Ru] +[S+] +[Ca] +[NH2+] +[NH+] +[Se] +[Os] +%10 +[NH] +[H+] +[se] +[Cu+2] +[SH] +[SiH2] +[Rh] +[Cl+3] +[Cl+] +[P@] +[SiH3] +[H-] +[P@@] +[PH+] +[S-] +[Co] +[B] +[Hg] +[SnH] +[NH2] +%11 +[Mg+2] +[I+] +[Ce] +[Pb] +[NH-] +[Mo] +[Ge] +[N@@] +[cH-] +[PH2] +[C+] +[In] +[W] +[n] +%12 +[s+] +[n-] +[N@] +[Zr] +[OH2+] +[Au] +[Ar] +[As] +[Ba] +[CH-] +[Ir] +[o+] +[Bi] +[CH2-] +[AlH2-] +%13 +[Ag+] +[Ba+2] +[KH] +[S@@+] +[Sb] +[OH3+] +[Sc+3] +p +[LiH] +[Os-2] +[Hg+2] +[S@+] +[Rh-3] +[Pd+] +[Cu+] +[te] +[Fe-3] +[V] +[OH+] +[Ca+2] +[Te] +[Cs+] +[Re] +[La] +[Ce+3] +[SiH4] +[Yb+3] +[Si@H] +[PH4+] +[Zr+2] +[Sm] +%14 +[Rh+2] +[I+2] +[Si-] +[Pd-] +[Sb-] +[O-2] +[BH] +[CH3-] +[Zn+2] +[SeH] +[Gd] +[Rh+3] +[Yb] +[N@+] +[SH+] +[Br+2] +[S-2] +[Pt-2] +[Ir+3] +[Th] +[CH+] +[Cu-2] +[Cr+3] +[Al+] +[N@@+] +%15 +[PH3+] +[Hf] +[CH2+] +[Al-] +[Pd-2] +[Ga] +%16 +[Si@] +[Cl+2] +[NH2-] +[Co+2] +[Gd+3] +[Bi+3] +[AlH-] +[Y+3] +[cH+] +[AlH3] +[Rb] +[Sn+2] +[Cd] +[Hg+] +[Zr+4] +[Ru-2] +[SH-] +[Ni+2] +[Si@@] +[SnH4] +[SH2] +[c+] +[Ti+3] +[S+2] +[Pb+2] +[Cr+2] +[PH3] +[SnH3] +[Cu-] +%17 +[Ir+] +[Sr] +[PdH2] +[Rh+] +[Sr+2] +[P+3] +[Y] +[Ru+2] +[Tl+] +[Zr+3] +[Nb] +[B+] +[Al+3] +[Au+] +[V+2] +[IH2+] +[Eu] +[B+3] +[IH] +%18 +[Br+] +[Pt+2] +[Sc] +[Cd+2] +%19 +[SnH2] +[Ce+4] +[SH3+] +[Fe+3] +[CaH2] +%20 +[Ta] +[Mg+] +[AlH2] +[Tl] +[Tl+3] +[Hf+2] +[IH+] +[Sb+3] +[Si+] +[Xe] +[H+2] +[PH5] +[Zn+] +[TeH] +[ClH+] +[Er+3] +[Sm+3] +b +[SH2+] +[sH+] +%21 +[Eu+3] +[Ac] +[Lu] +[YH] +%22 +[Ti+2] +[Tb] +[Si-2] +[Be] +[se+] +[N+3] +[Ti+4] +[Ga+3] +[pH] +[N@@H+] +[Pr] +[Au-] +[Ir-] +[U+2] +[N@H+] +[Mn+2] +[InH2] +[U] +%23 +[Ho] +[RuH] +[si] +[PH2+] +[GeH] +[F+] +%24 +[Si@@H] +[SeH2] +[At] +[Nd] +[BH2-] +[Be+2] +[Fm] +[Nd+3] +[RuH2] +[Dy+3] +[SeH-] +[Ru+] +[siH] +[B+2] +[Tc] +[LaH] +[Zr+] +[Si+4] +[P@@H] +[p] +[BH2] +[AlH2+] +[IH-] +[NiH] +[Ir+2] +[As-] +[As+] +[Dy] +%25 +[BrH+] +[Na-] +[B@-] +[Ni-2] +[IH3] +[Fe-4] +[PbH2] +[GeH2] +[SiH-] +[SH4] +[FH+] +[Hf+3] +[Y-] +[UH] +[Er] +[Co+3] +[Hf+4] +[Tb+3] +[AsH2] +[VH] +[PH4] +[Sn-] +[BH+] +[IrH] +[B@@-] +[SbH2] +[Pr+3] +[Se-] +[Se+] +[Ni+] +[Se@+] +[Ge-] +[Bi+] +[PH-] +[C+4] +[Po] +[Nb+5] +[Cm] +[GeH3] +[La+3] +[ClH2+] +[NaH3] +[V+5] +[Mg-] +[Sn+4] +[In+3] +[K-] +[Rh-2] +[Ac-] +[Ce+2] +[SeH+] +[W+] +[p+] +[Au-2] +[Ag+3] +[AsH] +[Fe-] +[p-] +[Tl+2] +[Sn+] +[CuH] +[Mn+4] +[SH5] +[Se@@+] +[H-2] +[Li-2] +[Ga+2] +[Yb+2] +[Ru-] +[CH3+] +[WH] +[Sn-2] +[Pt-] +[Mn+3] +[Bi+2] +[Te+] +[Rh-] +[Sg] +[Y+] +[NaH+] +[SnH2-] +[Au+3] +[Cr+] +[Co+] +[Fr] +[He] +[Mo+2] +[PH4-] +[TeH3] +[RgH] +[P@H] +[P-3] +[W+4] +[Cf] +[Mn+] +[Ga+] +[MgH] +[InH] +[SH3] +[Ta+2] +[IrH2] +[Ni-] +[OH+2] +[Mo+4] +[Re+] +[te+] +[N-2] +[N+2] +[IH2] +[Db] +[Pt+4] +[P@H+] +[ClH+2] +[Al+2] +[Re+4] +[TiH] +[BiH3] +[AsH4+] +[As+3] +[C-4] +[Zn-] +[AuH] +[Xe+] +[Sc+2] +[Fe+] +[CaH] +[PtH2] +[Cl-2] +[PtH] +[oH+] +[Lu+3] +[I-2] +[Sb+2] +[NaH4] +[RuH3] +[F-2] +[VH2] +[Pa] +[BiH2] +[AsH3] +[Br-2] +[Ca+] +[C-2] +[YH4] +[Si+2] +[Se+2] +[TeH2] +[CoH] +[Sb+] +[YH3] +[Cn] +[PH2-] +[Os+2] +[Cr-] +[Ru+3] +[YH2] +[Na-2] +[GeH4] +[V+4] +[Na+2] +[InH3] +[SbH] +[BrH2+] +[Fe-2] +[P+5] +[Ra] +[No] +[Te-] +[Ba+] +[Tm] +[GaH] +[CuH2] +[NH3+2] +[CaH3] +[Ne] +[Se-2] +[AlH3-] +[CeH] +[Mg-2] +[Li-] +[Hs] +[N-3] +[Ta+3] +[NiH2] +[Ru+6] +[YH5] +[LiH+] +[Rf] +[CuH2-] +[Pd-3] +[Cd+] +[OH2+2] +[W-2] +[FH2+] +[BrH+2] +[V+3] +[Ta-] +[C+3] +[MnH] +[BaH] +[Tm+3] +[C+2] +[Ti+] +[P+2] +[Ge+4] +[RaH] +[Sc+] +[PtH+] +[Gd+2] +[CmH5] +[Ce+] +[Pu] +[Am] +[NaH+2] +[AsH+] +[Ac+3] +[Ag-2] +[Ho+3] +[Co-2] +[Bh-] +[TiH2] +[Ta+5] +[Tc+4] +[TaH3] +[Ag+2] +[FH+2] +[BaH2] +[PtH+2] +[Hf+] +[Br+3] +[SbH3] +[AcH] +[NH+3] +[YH7] +[Pt+] +[Eu+2] +[Fe+4] +[Sm+2] +[Es] diff --git a/worker/Dockerfile b/worker/Dockerfile new file mode 100644 index 0000000000000000000000000000000000000000..3d219aa2f5f4327ca1115cf1ffa5d5bb8ce73b71 --- /dev/null +++ b/worker/Dockerfile @@ -0,0 +1,33 @@ +FROM public.ecr.aws/docker/library/python:3.13-slim + +# Set working directory +WORKDIR /app + +# Install system dependencies (only what you actually need) +RUN apt-get update \ + && apt-get install -y \ + libexpat1 \ + libxrender1 \ + libxext6 \ + && rm -rf /var/lib/apt/lists/* + +# Copy dependency definitions first (for cache friendliness) +COPY requirements.txt . + +# Install Python dependencies +RUN pip install --no-cache-dir --upgrade pip && \ + pip install --no-cache-dir -r requirements.txt + +# Copy the application code +COPY . . + +# Ensure entrypoint is executable +RUN chmod +x start.sh + +# Change to non-root user +RUN useradd -m -u 1000 rxn && \ + chown -R rxn:rxn /app + +USER rxn + +ENTRYPOINT ["./start.sh"] diff --git a/worker/celeryconfig.py b/worker/celeryconfig.py new file mode 100644 index 0000000000000000000000000000000000000000..8dd0adf90746da2570e0f861b53fedd78950e6fa --- /dev/null +++ b/worker/celeryconfig.py @@ -0,0 +1,20 @@ +import os + +broker_url = os.getenv("CELERY_BROKER_URL", "amqp://ubuntu:ubuntu@broker:5672//") +result_backend = os.getenv("CELERY_RESULT_BACKEND", "redis://:ubuntu@redis:6379/0") + +accept_content = ["json"] +task_serializer = "json" +result_serializer = "json" + +task_ignore_result = False + +task_track_started = True + +worker_hijack_root_logger = False + +task_routes = { + "tasks.product_prediction": {"queues": "product_prediction"}, + "tasks.retro_prediction": {"queues": "retro_prediction"}, + "tasks.retro_prediction_tree": {"queues": "retro_prediction"}, +} diff --git a/worker/model_loading.py b/worker/model_loading.py new file mode 100644 index 0000000000000000000000000000000000000000..8c74419ad8bd36891a6715f009d88c7cbcad3f70 --- /dev/null +++ b/worker/model_loading.py @@ -0,0 +1,242 @@ +from typing import List, Dict, Optional +import torch +from lightning import seed_everything +from transformers_model.model import LitVanillaTransformer + + +# Loads a LitVanillaTransformer set for forward (product) prediction +def load_forward_model( + ckpt_path: str, + vocab_path: str, + num_beams: int, + topn: int, + device: Optional[str] = None, +) -> LitVanillaTransformer: + return load_lit_transformer_for_inference( + ckpt_path=ckpt_path, + vocab_path=vocab_path, + task="forward", + num_beams=num_beams, + topn=topn, + gen_max_length=256, + device=device, + ) + + +# Loads a LitVanillaTransformer set for retrosynthesis prediction +def load_retro_model( + ckpt_path: str, + vocab_path: str, + num_beams: int, + topn: int, + device: Optional[str] = None, +) -> LitVanillaTransformer: + return load_lit_transformer_for_inference( + ckpt_path=ckpt_path, + vocab_path=vocab_path, + task="forward", + num_beams=num_beams, + topn=topn, + gen_max_length=256, + device=device, + ) + + +# Loads the LitVanillaTransformer for prediction +def load_lit_transformer_for_inference( + ckpt_path: str, + vocab_path: str, + task: str = "forward", + num_beams: int = 3, + topn: int = 3, + gen_max_length: int = 256, + device: Optional[str] = None, +) -> LitVanillaTransformer: + """ + Load the trained LitVanillaTransformer exactly like LightningCLI would, + and prepare it for inference/generation. + """ + torch.set_float32_matmul_precision("high") + + # Device handling when not set + if device is None: + device = "cuda" if torch.cuda.is_available() else "cpu" + + # Load from checkpoint while ensuring the same hyperparameters used at predict time. + lit = LitVanillaTransformer.load_from_checkpoint( + ckpt_path, + vocab_path=vocab_path, + task=task, + num_beams=num_beams, + topn=topn, + max_length=gen_max_length, + device=device, + ) + lit.eval() + + # Align internal config device with the actual target device to avoid mask/device mismatches + target_device = torch.device(device) + lit.to(target_device) + if hasattr(lit, "model") and hasattr(lit.model, "config"): + lit.model.config.device = target_device # your encoder/decoder build masks on config.device + + # Return the LitVanillaTransformer + return lit + + +# Uses the transformers model to run predictions with the SMILES inputs +@torch.inference_mode() +def predict_smiles( + lit: LitVanillaTransformer, # Model used for inference + inputs: List[str], # List of SMILES for prediction + *, + batch_size: int = 512, + truncation: bool = True, + padding: str = "max_length", + data_max_length: int = 278, + seed: Optional[int] = 42, +) -> List[Dict]: + + # Sets the seed for reproducibility + if seed is not None: + seed_everything(seed, workers=True) + + # Gets the device in use + device = next(lit.parameters()).device + + # Gets the tokenizer from the LitVanillaTransformer + tokenizer = lit.tokenizer + + # List to hold prediction results + all_results: List[Dict] = [] + + # Prediction loop for batch prediction + for start in range(0, len(inputs), batch_size): + + # Prepare the batch inputs according to batch size + batch_src = inputs[start : start + batch_size] + print(f"DEBUG - batch_src: {batch_src}") + + # Prepare inputs by tokenizing exactly like in the DataModule + encoder_input_ids = tokenizer( + batch_src, + truncation=truncation, + padding=padding, + max_length=data_max_length, + return_token_type_ids=False, + return_tensors="pt", + )["input_ids"] + + # Move inputs to device if one is set + if device is not None: + encoder_input_ids = encoder_input_ids.to(device) + + # Generate the model predictions + outputs = lit.model.generate( + encoder_input_ids, + do_sample=False, + max_length=lit.hparams.max_length, + num_beams=lit.hparams.num_beams, + num_return_sequences=lit.hparams.topn, + return_dict_in_generate=True, + output_scores=True, + ) + + # Get predicted sequences and scores + sequences = outputs.sequences # shape [B*topn, T] + scores = outputs.sequences_scores # 1 score per sequence + + # Decode and detokenize the sequences to obtain the predicted SMILES string + pred_texts = tokenizer.batch_decode(sequences, skip_special_tokens=True) + pred_texts = tokenizer.batch_detokenize_smiles_string(pred_texts) + + # Also decode and detokenize the encoder inputs to a SMILES string + src_texts_dec = tokenizer.batch_decode(encoder_input_ids, skip_special_tokens=True) + src_texts_dec = tokenizer.batch_detokenize_smiles_string(src_texts_dec) + print(f"DEBUG - src_texts_dec: {src_texts_dec}") + + # Get topn from the hyperparameters to know exactly how many results are grouped for the same source + topn = lit.hparams.topn + + # List to hold each batch prediction results + batch_results: List[Dict] = [] + + # Loop through the inputs + for i, src in enumerate(src_texts_dec): + + # Get starting index for each unique input + start_j = i * topn + + # Loop through all predictions for a given input + for j in range(start_j, start_j + topn): + + # Build each prediction result object + item = { + "source": src, + "predicted_target": pred_texts[j], + "confidence": scores[j].item(), # you keep raw score in your code + } + + # Add the prediction result to the batch results + batch_results.append(item) + + # Add this batch results to the total results list + all_results.extend(batch_results) + + # Return the prediction results + return all_results + + +# Uses the transformers model to check the likelihood of a reaction using the input and output SMILES +@torch.inference_mode() +def reaction_likelihood( + lit: LitVanillaTransformer, + src_text: str, + tgt_text: str, + truncation: bool = True, + padding: str = "max_length", + data_max_length: int = 278, +) -> float: + # Get the model, tokenizer and device from LitVanillaTransformer + model = lit.model + tokenizer = lit.tokenizer + device = next(model.parameters()).device + + # Tokenizes the SMILES sources + src_ids = tokenizer( + src_text, + truncation=truncation, + padding=padding, + max_length=data_max_length, + return_tensors="pt", + ).input_ids.to(device) + + # Tokenizes the SMILES targets + tgt_ids = tokenizer( + tgt_text, + truncation=truncation, + padding=padding, + max_length=data_max_length, + return_tensors="pt", + ).input_ids.to(device) + + # Gets the outputs from the model + outputs = model( + encoder_input_ids=src_ids, + decoder_input_ids=tgt_ids, + labels=tgt_ids, + return_dict=True, + ) + + # loss = mean NEGATIVE log-prob per token + loss = outputs.loss + + # Number of actual tokens (excluding padding) + pad_id = tokenizer.pad_token_id + num_tokens = (tgt_ids != pad_id).sum() + + # Calculates the log likelihood + log_likelihood = -loss * num_tokens + + # Returns likelihood probability + return float(torch.exp(log_likelihood)) diff --git a/worker/prediction.py b/worker/prediction.py new file mode 100644 index 0000000000000000000000000000000000000000..aa38f143ec6530651dab6e27bea3e65ebdd513f1 --- /dev/null +++ b/worker/prediction.py @@ -0,0 +1,166 @@ +import logging +from collections import Counter +from typing import Dict, List +import numpy as np +from rdkit import Chem +from utils import standardize_molecules +from model_loading import predict_smiles + +logger = logging.getLogger() + +RXN_SMILES_SEPARATOR = ">>" + + +def matching_atomic_numbers( + product_atomic_numbers: List[int], + reactant_atomic_numbers: List[int], + ignore_stoichiometry: bool = True, +) -> bool: + """ + Check for matching atomic numbers in products and reactants. + + Args: + product_atomic_numbers (List[int]): list of atomic numbers for the atoms in the product(s). + reactant_atomic_numbers (List[int]): list of atomic numbers for the atoms in the reactant(s). + ignore_stoichiometry (bool): whether stoichiometry is ignored. Default to True. + + Returns: + bool: whether the atomic numbers are matching. + """ + container_type = set if ignore_stoichiometry else list + atomic_numbers_difference = Counter( + container_type(product_atomic_numbers) + ) - Counter(container_type(reactant_atomic_numbers)) + return sum(atomic_numbers_difference.values()) == 0 + + +def forward_reaction( + smiles_list: List[str], + lit, +) -> List[Dict]: + """ + Get structured results and pictures from a batch reaction prediction. + + Args: + smiles_list (List[str]): List of reactants SMILES. + lit (Translator): predictive model. + + Returns: + Dict[str, List[Dict[str, Any]]]: Dict of predictions lists from the model. + """ + + # Make product predictions + results_list = predict_smiles(lit, smiles_list, batch_size=512) + + # Get the products, reactants and confidences + reactants_list = [result["source"] for result in results_list] + predicted_list = [result["predicted_target"] for result in results_list] + confidence_list = [result["confidence"] for result in results_list] + + # Start list to save the results + final_results_list = [] + + # Loop the prediction results data + for ( + predicted_smiles, + confidence, + reactants, + ) in zip( + predicted_list, + confidence_list, + reactants_list, + ): + try: + # process predicted SMILES and get molecules + predicted_mol = Chem.MolFromSmiles(predicted_smiles) + predicted_can_smiles = standardize_molecules( + predicted_smiles, + canonicalize=True, + sanitize=True, + fragment_bond="~", + ordered_precursors=False, + ) + reactants_mol = Chem.MolFromSmiles(reactants) + # sanity check on atomic numbers + products_atoms = [atom.GetAtomicNum() for atom in predicted_mol.GetAtoms()] + reactants_atoms = [atom.GetAtomicNum() for atom in reactants_mol.GetAtoms()] + atoms_ok = matching_atomic_numbers(products_atoms, reactants_atoms) + if not atoms_ok: + predicted_can_smiles = "" + logger.info("Atoms not ok!") + # represent reaction SMARTS + reaction_smiles = reactants + RXN_SMILES_SEPARATOR + predicted_can_smiles + + except Exception as errrr: + logger.error(errrr) + reaction_smiles = ( + reactants + RXN_SMILES_SEPARATOR + predicted_smiles + ).replace(" ", "") + message = "Invalid SMILES predicted" + logger.info("Predicted invalid reaction") + + # Create the product prediction object and append it to results list + final_results_list.append( + { + "reactants": reactants, + "product": predicted_smiles, + "confidence": float(np.exp(confidence)), # Convert confidence to probability + "smiles": reaction_smiles, + } + ) + + # Return product prediction results + return final_results_list + + +def retro_reaction( + smiles: str, + lit, +) -> List[Dict]: + """ + Get structured results and pictures from a batch reaction prediction. + + Args: + smiles (str): Reactants SMILES string. + lit (Translator): predictive model. + + Returns: + Dict[str, List[Dict[str, Any]]]: Dict of predictions lists from the model. + """ + + # Make retro predictions + results_list = predict_smiles(lit, [smiles], batch_size=512) + + # Get the products, reactants and confidences + product_list = [result["source"] for result in results_list] + predicted_list = [result["predicted_target"] for result in results_list] + confidence_list = [result["confidence"] for result in results_list] + + # Start list to save the results + final_results_list = [] + + # Loop the prediction results data + for ( + predicted_smiles, + confidence, + product, + ) in zip( + predicted_list, + confidence_list, + product_list, + ): + # Create the full reaction SMILES + reaction_smiles = predicted_smiles + RXN_SMILES_SEPARATOR + product + + # Create the retro prediction object and append it to results list + final_results_list.append( + { + "reactants": predicted_smiles, + "product": product, + "confidence": float(np.exp(confidence)), # Convert confidence to probability + "smiles": reaction_smiles, + } + ) + + # Return retro prediction results + return final_results_list diff --git a/worker/requirements.txt b/worker/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..a19a3abb898b83ee8703aca5497c728556928b5b --- /dev/null +++ b/worker/requirements.txt @@ -0,0 +1,11 @@ +celery[redis]<5.4.0 +lightning +Levenshtein +rdkit +redis +torch +transformers<5 +rxn-chem-utils +rxn-reaction-preprocessing +numpy +jsonargparse[signatures] \ No newline at end of file diff --git a/worker/scscore/__init__.py b/worker/scscore/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/worker/scscore/models/full_reaxys_model_1024bool/model.ckpt-10654.as_numpy.pickle b/worker/scscore/models/full_reaxys_model_1024bool/model.ckpt-10654.as_numpy.pickle new file mode 100644 index 0000000000000000000000000000000000000000..3e8a25ac087509591dc5f54e4339680522140982 --- /dev/null +++ b/worker/scscore/models/full_reaxys_model_1024bool/model.ckpt-10654.as_numpy.pickle @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:d3369182fc010850549712c0b67be4ba6e5947bf0118f854af78e98d3c89789a +size 2676496 diff --git a/worker/scscore/standalone_model_numpy.py b/worker/scscore/standalone_model_numpy.py new file mode 100644 index 0000000000000000000000000000000000000000..1993ec237c8eda449c81c9f309ebc0331f212100 --- /dev/null +++ b/worker/scscore/standalone_model_numpy.py @@ -0,0 +1,128 @@ +''' +This is a standalone, importable SCScorer model. It does not have tensorflow as a +dependency and is a more attractive option for deployment. The calculations are +fast enough that there is no real reason to use GPUs (via tf) instead of CPUs (via np) +''' + +import numpy as np +import rdkit.Chem as Chem +import rdkit.Chem.AllChem as AllChem +import json +import gzip +import six + +import os +project_root = os.path.dirname(os.path.dirname(__file__)) + +score_scale = 5.0 +min_separation = 0.25 + +FP_len = 1024 +FP_rad = 2 + +def sigmoid(x): + return 1 / (1 + np.exp(-x)) + +class SCScorer(): + def __init__(self, score_scale=score_scale): + self.vars = [] + self.score_scale = score_scale + self._restored = False + + def restore(self, weight_path=os.path.join(project_root,'scscore' , 'models', 'full_reaxys_model_1024bool', 'model.ckpt-10654.as_numpy.pickle'), FP_rad=FP_rad, FP_len=FP_len): + self.FP_len = FP_len; self.FP_rad = FP_rad + self._load_vars(weight_path) + print('Restored variables from {}'.format(weight_path)) + + if 'uint8' in weight_path or 'counts' in weight_path: + def mol_to_fp(self, mol): + if mol is None: + return np.array((self.FP_len,), dtype=np.uint8) + fp = AllChem.GetMorganFingerprint(mol, self.FP_rad, useChirality=True) # uitnsparsevect + fp_folded = np.zeros((self.FP_len,), dtype=np.uint8) + for k, v in six.iteritems(fp.GetNonzeroElements()): + fp_folded[k % self.FP_len] += v + return np.array(fp_folded) + else: + def mol_to_fp(self, mol): + if mol is None: + return np.zeros((self.FP_len,), dtype=np.float32) + return np.array(AllChem.GetMorganFingerprintAsBitVect(mol, self.FP_rad, nBits=self.FP_len, + useChirality=True), dtype=np.bool) + self.mol_to_fp = mol_to_fp + + self._restored = True + return self + + def smi_to_fp(self, smi): + if not smi: + return np.zeros((self.FP_len,), dtype=np.float32) + return self.mol_to_fp(self, Chem.MolFromSmiles(smi)) + + def apply(self, x): + if not self._restored: + raise ValueError('Must restore model weights!') + # Each pair of vars is a weight and bias term + for i in range(0, len(self.vars), 2): + last_layer = (i == len(self.vars)-2) + W = self.vars[i] + b = self.vars[i+1] + x = np.matmul(x, W) + b + if not last_layer: + x = x * (x > 0) # ReLU + x = 1 + (score_scale - 1) * sigmoid(x) + return x + + def get_score_from_smi(self, smi='', v=False): + if not smi: + return ('', 0.) + fp = np.array((self.smi_to_fp(smi)), dtype=np.float32) + if sum(fp) == 0: + if v: print('Could not get fingerprint?') + cur_score = 0. + else: + # Run + cur_score = self.apply(fp) + if v: print('Score: {}'.format(cur_score)) + mol = Chem.MolFromSmiles(smi) + if mol: + smi = Chem.MolToSmiles(mol, isomericSmiles=True, kekuleSmiles=True) + else: + smi = '' + return (smi, cur_score) + + def _load_vars(self, weight_path): + if weight_path.endswith('pickle'): + import pickle + with open(weight_path, 'rb') as fid: + self.vars = pickle.load(fid, encoding="latin1") + self.vars = [x.tolist() for x in self.vars] + elif weight_path.endswith('json.gz'): + with gzip.GzipFile(weight_path, 'r') as fin: # 4. gzip + json_bytes = fin.read() # 3. bytes (i.e. UTF-8) + json_str = json_bytes.decode('utf-8') # 2. string (i.e. JSON) + self.vars = json.loads(json_str) + self.vars = [np.array(x) for x in self.vars] + + +if __name__ == '__main__': + model = SCScorer() + model.restore(os.path.join(project_root, 'models', 'full_reaxys_model_1024bool', 'model.ckpt-10654.as_numpy.json.gz')) + smis = ['CCCOCCC', 'CCCNc1ccccc1'] + for smi in smis: + (smi, sco) = model.get_score_from_smi(smi) + print('%.4f <--- %s' % (sco, smi)) + + model = SCScorer() + model.restore(os.path.join(project_root, 'models', 'full_reaxys_model_2048bool', 'model.ckpt-10654.as_numpy.json.gz'), FP_len=2048) + smis = ['CCCOCCC', 'CCCNc1ccccc1'] + for smi in smis: + (smi, sco) = model.get_score_from_smi(smi) + print('%.4f <--- %s' % (sco, smi)) + + model = SCScorer() + model.restore(os.path.join(project_root, 'models', 'full_reaxys_model_1024uint8', 'model.ckpt-10654.as_numpy.json.gz')) + smis = ['CCCOCCC', 'CCCNc1ccccc1'] + for smi in smis: + (smi, sco) = model.get_score_from_smi(smi) + print('%.4f <--- %s' % (sco, smi)) diff --git a/worker/simplicity.py b/worker/simplicity.py new file mode 100644 index 0000000000000000000000000000000000000000..98ff1f9cc1e2edcfb734868de5adc131635ca88a --- /dev/null +++ b/worker/simplicity.py @@ -0,0 +1,81 @@ +""" +SCScore-based synthetic complexity and simplicity scoring. + +Uses the pretrained SCScore model by Coley et al. +Source: https://github.com/connorcoley/scscore +License: MIT +""" + +from functools import lru_cache +from typing import Optional + +from scscore.standalone_model_numpy import SCScorer + + +# --- Global singleton (per process) ----------------------------------------- + +_SCORER: Optional[SCScorer] = None + + +def _get_scorer() -> SCScorer: + """ + Lazily initialize and return the SCScorer singleton. + + This ensures: + - model weights are loaded once per process + - safe behavior in Celery workers + """ + global _SCORER + if _SCORER is None: + scorer = SCScorer() + scorer.restore() # loads pretrained numpy weights + _SCORER = scorer + return _SCORER + + +# --- Public API ------------------------------------------------------------- + +@lru_cache(maxsize=100_000) +def get_scscore(smiles: str) -> float: + """ + Return the SCScore for a molecule. + + Parameters + ---------- + smiles : str + Canonical (or at least valid) SMILES. + + Returns + ------- + float + SCScore in the range [1.0, 5.0] + """ + if not smiles: + return 5.0 # treat empty as maximally complex + + scorer = _get_scorer() + _, score = scorer.get_score_from_smi(smiles) + return float(score.item()) + + +def simplicity_score_from_scscore(sc: float) -> float: + """ + Convert SCScore to simplicity score. + + s(X) = 1 - (SC(X) - 1) / 4 + + Maps: + SCScore = 1 -> simplicity = 1.0 (very simple) + SCScore = 5 -> simplicity = 0.0 (very complex) + """ + return max(0.0, min(1.0, 1.0 - (sc - 1.0) / 4.0)) + + +def simplicity(smiles: str) -> float: + """ + Convenience wrapper: SMILES -> simplicity score. + + Equivalent to: + simplicity_score_from_scscore(get_scscore(smiles)) + """ + return simplicity_score_from_scscore(get_scscore(smiles)) diff --git a/worker/spectators.py b/worker/spectators.py new file mode 100644 index 0000000000000000000000000000000000000000..a44360e33d070e1a5932d3583bf3974a49392e27 --- /dev/null +++ b/worker/spectators.py @@ -0,0 +1,117 @@ +COMMON_SOLVENT_SMILES = { + "O", # water + "CS(=O)C", # DMSO + "CC#N", # acetonitrile + "CO", # methanol + "CCO", # ethanol + "CC(C)O", # isopropanol + "C[N+](=O)[O-]", # nitromethane + "CN(C)C=O", # DMF + "CN(C)C(C)=O", # DMA + "O=C1N(CCCN1)C", # NMP + "OCCO", # ethylene glycol + "NC=O", # formamide + "O=S1(=O)CCCC1", # sulfolane + "COCCO", # 2-methoxyethanol + "CC(=O)C", # acetone + "C1CCOC1", # THF + "O1CCOCC1", # dioxane + "ClCCl", # DCM + "ClC(Cl)Cl", # chloroform + "ClCCCl", # 1,2-dichloroethane + "COC(=O)OC", # dimethyl carbonate + "COc1ccccc1", # anisole + "CCOC(=O)C", # ethyl acetate + "CCOCC", # diethyl ether + "Cc1ccccc1", # toluene + "Cc1ccccc1C", # o-xylene + "Cc1cccc(C)c1", # m-xylene + "Cc1ccc(C)cc1", # p-xylene + "Cc1cc(C)cc(C)c1", # mesitylene + "ClC(Cl)(Cl)Cl", # carbon tetrachloride + "C1CCCCC1", # cyclohexane + "CCCCCC", # n-hexane + "CCCCCCC", # heptane + "CCCCC", # pentane + "c1ccccc1", # benzene +} + +INERT_ATMOSPHERE_GASES = { + "N#N", # nitrogen + "[Ar]", # argon (rare but possible) +} + +REACTIVE_GASES = { + "O=O", # oxygen + "[H][H]", # hydrogen +} + +COMMON_ION_SMILES = { + "[Cl-]", + "[Br-]", + "[I-]", + "[F-]", + "[NH4+]", + "[Na+]", + "[K+]", + "[Li+]", +} + +COMMON_ELEMENTAL_SPECTATORS = { + "Cl", # chlorine artefact + "[Cl]", # sometimes appears bracketed + "Br", + "[Br]", + "I", + "[I]", + "[Mg]", + "Mg", + "[Zn]", + "Zn", + "[Cu]", + "Cu", + "[Fe]", + "Fe", + "[Pd]", +} + +COMMON_HALOGEN_SMILES = { + "ClCl", # Cl2 + "BrBr", # Br2 + "II", # I2 (non-standard but common in datasets) +} + + +def is_solvent(smiles: str) -> bool: + return smiles in COMMON_SOLVENT_SMILES + + +def is_gas(smiles: str) -> bool: + return smiles in INERT_ATMOSPHERE_GASES or smiles in REACTIVE_GASES + + +def is_ion(smiles: str) -> bool: + return smiles in COMMON_ION_SMILES + + +def is_halogen(smiles: str) -> bool: + return smiles in COMMON_HALOGEN_SMILES + + +# Checks if the SMILES belongs to a spectator that shouldn't be expanded via retrosynthesis +def is_spectator(smiles: str) -> bool: + return ( + is_solvent(smiles) + or is_gas(smiles) + or is_ion(smiles) + or is_halogen(smiles) + or smiles in COMMON_ELEMENTAL_SPECTATORS + ) + + +# Strips spectators that do not generate unique reactions +def strip_neutral_spectators(reactants: list[str]) -> list[str]: + return sorted( + r for r in reactants + if not is_solvent(r) and r not in INERT_ATMOSPHERE_GASES + ) diff --git a/worker/start.sh b/worker/start.sh new file mode 100644 index 0000000000000000000000000000000000000000..0dce3e14cf696b02e5c9d2c2475cf608a6811b59 --- /dev/null +++ b/worker/start.sh @@ -0,0 +1,30 @@ +#!/bin/bash +set -euo pipefail + +# Required +: "${CELERY_QUEUE:?CELERY_QUEUE is not set}" + +# Optional (with sane defaults) +CELERY_CONCURRENCY="${CELERY_CONCURRENCY:-1}" +CELERY_MAX_TASKS_PER_CHILD="${CELERY_MAX_TASKS_PER_CHILD:-1}" +RXN_LOG_LEVEL="${RXN_LOG_LEVEL:-debug}" + +echo "Starting Celery worker" +echo +echo " Queue: ${CELERY_QUEUE}" +echo " Concurrency: ${CELERY_CONCURRENCY}" +echo " Log level: ${RXN_LOG_LEVEL}" + +mkdir -p celery + +exec celery -A tasks worker -E \ + --concurrency "${CELERY_CONCURRENCY}" \ + -l "${RXN_LOG_LEVEL}" \ + --without-gossip \ + --without-mingle \ + --without-heartbeat \ + -Q "${CELERY_QUEUE}" \ + -Ofair \ + --max-tasks-per-child "${CELERY_MAX_TASKS_PER_CHILD}" \ + --pidfile=celery/%n.pid \ + --logfile=celery/%n.log diff --git a/worker/tasks.py b/worker/tasks.py new file mode 100644 index 0000000000000000000000000000000000000000..cff7983f3fb868b40f705379c16f7edde313866c --- /dev/null +++ b/worker/tasks.py @@ -0,0 +1,588 @@ +import logging +import os +import sys +import traceback +import math +import time +from heapq import nlargest +from logging.config import dictConfig +from typing import Callable, List, TypeVar, Dict + +import torch +from celery import Celery, Task +from celery.signals import setup_logging +from rxn.chemutils.conversion import canonicalize_smiles +from rxn.chemutils.exceptions import InvalidSmiles +from prediction import forward_reaction, retro_reaction +from model_loading import load_forward_model, load_retro_model, reaction_likelihood +from spectators import is_spectator, strip_neutral_spectators +from simplicity import simplicity + + +# Avoid CPU over-subscription. +torch.set_num_threads(int(os.getenv("RXN_TORCH_NUM_THREADS", "1"))) + +logger = logging.getLogger(__name__) + +app = Celery() +app.config_from_object("celeryconfig") + +T = TypeVar("T") + +def handle_exception(message: str) -> Dict: + exc_type, exc_value, exc_traceback = sys.exc_info() + return { + "title": f"{message}: {getattr(exc_type, '__name__', 'Exception')}", + "traceback": traceback.format_exception(exc_type, exc_value, exc_traceback), + "detail": str(exc_value), + } + + +@setup_logging.connect +def configure_logging(**kwargs): + dictConfig({ + "version": 1, + "disable_existing_loggers": False, + "formatters": { + "rxn": { + "format": "%(asctime)s %(levelname)-8s [%(filename)s:%(funcName)s:%(lineno)d] %(message)s" + } + }, + "handlers": { + "worker_file": { + "class": "logging.handlers.RotatingFileHandler", + "filename": "worker.log", + "formatter": "rxn", + "maxBytes": 10 * 1024 * 1024, + "backupCount": 5, + }, + "celery_file": { + "class": "logging.handlers.RotatingFileHandler", + "filename": "celery.log", + "formatter": "rxn", + "maxBytes": 10 * 1024 * 1024, + "backupCount": 5, + }, + }, + "loggers": { + "celery": {"level": "DEBUG", "handlers": ["celery_file"], "propagate": False}, + "amqp": {"level": "DEBUG", "handlers": ["celery_file"], "propagate": False}, + "kombu": {"level": "DEBUG", "handlers": ["celery_file"], "propagate": False}, + "": {"level": "INFO", "handlers": ["worker_file"]}, + }, + }) + + +def run_task(task: Task, name: str, fn: Callable[[], T]) -> T: + """ + Wraps the running and error handling of a celery task. + """ + logger.info("Starting %s", name) + + started = time.time() + task.update_state(state="STARTED") + try: + result = fn() + finished = time.time() + return { + "result": result, + "time": finished - started, + } + except Exception as exception: + logger.exception("Error in %s", name) + raise + + +@app.task(bind=True) +def product_prediction(self: Task, reactants_list: List[str], **kwargs): + def callback(): + return product_prediction_body(reactants_list, **kwargs) + + return run_task(task=self, name="product_prediction", fn=callback) + + +def product_prediction_body(reactants_list: List[str], **kwargs): + logger.info("product_prediction: reactants_list=%s, kwargs=%s", reactants_list, kwargs) + + # Checkpoint and vocab for the forward model. Should use the included one. + ckpt_forward = kwargs.get("ckpt_forward", "Pistachio2025Q2-Forward") + ckpt_forward_path = kwargs.get("ckpt_forward_path", f"models/forward/{ckpt_forward}.ckpt") + vocab = kwargs.get("vocab", "Pistachio2025Q2") + vocab_path = kwargs.get("vocab_path", f"vocab/{vocab}.txt") + # Forward algorithm parameters + nbeams = kwargs.get("num_beams", 3) + topn = kwargs.get("topn", 1) + # Device used for predicting, defaults to "cuda" if available + device = kwargs.get("device", None) + + # Forward model used for product predictions + forward_lit = load_forward_model( + ckpt_path=ckpt_forward_path, + vocab_path=vocab_path, + num_beams=nbeams, + topn=topn, + device=device, + ) + + # Return the prediction results for a batch of product predictions + return forward_reaction( + smiles_list=reactants_list, + lit=forward_lit, + ) + + +@app.task(bind=True) +def retro_prediction(self: Task, product: str, **kwargs): + def callback(): + return retro_prediction_body(product, **kwargs) + + return run_task(task=self, name="retro_prediction", fn=callback) + + +def retro_prediction_body(product: str, **kwargs): + logger.info("retro_prediction: product=%s, kwargs=%s", product, kwargs) + + # Checkpoints to the retrosynthesis and forward models. Should use the included models. + ckpt_retro = kwargs.get("ckpt_retro", "Pistachio2025Q2-Retro") + ckpt_retro_path = kwargs.get("ckpt_retro_path", f"models/retrosynthesis/{ckpt_retro}.ckpt") + ckpt_forward = kwargs.get("ckpt_forward", "Pistachio2025Q2-Forward") + ckpt_forward_path = kwargs.get("ckpt_forward_path", f"models/forward/{ckpt_forward}.ckpt") + # Vocab path for the retrosynthesis and forward models. For the defaults models the vocab is the same. + vocab = kwargs.get("vocab", "Pistachio2025Q2") + vocab_path = kwargs.get("vocab_path", f"vocab/{vocab}.txt") + # Retrosynthesis algorithm parameters + topn = kwargs.get("topn", 15) # Number of retro prediction results per product SMILES + fap = kwargs.get("fap", 0.6) # Forward likelihood acceptance probability (not length averaged) + fld = kwargs.get("dfl", 0.2) # Forward likelihood delta between the top2 forward prediction results + nbeams = max(kwargs.get("num_beams", 15), topn) # Must never be lower than topn + # Device used for predicting, defaults to "cuda" if available + device = kwargs.get("device", None) + + + # Retro model used for retrosynthesis prediction + retro_lit = load_retro_model( + ckpt_path=ckpt_retro_path, + vocab_path=vocab_path, + num_beams=nbeams, + topn=topn, + device=device, + ) + + # Forward model used to get the forward likelihood and top‑2 forward prediction with the predicted reactants + forward_lit = load_forward_model( + ckpt_path=ckpt_forward_path, + vocab_path=vocab_path, + num_beams=2, + topn=2, + device=device, + ) + + # Return the prediction results for a single retrosynthesis step + return retro_single_step( + product=product, + retro_lit=retro_lit, + forward_lit=forward_lit, + fap=fap, + fld=fld, + ) + + +# Predicts the retrosynthesis reaction for a given product +def retro_single_step( + product: str, + retro_lit, + forward_lit, + fap: float = 0.6, + fld: float = 0.2, +): + # Canonicalize the product before doing retrosynthesis + product = canonicalize_smiles(product) + + # Get retrosynthesis prediction results + retro_results = retro_reaction( + smiles=product, + lit=retro_lit, + ) + + # Empty list for accepted results + accepted_results = [] + + # Loop retro prediction results + for result in retro_results: + + logger.info("Evaluating retro result: %s", result) + + # Get the predicted reactants + reactants = result.get("reactants") + + # Step 1: Canonicalize reactants + # If it fails to canonicalize -> skips to next result + try: + reactants = canonicalize_smiles(reactants) + except InvalidSmiles as e: + logger.info("Invalid SMILES '%s': %s", reactants, e) + continue + + # Step 2: Skips to next result if product is among the reactants + if product in reactants.split("."): + logger.info("Product appears among reactants, skipping.") + continue + + # Step 3: Compute forward likelihood + likelihood = reaction_likelihood(forward_lit, reactants, product) + + # Step 4: Results with strong likelihood (above the fap that was set) are added to the accepted results list + if likelihood > fap: + accepted_results.append(result) + continue + + # Step 5: Generate top‑2 forward predictions + forward_preds = forward_reaction( + smiles_list=[reactants], + lit=forward_lit, + ) + + # If it wasn't able to generate the 2 results, it skips to the next result + if len(forward_preds) < 2: + continue + + top1, top2 = forward_preds[0], forward_preds[1] + + # Step 6: The top‑1 prediction must match target product + if top1["product"] != product: + continue + + # Step 7: The different between the top1 and top2 likelihood bust be higher than fld to accept the result + likelihood_top1 = reaction_likelihood(forward_lit, reactants, top1["product"]) + likelihood_top2 = reaction_likelihood(forward_lit, reactants, top2["product"]) + if likelihood_top1 > likelihood_top2 + fld: + accepted_results.append(result) + + # Return the accepted results + return accepted_results + + +@app.task(bind=True) +def retro_prediction_tree(self: Task, product: str, **kwargs): + def callback(): + return retro_prediction_tree_body(product, **kwargs) + + return run_task(task=self, name="retro_prediction_tree", fn=callback) + + +def retro_prediction_tree_body(product: str, **kwargs): + logger.info("retro_prediction_tree: product=%s, kwargs=%s", product, kwargs) + + # Checkpoints to the retrosynthesis and forward models. Should use the included models. + ckpt_retro = kwargs.get("ckpt_retro", "Pistachio2025Q2-Retro") + ckpt_retro_path = kwargs.get("ckpt_retro_path", f"models/retrosynthesis/{ckpt_retro}.ckpt") + ckpt_forward = kwargs.get("ckpt_forward", "Pistachio2025Q2-Forward") + ckpt_forward_path = kwargs.get("ckpt_forward_path", f"models/forward/{ckpt_forward}.ckpt") + # Vocab path for the retrosynthesis and forward models. For the defaults models the vocab is the same. + vocab = kwargs.get("vocab", "Pistachio2025Q2") + vocab_path = kwargs.get("vocab_path", f"vocab/{vocab}.txt") + # Retrosynthesis algorithm parameters + topn = kwargs.get("topn", 15) # Number of retro prediction results per product SMILES + fap = kwargs.get("fap", 0.6) # Forward likelihood acceptance probability (not length averaged) + fld = kwargs.get("dfl", 0.2) # Forward likelihood delta between the top2 forward prediction results + nbeams = max(kwargs.get("num_beams", 15), topn) # Must never be lower than topn + # Retrosynthesis tree beam search parameters + max_depth = kwargs.get("max_depth", 3) # This is the max depth of the retrosynthesis tree + beam_width = kwargs.get("beam_width", 5) # This is the max amount of nodes being expanded in each step + # Device used for predicting, defaults to "cuda" if available + device = kwargs.get("device", None) + + # Retro model used for retrosynthesis prediction + retro_lit = load_retro_model( + ckpt_path=ckpt_retro_path, + vocab_path=vocab_path, + num_beams=nbeams, + topn=topn, + device=device, + ) + + # Forward model used to get the forward likelihood and top‑2 forward prediction with the predicted reactants + forward_lit = load_forward_model( + ckpt_path=ckpt_forward_path, + vocab_path=vocab_path, + num_beams=2, + topn=2, + device=device, + ) + + # Return a retrosynthesis tree using a custom beam search algorithm and retro predictions + return retrosynthesis_beam_search( + product=product, + retro_lit=retro_lit, + forward_lit=forward_lit, + max_steps=max_depth, + beam_width=beam_width, + fap=fap, + fld=fld, + ) + + +def retrosynthesis_beam_search( + product: str, + retro_lit, + forward_lit, + beam_width: int = 5, # The node limit for expansion each step + max_steps: int = 5, # The maximum depth of the retrosynthesis tree + fap: float = 0.6, + fld: float = 0.2, + is_available_fn=None, # Availability function can be added (future-proofing) +): + # Canonicalize product smiles + product = canonicalize_smiles(product) + + # List of beam paths to expand + # We start with a single node witch corresponds to the product from where all retrosynthesis predictions start from + beams = [{ + "open_nodes": {product}, # The product is the only node molecule at the start + "arcs": [], # Arcs are the list of reactions. No arcs yet since there aren't any predictions yet + "score": 0.0, # Sum of all arc log scores. Starts with zero + "expanded": set(), # Molecule nodes already expanded. None yet + "steps": 0, # Current step. Starts at step 0, can go until max_steps/max_depth + }] + + # List of completed beam paths + completed = [] + + # Loop max_steps/max_depth + for _ in range(max_steps): + + # New beam path candidates + candidates = [] + + # Loop paths in the list of beam paths to expand + for path in beams: + + # If a function to check the commercial availability of the molecules exists then check if all molecules + # in the list of nodes are available. If everything is available the beam path stops here as there is no + # need to expand anything further + if is_available_fn and all(is_available_fn(m) for m in path["open_nodes"]): + completed.append(path) + continue + + # Creates a list of expandable nodes. The expandable nodes are the molecules that weren't expanded already + # and that are not solvents. They also need to follow the active branch path. + # On the first iteration the product is the only expandable node + expandable = [ + m for m in path["open_nodes"] + if m not in path["expanded"] and not is_spectator(m) + and ("active_branch" not in path or m in path["active_branch"]) + ] + + # If there's nothing to expand, then the beam path is complete and is added to the completed beam paths list + if not expandable: + completed.append(path) + continue + + # Get the higher complexity molecule from the expandable nodes + scores = {m: simplicity(m) for m in expandable} + molecule = min(scores, key=scores.get) + + # Expand the molecule (retrosynthesis prediction) + expansions = expand_molecule( + molecule, + retro_lit=retro_lit, + forward_lit=forward_lit, + fap=fap, + fld=fld, + ) + + # If there were no retrosynthesis prediction results, then the beam can't be expanded any further + # The beam path is added to the completed beam paths list, and it skips to the next beam path + if not expansions: + completed.append(path) + continue + + # Loop the retrosynthesis prediction results (arcs) from the expansion + for arc in expansions: + + # Cycle detection + # If any predicted reactant matches the original product or a molecule, then + # there is a loop and the arc is ignored + if any(r == product for r in arc["reactants"]): + continue + + # Create a new beam path candidate + new_path = { + # The new beam open nodes are the previous open nodes minus the molecule that was expanded + # The molecule removed is replaced by the reactants resultant from the retrosynthesis prediction + "open_nodes": (path["open_nodes"] - {molecule}) | set(arc["reactants"]), + # The arcs are the reactions that make the beam path + # Here we update it adding the new reaction + "arcs": path["arcs"] + [{**arc}], + # The beam score is the sum of all the arc log scores, here it sums the existing score with the + # arc log score with 1e-12 serving as log-flooring to prevent errors if the score is zero + "score": path["score"] + math.log(arc["arc_score"] + 1e-12), + # The expanded molecule is added to the expanded list + "expanded": path["expanded"] | {molecule}, + # The step is incremented to keep track on how deep in the tree it currently is + "steps": path["steps"] + 1, + # Makes the retrosynthesis tree go down in a single branch instead of multiple + # Only do retrosynthesis on latest iteration reactants + "active_branch": set(arc["reactants"]), + } + + # The new beam path is added to the list of candidates + candidates.append(new_path) + + # If there are no new candidates the beam search is over + if not candidates: + break + + # Only keep unique candidates by checkin the route signature + # If there is a duplicate, keep the one with the highest score + unique = {} + for candidate in candidates: + sig = route_signature(candidate) + if sig not in unique or candidate["score"] > unique[sig]["score"]: + unique[sig] = candidate + + # Keep only the best beam_width results for next expansion loop or to add to the final results if on last loop + beams = nlargest( + beam_width, + unique.values(), + key=lambda p: p["score"] + ) + + # Add remaining beams to the completed beam paths list + completed.extend(beams) + + # Only keep unique candidates by checkin the route signature + # If there is a duplicate, keep the one with the highest score + # We deduplicate again here since it could reach the same routes in different ways + unique = {} + for p in completed: + sig = route_signature(p) + if sig not in unique or p["score"] > unique[sig]["score"]: + unique[sig] = p + + # Update the completed list with only the unique values + completed = list(unique.values()) + + # Return the beam search retrosynthesis tree from all the predictions + return [ + # The results are serialized before being returned + # All sets are converted to lists and the log sum is converted to a probability score + serialize_pathway(p) + # The final completed beam paths are sorted from highest to lowest score + for p in sorted(completed, key=lambda p: p["score"], reverse=True) + ] + + +# Expands a beam path by doing retrosynthesis on a molecule +def expand_molecule( + product: str, # Molecule to expand (retrosynthesis) + retro_lit, # Retrosynthesis model + forward_lit, # Forward model + fap: float, # Forward likelihood acceptance probability (not length averaged) + fld: float, # Forward likelihood delta between the top2 forward prediction results +): + """ + Generate solvent-agnostic retrosynthetic steps for a product. + Reactions differing only by solvents are deduplicated, keeping + the highest-scoring representative. + """ + + # Single step retro prediction + retro_predictions = retro_single_step( + product=product, + retro_lit=retro_lit, + forward_lit=forward_lit, + fap=fap, + fld=fld, + ) + + # Product simplicity score + sc = simplicity(product) + + # Hold the unique reactions after deduplication + unique = {} + + # Loop prediction results from retrosynthesis + for prediction in retro_predictions: + + # Canonicalize reactants + reactants = canonicalize_smiles(prediction["reactants"]).split(".") + reactants_smiles = ".".join(reactants) + + # Get forward likelihood + p = reaction_likelihood(forward_lit, reactants_smiles, product) + + # Reactants simplicity score + sr = 1.0 + for r in reactants: + sr *= simplicity(r) + + # Score of the overall retro prediction + arc_score = p * (sr / sc) + + # Create the object for this arc candidate (retrosynthesis reaction) + candidate = { + "product": product, + "reactants": reactants, + "forward_likelihood": p, + "arc_score": arc_score, + "confidence": prediction.get("confidence"), + "metadata": prediction, + } + + # Only keep unique candidates by checkin the reaction signature + # If there is a duplicate, keep the one with the highest score + # We strip the solvents on the reaction signature so that they don't influence the reactions uniqueness + sig = reaction_signature(product, reactants) + if sig not in unique or arc_score > unique[sig]["arc_score"]: + unique[sig] = candidate + + # Return the list of unique reaction + return list(unique.values()) + + +# Serializes the beam paths so they can be properly sent via the celery worker +def serialize_pathway(path): + return { + "open_nodes": list(path["open_nodes"]), + "expanded": list(path["expanded"]), + "score": math.exp(path["score"]), # Converted to probability + "steps": path["steps"], + "arcs": path["arcs"], + } + +# Route signature function to ensure all routes are unique +def route_signature(path): + """ + Canonical, hashable signature of a retrosynthesis route. + """ + # This will return a route signature with the following format: + # ( + # ( + # "", + # ("", "", ..., "") + # ), + # ( + # "", + # ("", "", ..., "") + # ), + # ... + # ) + return tuple( + (arc["product"], tuple(sorted(arc["reactants"]))) + for arc in path["arcs"] + ) + +# Route signature function to ensure all reactions are unique +# We remove the solvents before creating a signature since their presence does not equate to a different reaction +def reaction_signature(product, reactants): + """ + Canonical representation of a reaction ignoring solvents. + """ + # This will return a route signature with the following format: + # ( + # "", + # ("", "", ..., "") + # ) + core_reactants = tuple(strip_neutral_spectators(reactants)) + return product, core_reactants diff --git a/worker/transformers_model/__init__.py b/worker/transformers_model/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/worker/transformers_model/configuration.py b/worker/transformers_model/configuration.py new file mode 100644 index 0000000000000000000000000000000000000000..acf2549c894fa87c6f023acc0cd611eebd1c5adf --- /dev/null +++ b/worker/transformers_model/configuration.py @@ -0,0 +1,90 @@ +"""Vanilla transformer configuration.""" + +from transformers.configuration_utils import PretrainedConfig + + +class VanillaTransformerConfig(PretrainedConfig): + """VanillaTransformerConfig implementation.""" + + model_type = "vanilla_transformer" + + def __init__( + self, + vocabulary_size: int = 569, + vocab_size: int = 569, + embedding_dim: int = 256, + ffnn_hidden_dim: int = 2048, + dropout: float = 0.1, + activation: str = "relu", + num_attention_heads: int = 8, + num_encoder_layers: int = 4, + num_decoder_layers: int = 4, + num_hidden_layers: int = 4, + attention_mask: float = float("-inf"), + init_std: float = 0.02, + max_position_embeddings: int = 5000, + pad_token_id: int = 0, + bos_token_id: int = 12, + eos_token_id: int = 13, + decoder_start_token_id: int = 12, + is_encoder_decoder: bool = True, + return_dict_in_generate: bool = False, + num_beams: int = 1, + max_length: int = 278, + device: str = "cpu", + **kwargs + ) -> None: + """Initialization of the configuation. + + Args: + vocabulary_size: number of different tokens that can be represented. Defaults to 569. + vocab_size: number of different tokens that can be represented. Defaults to 569. + embedding_dim: dimensionality of the layers and the pooler layer. Defaults to 256. + ffnn_hidden_dim: Dimensionality of the feed-forward layer in language modeling head. Defaults to 2048. + dropout: the dropout probability for all fully connected layers in the embeddings, encoder, and pooler. Defaults to 0.1. + activation: activation function to be used in the encoder and decoder layers. Defaults to "relu". + num_attention_heads: number of attention heads for each attention layer in the Transformer encoder and decoder. Defaults to 8. + num_encoder_layers: number of encoder layers. Defaults to 4. + num_decoder_layers: number of decoder layers. Defaults to 4. + attention_mask: value used to mask attention values. Defaults to -inf. + init_std: standard deviation of the truncated_normal_initializer for initializing all weight matrices. Defaults to 0.02. + max_position_embeddings: maximum length used to compute the positional encoding. Defaults to 5000. + pad_token_id: token value reserved for padding. Defaults to 0. + bos_token_id: token value reserved for the beginning of a sentence. Defaults to 2. + eos_token_id: token value reserved for the end of a sentence. Defaults to 3. + decoder_start_token_id: token value given to the decoder as a start point for the generation. Defaults to 2. + is_encoder_decoder: defines the structure of the model for the generation routine. Defaults to True. + return_dict_in_generate: whether or not to return a [`ModelOutput`] instead of a plain tuple during generation. Defaults to False. + num_beams: number of beams for beam search that will be used by default in the `generate` method of the model (1 means no beam search). Default to 1. + max_length: maximum length of the sequence to be generated. Defaults to 278. + device: device on which the model will be allocated. Defaults to "cpu". + """ + + # model hyperparameters configuration + self.vocabulary_size = vocabulary_size + self.vocab_size = vocab_size + self.embedding_dim = embedding_dim + self.ffnn_hidden_dim = ffnn_hidden_dim + self.dropout = dropout + self.activation = activation + self.num_attention_heads = num_attention_heads + self.num_encoder_layers = num_encoder_layers + self.num_decoder_layers = num_decoder_layers + self.num_hidden_layers = num_hidden_layers + self.attention_mask = attention_mask + self.init_std = init_std + self.max_position_embeddings = max_position_embeddings + self.device = device + + super().__init__( + model_type="vanilla_transformer", + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + is_encoder_decoder=is_encoder_decoder, + return_dict_in_generate=return_dict_in_generate, + decoder_start_token_id=decoder_start_token_id, + max_length=max_length, + num_beams=num_beams, + **kwargs, + ) diff --git a/worker/transformers_model/logging_utils.py b/worker/transformers_model/logging_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..fad3262a653118c797685a34ff2b558e6a4fc101 --- /dev/null +++ b/worker/transformers_model/logging_utils.py @@ -0,0 +1,42 @@ +"""Utility functions for the logging.""" + +import logging +from typing import Optional + +LOGGING_LEVELS = { + "debug": logging.DEBUG, + "info": logging.INFO, + "warning": logging.WARNING, + "error": logging.ERROR, + "critical": logging.CRITICAL, +} + + +def setup_logger( + logger_name: str, logging_level: str = "info", logging_file: Optional[str] = None +) -> logging.Logger: + """Configures the logging for the application. + + Args: + name (str): The name of the logger. + logging_level (str): The logging level. Defaults to "info". + logging_file (Optional[str]): The file to log to. Defaults to None. + + Returns: + logging.Logger: The configured logger. + + Raises: + ValueError: If the logging level is invalid. + """ + logging_level = logging_level.lower() + if logging_level not in LOGGING_LEVELS: + raise ValueError( + f"Invalid logging level: {logging_level}. Available levels: {LOGGING_LEVELS.keys()}" + ) + LOGGER = logging.getLogger(logger_name) + logging.basicConfig( + level=LOGGING_LEVELS[logging_level], + format="%(name)s - %(levelname)s - %(message)s", + filename=logging_file, + ) + return LOGGER diff --git a/worker/transformers_model/model.py b/worker/transformers_model/model.py new file mode 100644 index 0000000000000000000000000000000000000000..2faff984228667a844456c1d99f1458542e0a288 --- /dev/null +++ b/worker/transformers_model/model.py @@ -0,0 +1,497 @@ +"""Pytorch Lightning implementation for VanillaTransformer.""" + +import math +import json +import Levenshtein +import lightning as L +import torch +import torch.nn.functional as F +import torch.optim as optim +from typing import Dict +from torch import Tensor +from rdkit import Chem, DataStructs, RDLogger +from rdkit.Chem.rdFingerprintGenerator import GetMorganGenerator +from functools import lru_cache +from transformers import get_cosine_schedule_with_warmup +from types import MethodType +from transformers_model.smiles_tokenizer import SmilesTokenizer +from transformers_model.configuration import VanillaTransformerConfig +from transformers_model.transformer import VanillaTransformer +from transformers_model.logging_utils import setup_logger + +logger = setup_logger(__name__, logging_level="info") +RDLogger.DisableLog('rdApp.*') + + +class LitVanillaTransformer(L.LightningModule): + """Pytorch lightning model for VanillaTransformer.""" + + def __init__( + self, + vocab_path: str, + task: str, + learning_rate: float = 0.0001, + adam_beta1: float = 0.9, + adam_beta2: float = 0.98, + adam_epsilon: float = 1e-9, + adam_weight_decay: float = 0.01, + embedding_dim: int = 256, + ffnn_hidden_dim: int = 2048, + dropout: float = 0.1, + activation: str = "relu", + num_attention_heads: int = 8, + num_encoder_layers: int = 4, + num_decoder_layers: int = 4, + attention_mask: float = float("-inf"), + init_std: float = 0.02, + max_position_embeddings: int = 5000, + num_beams: int = 1, + topn: int = 1, + max_length: int = 256, + warmup_ratio: float = 0.06, + predictions_path: str = "data/predictions.jsonl", + ignore_nan: bool = False, + device: str = "cuda" + + ) -> None: + """Construct an LM lightning module. + """ + super().__init__() + + self.save_hyperparameters() + + self.model: VanillaTransformer + self.tokenizer = SmilesTokenizer(vocab_path) + self.task = task + self.learning_rate = learning_rate + self.adam_beta1 = adam_beta1 + self.adam_beta2 = adam_beta2 + self.adam_epsilon = adam_epsilon + self.adam_weight_decay = adam_weight_decay + self.embedding_dim = embedding_dim + self.ffnn_hidden_dim = ffnn_hidden_dim + self.dropout = dropout + self.activation = activation + self.num_attention_heads = num_attention_heads + self.num_encoder_layers = num_encoder_layers + self.num_decoder_layers = num_decoder_layers + self.attention_mask = attention_mask + self.init_std = init_std + self.max_position_embeddings = max_position_embeddings + self.num_beams = num_beams + self.topn = max(self.num_beams, topn) + self.max_length = max_length + self.warmup_ratio = warmup_ratio + self.predictions_path = predictions_path + self.ignore_nan = ignore_nan + self.device_pref = device + self.invalid_pred = 0 + + self.init_model() + + def init_model(self) -> None: + """Initialize a VanillaTransformer.""" + + config_args = { + "pad_token_id": self.tokenizer.pad_token_id, + "eos_token_id": self.tokenizer.sep_token_id, + "bos_token_id": self.tokenizer.cls_token_id, + "decoder_start_token_id": self.tokenizer.cls_token_id, + "vocabulary_size": self.tokenizer.vocab_size, + "vocab_size": self.tokenizer.vocab_size, + "embedding_dim": self.embedding_dim, + "ffnn_hidden_dim": self.ffnn_hidden_dim, + "dropout": self.dropout, + "activation": self.activation, + "num_attention_heads": self.num_attention_heads, + "num_encoder_layers": self.num_encoder_layers, + "num_decoder_layers": self.num_decoder_layers, + "attention_mask": self.attention_mask, + "init_std": self.init_std, + "max_position_embeddings": self.max_position_embeddings, + "num_beams": self.num_beams, + "device": self.device_pref, + } + config = VanillaTransformerConfig(**config_args) + self.model = VanillaTransformer(config) + + def forward(self, x: Tensor) -> Tensor: # type: ignore + """Forwards through the model. + """ + return self.model(**kwargs) + + @lru_cache() + def total_steps(self): + return len(self.trainer.datamodule.train_dataloader()) // self.trainer.accumulate_grad_batches * self.trainer.max_epochs + + def configure_optimizers( + self, + ) -> Dict[str, object]: + """Create and return the optimizer. + + Returns: + output (dict of str: Any): + - optimizer: the optimizer used to update the parameter. + """ + + # definition of the optimizer + optimizer = optim.AdamW( + params=self.parameters(), + lr=self.learning_rate, + betas=(self.adam_beta1, self.adam_beta2), + eps=self.adam_epsilon, + weight_decay=self.adam_weight_decay, + ) + + total_steps = self.total_steps() + print("Total steps: ", total_steps) + warmup_steps = int(self.total_steps() * self.warmup_ratio) + print("Warmup steps: ", warmup_steps) + + scheduler = get_cosine_schedule_with_warmup( + optimizer, + num_warmup_steps=warmup_steps, + num_training_steps=total_steps, + ) + + output = { + "optimizer": optimizer, + "lr_scheduler": { + "scheduler": scheduler, + "interval": "step", + "frequency": 1 + } + } + + return output # type: ignore + + def training_step(self, batch: Dict[str, Tensor], batch_idx: int) -> Tensor: # type: ignore + """ + Training step which encompasses the forward pass and the computation of the loss value. + + Args: + batch: dictionary containing the input_ids and the attention_type. + batch_idx: index of the current batch, unused. + + Returns: + loss computed on the batch. + """ + loss = self.model(**batch).loss # type:ignore + self.log("train_loss", loss, on_step=False, on_epoch=True, logger=True) + + + current_lr = self.trainer.optimizers[0].param_groups[0]['lr'] + self.log("current_lr", current_lr, on_step=False, on_epoch=True, logger=True) + + # stop training if loss is NaN + if not self.ignore_nan: + if torch.isnan(loss): + raise ValueError("Train loss is NaN") + + return loss + + def on_validation_epoch_start(self): + if self.task == "forward": + self.invalid_pred = 0 + + def validation_step(self, batch: Dict[str, Tensor], batch_idx: int) -> Tensor: # type: ignore + """ + Validation step which encompasses the forward pass and the computation of the loss value. + + Args: + batch: dictionary containing the input_ids and the attention_type. + batch_idx: index of the current batch, unused. + + Returns: + loss computed on the batch. + """ + + # compute loss + outputs = self.model(**batch) + loss = outputs.loss + self.log("val_loss", loss, on_step=False, on_epoch=True, logger=True) + + # stop training if loss is NaN + if not self.ignore_nan: + if torch.isnan(loss): + raise ValueError("Validation loss is NaN") + + + sources = batch["encoder_input_ids"] + targets = batch["decoder_input_ids"] + + # generating predictions + predictions = self.model.generate( + sources, + do_sample=False, + max_length=self.max_length, + num_beams=self.num_beams, + ) + + # decode predictions and labels into untokenized strings + predicted_texts = self.tokenizer.batch_decode(predictions, skip_special_tokens=True) + predicted_texts = self.tokenizer.batch_detokenize_smiles_string(predicted_texts) + target_texts = self.tokenizer.batch_decode(targets, skip_special_tokens=True) + target_texts = self.tokenizer.batch_detokenize_smiles_string(target_texts) + + # calculate accuracy + correct = sum(pred == tgt for pred, tgt in zip(predicted_texts, target_texts)) + accuracy = correct / len(target_texts) + self.log("val_accuracy", accuracy, prog_bar=True, on_step=False, on_epoch=True, logger=True) + + # compute Levenshtein distance + levenshtein_distances = [ + Levenshtein.distance(pred, tgt) + for pred, tgt in zip(predicted_texts, target_texts) + ] + avg_levenshtein_distance = sum(levenshtein_distances) / len(levenshtein_distances) + self.log("val_levenshtein_distance_avg", avg_levenshtein_distance, on_step=False, on_epoch=True, logger=True) + + # compute Levenshtein ratio + levenshtein_ratios = [ + Levenshtein.ratio(pred, tgt) + for pred, tgt in zip(predicted_texts, target_texts) + ] + avg_levenshtein_ratio = sum(levenshtein_ratios) / len(levenshtein_ratios) + self.log("val_levenshtein_ratio_avg", avg_levenshtein_ratio, on_step=False, on_epoch=True, logger=True) + + if self.task == "forward": + similarities = [] + generator = GetMorganGenerator(radius=2, fpSize=2048) + for pred, tgt in zip(predicted_texts, target_texts): + mol1 = Chem.MolFromSmiles(pred) + mol2 = Chem.MolFromSmiles(tgt) + + if mol1 is None or mol2 is None: + if mol1 is None: + self.invalid_pred += 1 + continue + + fp1 = generator.GetFingerprint(mol1) + fp2 = generator.GetFingerprint(mol2) + sim = DataStructs.TanimotoSimilarity(fp1, fp2) + similarities.append(sim) + + avg_valid_tanimoto = sum(similarities) / len(similarities) if len(similarities) > 0 else 0 + avg_total_tanimoto = sum(similarities) / len(predicted_texts) + self.log("val_tanimoto_valid_avg", avg_valid_tanimoto, on_step=False, on_epoch=True, logger=True) + self.log("val_tanimoto_total_avg", avg_total_tanimoto, on_step=False, on_epoch=True, logger=True) + return {"val_loss": loss, "val_accuracy": accuracy, "val_levenshtein_distance_avg": avg_levenshtein_distance, "val_levenshtein_ratio_avg": avg_levenshtein_ratio, "val_tanimoto_valid_avg": avg_valid_tanimoto, "val_tanimoto_total_avg": avg_total_tanimoto} + else: + return {"val_loss": loss, "val_accuracy": accuracy, "val_levenshtein_distance_avg": avg_levenshtein_distance, "val_levenshtein_ratio_avg": avg_levenshtein_ratio} + + def on_validation_epoch_end(self): + if self.task == "forward": + self.log("val_tanimoto_invalid_pred", self.invalid_pred) + self.invalid_pred = 0 + + def test_step(self, batch: Dict[str, Tensor], batch_idx: int) -> float: # type: ignore + """ + Test step which encompasses the forward pass and the computation of the accuracy and the loss value. + + Args: + batch: dictionary containing the input_ids and the attention_type. + batch_idx: index of the current batch, unused. + + Returns: + accuracy computed on the batch. + """ + + sources = batch["encoder_input_ids"] + targets = batch["decoder_input_ids"] + + # generating predictions + predictions = self.model.generate( + sources, + do_sample=False, + max_length=self.max_length, + num_beams=self.num_beams, + ) + + # decode predictions and labels into untokenized strings + predicted_texts = self.tokenizer.batch_decode(predictions, skip_special_tokens=True) + predicted_texts = self.tokenizer.batch_detokenize_smiles_string(predicted_texts) + target_texts = self.tokenizer.batch_decode(targets, skip_special_tokens=True) + target_texts = self.tokenizer.batch_detokenize_smiles_string(target_texts) + + # calculate accuracy + correct = sum(pred == tgt for pred, tgt in zip(predicted_texts, target_texts)) + accuracy = correct / len(target_texts) + self.log("accuracy", accuracy, prog_bar=True, on_step=False, on_epoch=True, logger=True) + + # compute Levenshtein distance + levenshtein_distances = [ + Levenshtein.distance(pred, tgt) + for pred, tgt in zip(predicted_texts, target_texts) + ] + avg_levenshtein_distance = sum(levenshtein_distances) / len(levenshtein_distances) + self.log("levenshtein_distance_avg", avg_levenshtein_distance, on_step=False, on_epoch=True, logger=True) + + # compute Levenshtein ratio + levenshtein_ratios = [ + Levenshtein.ratio(pred, tgt) + for pred, tgt in zip(predicted_texts, target_texts) + ] + avg_levenshtein_ratio = sum(levenshtein_ratios) / len(levenshtein_ratios) + self.log("levenshtein_ratio_avg", avg_levenshtein_ratio, on_step=False, on_epoch=True, logger=True) + + if self.task == "forward": + similarities = [] + generator = GetMorganGenerator(radius=2, fpSize=2048) + for pred, tgt in zip(predicted_texts, target_texts): + mol1 = Chem.MolFromSmiles(pred) + mol2 = Chem.MolFromSmiles(tgt) + + if mol1 is None or mol2 is None: + if mol1 is None: + self.invalid_pred += 1 + continue + + fp1 = generator.GetFingerprint(mol1) + fp2 = generator.GetFingerprint(mol2) + sim = DataStructs.TanimotoSimilarity(fp1, fp2) + similarities.append(sim) + + avg_valid_tanimoto = sum(similarities) / len(similarities) if len(similarities) > 0 else 0 + avg_total_tanimoto = sum(similarities) / len(predicted_texts) + self.log("tanimoto_valid_avg", avg_valid_tanimoto, on_step=False, on_epoch=True, logger=True) + self.log("tanimoto_total_avg", avg_total_tanimoto, on_step=False, on_epoch=True, logger=True) + return {"accuracy": accuracy, "levenshtein_distance_avg": avg_levenshtein_distance, "levenshtein_ratio_avg": avg_levenshtein_ratio, "tanimoto_valid_avg": avg_valid_tanimoto, "tanimoto_total_avg": avg_total_tanimoto} + else: + return {"accuracy": accuracy, "levenshtein_distance_avg": avg_levenshtein_distance, "levenshtein_ratio_avg": avg_levenshtein_ratio} + + def on_predict_start(self): + # Clear the predictions file at the start of prediction + with open("predictions.jsonl", "w") as f: + f.truncate(0) # or just pass if you want to create/overwrite + + def predict_step(self, batch: Dict[str, Tensor], batch_idx: int) -> Dict: # type: ignore + """ + Predict step. + + Args: + batch: dictionary containing the input_ids and the attention_type. + batch_idx: index of the current batch, unused. + + Returns: + Predictions + """ + + sources = batch["encoder_input_ids"] + + # generating the predicted sequence + outputs = self.model.generate( + sources, + do_sample=False, + max_length=self.max_length, + num_beams=self.num_beams, + num_return_sequences=self.topn, + return_dict_in_generate=True, + output_scores=True, + ) + + predictions = outputs.sequences # shape: (B * top_n, T) + scores = outputs.sequences_scores + + cross_mats_all = self._collect_cross_attentions(sources, predictions) + + # decode predictions and labels into untokenized strings + predicted_texts = self.tokenizer.batch_decode(predictions, skip_special_tokens=True) + predicted_texts = self.tokenizer.batch_detokenize_smiles_string(predicted_texts) + target_texts = self.tokenizer.batch_decode(sources, skip_special_tokens=True) + target_texts = self.tokenizer.batch_detokenize_smiles_string(target_texts) + + results = [] + + # write predictions to file + with open(self.predictions_path, "a") as f: + for i, target_text in enumerate(target_texts): + start = i * self.topn + end = start + self.topn + for j in range(start, end): + result = { + "source": target_text, + "predicted_target": predicted_texts[j], + "confidence": scores[j].item() #math.exp(float(scores[j].item())) if isinstance(scores, torch.Tensor) else math.exp(float(scores[j])), + } + + if cross_mats_all and j < len(cross_mats_all): + # If JSON size becomes large, you may choose only the last layer: cross_mats_all[j][-1] + result["cross_attentions"] = cross_mats_all[j] + + # f.write(json.dumps(result) + "\n") + results.append(result) + + # return predictions + return results + + def _collect_cross_attentions(self, sources: torch.Tensor, predictions: torch.Tensor): + """ + Collect decoder→encoder cross-attention for each returned hypothesis. + Returns: List[List[torch.Tensor]] + out_all[b][l] is a tensor of shape [tgt_len, src_len] for hypothesis b and decoder layer l. + """ + model = self.model # VanillaTransformer + device = sources.device + + # 1) Encode sources (batch = original B) + enc_out = model.encode(input_ids=sources) + memory = enc_out.last_hidden_state # [B, src_len, dim] + + # 2) Match memory batch to predictions batch (B_eff = B * topn) + B_eff = predictions.size(0) + B_mem = memory.size(0) + if B_mem != B_eff: + memory = memory.expand(B_eff, -1, -1) # cheap view; use .repeat if necessary + + # 3) Temporarily override cross-attn forward to request weights + layers = list(model.decoder.decoder.layers) + orig_forwards = [] + for layer in layers: + mha = layer.multihead_attn + orig_forward = mha.forward + + def forward_with_weights(self_mha, query, key, value, **kwargs): + kwargs['need_weights'] = True + kwargs.setdefault('average_attn_weights', True) # -> [B, tgt_len, src_len] + out = orig_forward(query, key, value, **kwargs) + layer._last_cross_attn = out[1] + return out + + mha.forward = MethodType(forward_with_weights, mha) + orig_forwards.append(orig_forward) + + try: + # 4) Teacher forcing pass on generated sequences + _ = model.decoder( + input_ids=predictions, # [B_eff, tgt_len] + encoder_output=memory, # [B_eff, src_len, dim] + padding_mask=None, + ) + finally: + # 5) Restore original forward functions + for layer, orig in zip(layers, orig_forwards): + layer.multihead_attn.forward = orig + + # 6) Build per-hypothesis per-layer tensors + + out_all = [] + num_layers = len(layers) + + for b in range(B_eff): + last_mat = None + # search from last layer backwards for robustness + for li in range(num_layers - 1, -1, -1): + attn = getattr(layers[li], "_last_cross_attn", None) + if attn is not None and attn.numel() > 0: + # average_attn_weights=True -> [B_eff, tgt_len, src_len] + last_mat = attn[b].detach().cpu() # [tgt_len, src_len] + break + + if last_mat is None: + # fall back to an empty tensor (or you can raise a warning) + last_mat = torch.empty(0, 0) + + out_all.append(last_mat) + + return out_all # List[Tensor], each [tgt_len, src_len] diff --git a/worker/transformers_model/runtime_cli.py b/worker/transformers_model/runtime_cli.py new file mode 100644 index 0000000000000000000000000000000000000000..1165be66000e29c5f81113f748b387f83b0eff07 --- /dev/null +++ b/worker/transformers_model/runtime_cli.py @@ -0,0 +1,77 @@ + +# worker/rs2s_tasks/runtime_cli.py +import os +import torch +from lightning.pytorch.cli import LightningCLI +from lightning.pytorch import Trainer +from pathlib import Path +from transformers_model.model import LitVanillaTransformer +from transformers_model.smiles_datamodule import LitSmilesDataset + +FORWARD_CKPT_PATH = os.getenv( + "FORWARD_CKPT_PATH", + "rs2s_tasks/epoch=59-val_accuracy=0.6476.ckpt" +) + +REQUESTED_DEVICE = os.getenv("DEVICE", "").lower() +def _choose_accelerator(): + if REQUESTED_DEVICE == "cuda" and torch.cuda.is_available(): + return "gpu" + return "cpu" + +_CLI = None +_TRAINER = None + +def get_cli_runtime(): + """Instantiate LightningCLI once and cache model/datamodule/trainer.""" + global _CLI, _TRAINER + if _CLI is None: + # Build CLI programmatically (no argv parsing). We mimic your predict setup. + _CLI = LightningCLI( + model_class=LitVanillaTransformer, + datamodule_class=LitSmilesDataset, + subclass_mode_model=False, + subclass_mode_data=False, + run=False, # <--- do not launch fit/test/predict automatically + args=[ + f"--model.vocab_path={str(Path(__file__).with_name('vocab.txt'))}", + "--model.task=forward", + "--model.device=cpu", + f"--data.vocab_path={str(Path(__file__).with_name('vocab.txt'))}", + "--data.batch_size=512", + "--seed_everything=42", + ], + ) + + # Load weights exactly like CLI predict + _CLI.model = LitVanillaTransformer.load_from_checkpoint( + checkpoint_path=FORWARD_CKPT_PATH, + strict=False, + map_location="cpu", # safe default + vocab_path=str(Path(__file__).with_name("vocab.txt")), + task="forward", + device="cpu", + ) + _CLI.model.eval().freeze() + + _CLI.datamodule = LitSmilesDataset( + batch_size=512, + vocab_path=str(Path(__file__).with_name("vocab.txt")), + ) + + # Build a lightweight Trainer for predict calls + _TRAINER = Trainer( + accelerator="cpu", #_choose_accelerator(), + devices=1, + logger=False, + enable_checkpointing=False, + inference_mode=True, + ) + + # Place model on the selected device (after construction) + #device = torch.device("cuda" if _TRAINER.accelerator.strategy.root_device.type == "cuda" else "cpu") + device = "cpu" + _CLI.model.to(device).eval() + torch.set_grad_enabled(False) + + return _CLI, _TRAINER diff --git a/worker/transformers_model/smiles_datamodule.py b/worker/transformers_model/smiles_datamodule.py new file mode 100644 index 0000000000000000000000000000000000000000..fae742ae41bf1ecd9a81913655ed09cbbd7747fb --- /dev/null +++ b/worker/transformers_model/smiles_datamodule.py @@ -0,0 +1,288 @@ +"""MIT Smiles dataset routines-filtering, dataset building.""" + +import json +import os +from pathlib import PosixPath +from typing import Dict, List, Union, Optional + +import lightning as L +from torch import Tensor +from torch.utils.data import ConcatDataset, DataLoader, Dataset +from transformers import default_data_collator +from transformers.tokenization_utils_base import BatchEncoding + +from transformers_model.smiles_tokenizer import SmilesTokenizer +from transformers_model.logging_utils import setup_logger + +logger = setup_logger(__name__, logging_level="info") + + +class SmilesDataset(Dataset): + """Smiles dataset class.""" + + def __init__( + self, + filepath: str, + tokenizer: SmilesTokenizer, + padding_idx: int = 0, + ) -> None: + """Initialize the LM data module. + + Args: + filepath: path where the dataset is located. + tokenizer: tokenize function to be used in the module. + """ + + self.filepath = filepath + if not self.filepath.endswith(".jsonl") and not self.filepath.endswith(".json"): + raise ValueError(f"{filepath} is not a .jsonl or a json.") + self.tokenizer = tokenizer + self.length = SmilesDataset.count_examples(filepath) + self.padding_idx = padding_idx + + self.examples = self.examples_reader() + + def examples_reader(self) -> List[Dict[str, str]]: + """Read instances from a filepath. + + Returns: + list of instances. + """ + with open(self.filepath) as fp: + return [json.loads(line.strip()) for line in fp] + + @staticmethod + def count_examples(filepath: str) -> int: + """Count instances of a filepath. + + Args: + filepath: path of the dataset. + + Returns: + number of examples existed in the given filepath. + """ + + def _make_gen(reader): + while True: + b = reader(2**16) + if not b: + break + yield b + + with open(filepath, "rb") as f: + count = sum(buf.count(b"\n") for buf in _make_gen(f.raw.read)) # type: ignore + return count + + def __len__(self) -> int: + """Number of instances of the dataset. + + Returns: + number of instances + """ + return self.length + + def __getitem__(self, index) -> Dict[str, Tensor]: + """Get an item of the dataset. + + Args: + index: index of the item. + + Returns: + tokenized item. + """ + + example = self.examples[index] + source = example.get("source") + target = example.get("target") + + if source is None: + raise ValueError(f"Missing 'source' in example at index {index}") + + item = {} + source_item = self.tokenizer(source) + item["encoder_input_ids"] = source_item["input_ids"].squeeze(0) + item["encoder_padding_mask"] = (1 - source_item["attention_mask"]).squeeze(0).bool() + + if target is not None: + target_item = self.tokenizer(target) + item["decoder_input_ids"] = target_item["input_ids"].squeeze(0) + item["decoder_padding_mask"] = (1 - target_item["attention_mask"]).squeeze(0).bool() + + return item + + +class LitSmilesDataset(L.LightningDataModule): + """Pytorch-lightning-style data module for smiles dataset.""" + + def __init__( + self, + vocab_path: str, + train_path: Optional[str] = None, + validation_path: Optional[str] = None, + test_path: Optional[str] = None, + predict_path: Optional[str] = None, + truncation: bool = True, + padding: str = "max_length", + max_length: int = 278, + return_token_type_ids: bool = False, + batch_size: int = 32, + num_dataloader_workers: int = 8, + ) -> None: + + super().__init__() + self.tokenizer = SmilesTokenizer(vocab_path) + self.train_path = train_path + self.validation_path = validation_path + self.test_path = test_path + self.predict_path = predict_path + self.truncation = truncation + self.padding = padding + self.max_length = max_length + self.return_token_type_ids = return_token_type_ids + self.batch_size = batch_size + + cpus_count = os.cpu_count() + self.num_dataloader_workers = min(num_dataloader_workers, cpus_count or 1) + + self.data_collator = default_data_collator + + def setup(self, stage: Optional[str] = None): + """Setup the data module.""" + + if stage == "fit" or stage is None: + if self.train_path and self.validation_path: + self.smiles_dataset_train = self.build_dataset(self.train_path) + self.smiles_dataset_val = self.build_dataset(self.validation_path) + logger.info(f"Train set size: {len(self.smiles_dataset_train)}") + logger.info(f"Validation set size: {len(self.smiles_dataset_val)}") + elif stage == "fit": + raise ValueError("Both train_path and validation_path must be provided for training.") + + if stage == "validation" or stage is None: + if self.validation_path: + self.smiles_dataset_val = self.build_dataset(self.validation_path) + logger.info(f"Validation set size: {len(self.smiles_dataset_val)}") + elif stage == "validation": + raise ValueError("The validation_path must be provided for validation.") + + if stage == "test" or stage is None: + if self.test_path: + self.smiles_dataset_test = self.build_dataset(self.test_path) + logger.info(f"Test set size: {len(self.smiles_dataset_test)}") + elif stage == "test": + raise ValueError("The test_path must be provided for testing.") + + if stage == "predict" or stage is None: + if self.predict_path: + self.smiles_dataset_predict = self.build_dataset(self.predict_path) + logger.info(f"Predict set size: {len(self.smiles_dataset_predict)}") + elif stage == "predict": + raise ValueError("The predict_path must be provided for running prediction.") + + def build_dataset(self, path: Union[str, PosixPath]) -> Dataset: + """Builds the dataset. + + Args: + path: path of the dataset or the directory that contains it. + + Returns: + pytorch dataset. + """ + path = str(path) + if path.endswith(".jsonl") or path.endswith(".json"): + return SmilesDataset( + path, self.tokenize_function, padding_idx=self.tokenizer.pad_token_id # type: ignore + ) + elif os.path.isdir(path): + return ConcatDataset( + datasets=[ + SmilesDataset( + os.path.join(path, filename), + self.tokenize_function, # type: ignore + padding_idx=self.tokenizer.pad_token_id, + ) + for filename in os.listdir(path) + if filename.endswith(".jsonl") or filename.endswith(".json") + ] + ) + else: + raise TypeError(f"{path} type is not supported for dataset.") + + def tokenize_function(self, example: str) -> BatchEncoding: + """Tokenize the given examples. + + Args: + examples: list of examples. + + Returns: + tokenized examples. + """ + + return self.tokenizer( # type: ignore + example, + truncation=self.truncation, + padding=self.padding, + max_length=self.max_length, + return_token_type_ids=self.return_token_type_ids, + return_tensors="pt", + ) + + def train_dataloader(self) -> DataLoader: + """Creates the dataloader for the training step. + + Returns: + pytorch dataloader. + """ + return DataLoader( + self.smiles_dataset_train, + batch_size=self.batch_size, + num_workers=self.num_dataloader_workers, + collate_fn=self.data_collator, + pin_memory=True, + shuffle=True, + drop_last=True, + ) + + def val_dataloader(self) -> DataLoader: + """Creates the dataloader for the validation step. + + Returns: + pytorch dataloader. + """ + return DataLoader( + self.smiles_dataset_val, + batch_size=self.batch_size, + num_workers=self.num_dataloader_workers, + collate_fn=self.data_collator, + pin_memory=True, + drop_last=True, + ) + + def test_dataloader(self) -> DataLoader: + """Creates the dataloader for the test step. + + Returns: + pytorch dataloader. + """ + return DataLoader( + self.smiles_dataset_test, + batch_size=self.batch_size, + num_workers=self.num_dataloader_workers, + collate_fn=self.data_collator, + pin_memory=True, + drop_last=True, + ) + + def predict_dataloader(self) -> DataLoader: + """Creates the dataloader for the predict step. + + Returns: + pytorch dataloader. + """ + return DataLoader( + self.smiles_dataset_predict, + batch_size=self.batch_size, + num_workers=self.num_dataloader_workers, + collate_fn=self.data_collator, + pin_memory=True, + ) diff --git a/worker/transformers_model/smiles_tokenizer.py b/worker/transformers_model/smiles_tokenizer.py new file mode 100644 index 0000000000000000000000000000000000000000..6b9b13b6716029a9bd3da8f32c5e37f50c6a6af7 --- /dev/null +++ b/worker/transformers_model/smiles_tokenizer.py @@ -0,0 +1,265 @@ +"""Tokenizer for smiles. +Based on: https://github.com/huggingface/transformers/blob/main/src/transformers/models/bert/tokenization_bert.py +""" + +import collections +import os +import re +from typing import Dict, List, Optional, Tuple + +from rxn.chemutils.tokenization import SMILES_TOKENIZER_PATTERN +from transformers.models.bert import BertTokenizer + +from transformers_model.logging_utils import setup_logger + +logger = setup_logger(__name__, logging_level="info") + + +class SmilesTokenizer(BertTokenizer): + """Tokenizer for smiles. + Adapted from https://github.com/huggingface/transformers. + """ + + def __init__( + self, + vocab_file: str, + unk_token: str = "[UNK]", + sep_token: str = "[SEP]", + pad_token: str = "[PAD]", + cls_token: str = "[CLS]", + mask_token: str = "[MASK]", + **kwargs, + ) -> None: + """Initializes a SmilesTokenizer. + + Args: + vocab_file: path to a SMILES character per line vocabulary file. + unk_token: unknown token. Defaults to "[UNK]". + sep_token: separator token. Defaults to "[SEP]". + pad_token: pad token. Defaults to "[PAD]". + cls_token: CLS token. Defaults to "[CLS]". + mask_token: mask token. Defaults to "[MASK]". + """ + super().__init__( + vocab_file, + unk_token=unk_token, + sep_token=sep_token, + pad_token=pad_token, + cls_token=cls_token, + mask_token=mask_token, + **kwargs, + ) + + if not os.path.isfile(vocab_file): + raise ValueError(f"Can't find a vocab file at path '{vocab_file}'.") + + self.highest_unused_index = max( + [i for i, v in enumerate(self.vocab.keys()) if v.startswith("[unused")] + ) + self.basic_tokenizer = BasicSmilesTokenizer() + self.init_kwargs["model_max_length"] = self.model_max_length + + @property + def vocab_size(self) -> int: + """Gets the vocabulary size. + + Returns: + size of the vocabulary. + """ + return len(self.vocab) + + @property + def vocab_list(self) -> List[str]: + """Gets a list of all the vocabulary tokens. + + Returns: + list of all the vocabulary tokens. + """ + return list(self.vocab.keys()) + + def _tokenize(self, text: str) -> List[str]: + """Tokenizes a text input using the basic smiles tokenizer. + + Args: + text: a textual input. + + Returns: + list of tokens. + """ + split_tokens = [token for token in self.basic_tokenizer.tokenize(text)] + return split_tokens + + def _convert_token_to_id(self, token: str) -> int: + """Converts a token to the corresponding index in the vocabulary. + + Args: + token: a token. + + Returns: + index corresponding to the token in the vocabulary. + """ + return self.vocab.get(token, self.vocab[self.unk_token]) + + def _convert_id_to_token(self, index: int) -> str: + """Converts an index to the corresponding token in the vocabulary. + + Args: + index: an index. + + Returns: + token corresponding to the index in the vocabulary. + """ + return self.ids_to_tokens.get(index, self.unk_token) + + def convert_tokens_to_string(self, tokens: List[str]) -> str: + """Converts a sequence of tokens in a single string. + + Args: + tokens: some tokens. + + Returns: + untokenized string. + """ + out_string = " ".join(tokens).replace(" ##", "").strip() + return out_string + + def add_special_tokens_ids_single_sequence(self, token_ids: List[int]) -> List[int]: + """Adds special tokens at the extremes of the token ids list. + A BERT sequence has the following format: [CLS]_id ... token_ids ... [SEP]_id + + Args: + token_ids: a list of token indexes. + + Returns: + input token id sequence with added special token to the extremes. + """ + return [self.cls_token_id] + token_ids + [self.sep_token_id] + + def add_special_tokens_single_sequence(self, tokens: List[str]) -> List[str]: + """Adds special tokens at the extremes of the tokens list. + A BERT sequence has the following format: [CLS] ... tokens ... [SEP] + + Args: + tokens: a list of tokens. + + Returns: + input token sequence with added special token to the extremes. + """ + return [self.cls_token] + tokens + [self.sep_token] + + def add_special_tokens_sequence_pair( + self, token_0: List[str], token_1: List[str] + ) -> List[str]: + """Adds special tokens at the extremes and between two tokens lists. + A BERT sequence pair has the following format: [CLS] token_0 [SEP] token_1 [SEP] + + Args: + token_0: first token sequence. + token_1: second token sequence. + + Returns: + concatenated input tokens sequences with special tokens at the extremes and between them. + """ + sep = [self.sep_token] + cls = [self.cls_token] + return cls + token_0 + sep + token_1 + sep + + def add_special_tokens_ids_sequence_pair( + self, token_ids_0: List[int], token_ids_1: List[int] + ) -> List[int]: + """Adds special tokens at the extremes and between two token ids lists. + A BERT sequence pair has the following format: [CLS] A [SEP] B [SEP] + + Args: + token_ids_0: first token indexes sequence. + token_ids_1: second token indexes sequence. + + Returns: + concatenated input tokens sequences with special tokens at the extremes and between them. + """ + sep = [self.sep_token_id] + cls = [self.cls_token_id] + return cls + token_ids_0 + sep + token_ids_1 + sep + + def add_padding_tokens( + self, token_ids: List[str], length: int, right: bool = True + ) -> List[int]: + """Adds padding tokens until the sequence reaches a length of max_length. + + By default padding tokens are added to the right of the sequence. + + Args: + token_ids: token ids sequence. + length: maximum sequence length. + right: whether to adding the padding to the right. Defaults to True. + + Returns: + padded sequence. + """ + padding = [self.pad_token_id] * (length - len(token_ids)) + if right: + return token_ids + padding + else: + return padding + token_ids + + def detokenize_smiles_string(self, smiles_string: str) -> str: + return smiles_string.replace(" ", "") + + def batch_detokenize_smiles_string(self, smiles_list: List[str]) -> List[str]: + return [smiles.replace(" ", "") for smiles in smiles_list] + + def save_vocabulary( + self, vocab_path: str, filename_prefix: Optional[str] = None + ) -> Tuple[str, ...]: + """Saves the tokenizer's vocabulary to a file. + + Args: + vocab_path: _description_. + filename_prefix: prefix for vocab name. Defaults to None. + + Returns: + path of the file containing the saved vocabulary. + """ + + index = 0 + + if filename_prefix is not None: + vocab_file = os.path.join(vocab_path, f"{filename_prefix}-vocab.txt") + else: + vocab_file = os.path.join(vocab_path, "vocab.txt") + + with open(vocab_file, "w", encoding="utf-8") as writer: + for token, token_index in sorted(self.vocab.items(), key=lambda kv: kv[1]): + if index != token_index: + logger.warning( + f"Saving vocabulary to {vocab_file}: vocabulary indices are not consecutive. Please check that the vocabulary is not corrupted!" + ) + index = token_index + writer.write(token + "\n") + index += 1 + return (vocab_file,) + + +class BasicSmilesTokenizer: + """Runs basic SMILES tokenization.""" + + def __init__(self, regex_pattern: str = SMILES_TOKENIZER_PATTERN) -> None: + """Initializes a BasicSMILESTokenizer. + + Args: + regex_pattern: a regular expression pattern. Defaults to SMILES_TOKENIZER_PATTERN. + """ + self.regex_pattern = regex_pattern + self.regex = re.compile(self.regex_pattern) + + def tokenize(self, text: str) -> List[str]: + """Tokenizes a SMILES string (with the regex pattern). + + Args: + text: SMILES string. + + Returns: + list of tokens. + """ + tokens = [token for token in self.regex.findall(text)] + return tokens diff --git a/worker/transformers_model/transformer.py b/worker/transformers_model/transformer.py new file mode 100644 index 0000000000000000000000000000000000000000..fd4e27cd235b44963e3268111e4fa543289e63b7 --- /dev/null +++ b/worker/transformers_model/transformer.py @@ -0,0 +1,761 @@ +"""Vanilla transformer model.""" + +import math +from dataclasses import dataclass +from typing import Any, Dict, List, Optional, Tuple, Union + +import torch +import torch.utils.checkpoint +from torch import Tensor, nn +from torch.nn import ( + CrossEntropyLoss, + TransformerDecoder, + TransformerDecoderLayer, + TransformerEncoder, + TransformerEncoderLayer, +) +from transformers import PreTrainedModel, GenerationMixin +from transformers.utils import ModelOutput + +from transformers_model.configuration import VanillaTransformerConfig + + +class TokenEmbedding(nn.Module): + """Token embedding layer implementation.""" + + def __init__( + self, vocabulary_size: int, embedding_dim: int, padding_idx: int + ) -> None: + """Contructs a TokenEmbedding layer. + + Args: + vocabulary_size: number of different tokens that can be represented. + embedding_dim: dimensionality of the embeddings. + padding_idx: token value reserved for padding. + """ + super().__init__() + self.embedding = nn.Embedding(vocabulary_size, embedding_dim, padding_idx) + self.embedding_dim = embedding_dim + + def forward(self, tokens: Tensor) -> Tensor: + """Creation of the embedding. + + Args: + tokens: tokens to be embedded. + + Returns: + embedding of the input tokens. + """ + return self.embedding(tokens.long()) * math.sqrt(self.embedding_dim) + + +class PositionalEncoding(nn.Module): + """Position embedding layer implementation.""" + + def __init__( + self, embedding_dim: int, dropout: float = 0.1, max_len: int = 5000 + ) -> None: + """Initializes a PositionalEncoding layer. + + Args: + embedding_dim: dimensionality of the embeddings. + dropout: dropout probability. Defaults to 0.1. + max_len: upperbound to the length of the sequence for the generation of position. Defaults to 5000. + """ + super().__init__() + + self.dropout = nn.Dropout(p=dropout) + + position = torch.arange(0, max_len).reshape(max_len, 1) + div_term = torch.exp( + -torch.arange(0, embedding_dim, 2) * math.log(10000) / embedding_dim + ) + position_embedding = torch.zeros(max_len, embedding_dim) + position_embedding[:, 0::2] = torch.sin(position * div_term) + position_embedding[:, 1::2] = torch.cos(position * div_term) + position_embedding = position_embedding.unsqueeze(0) + + self.register_buffer("position_embedding", position_embedding) + + def forward(self, token_embedding: Tensor) -> Tensor: + """Creation of the positional embedding. + + Args: + token_embedding: token embedding to be merged to the positional embedding. + + Returns: + Embedding obtained as the sum of token and positional embeddings. + """ + + position_embedding = self.position_embedding[:, :token_embedding.size(1), :] + token_embedding = token_embedding + position_embedding + return self.dropout(token_embedding) + + +class ExtendedTransformerEncoderLayer(TransformerEncoderLayer): + """TransformerEncoderLayer extension to be compliant with different torch versions.""" + + def __init__( + self, + d_model: int, + nhead: int, + dim_feedforward: int = 2048, + dropout: float = 0.1, + activation: str = "relu", + batch_first: bool = True, + device=None, + ): + """Initializes a TrasformerEncoderLayer. + + Args: + d_model: the number of expected features in the input. + nhead: the number of heads in the multiheadattention models. + dim_feedforward: the dimension of the feedforward network model. Defaults to 2048. + dropout: the dropout value. Defaults to 0.1. + activation: the activation function of the intermediate layer. Defaults to "relu". + """ + super().__init__( # type: ignore + d_model=d_model, + nhead=nhead, + dim_feedforward=dim_feedforward, + dropout=dropout, + activation=activation, + batch_first=batch_first, + device=device, + ) + + +class ExtendedTransformerDecoderLayer(TransformerDecoderLayer): + """TransformerEncoderLayer extension to be compliant with different torch versions.""" + + def __init__( + self, + d_model: int, + nhead: int, + dim_feedforward: int = 2048, + dropout: float = 0.1, + activation: str = "relu", + batch_first: bool = True, + device=None, + ): + """Initializes a TrasformerDecoderLayer. + + Args: + d_model: the number of expected features in the input. + nhead: the number of heads in the multiheadattention models. + dim_feedforward: the dimension of the feedforward network model. Defaults to 2048. + dropout: the dropout value. Defaults to 0.1. + activation: the activation function of the intermediate layer. Defaults to "relu". + """ + super().__init__( # type: ignore + d_model=d_model, + nhead=nhead, + dim_feedforward=dim_feedforward, + dropout=dropout, + activation=activation, + batch_first=batch_first, + device=device, + ) + + +class ExtendedLayerNorm(nn.LayerNorm): + """LayerNorm extension to be compliant with different torch versions.""" + + def __init__( + self, + normalized_shape: Union[int, List, torch.Size], + eps: float = 1e-5, + elementwise_affine: bool = True, + device=None, + ): + """Initializes a LayerNorm. + + Args: + normalized_shape: input shape from an expected input of size. + eps: a value added to the denominator for numerical stability. Defaults to 1e-5. + elementwise_affine: a boolean value that when set to ``True``, this module has learnable per-element affine parameters initialized to ones (for weights) and zeros (for biases). Defaults to True. + """ + super().__init__( # type: ignore + normalized_shape=normalized_shape, + eps=eps, + elementwise_affine=elementwise_affine, + device=device, + ) + + +class VanillaTransformerPretrainedModel(PreTrainedModel, GenerationMixin): + """Abstract class to handle weights initialization.""" + + config_class = VanillaTransformerConfig + base_model_prefix = "model" + + def _init_weights(self, module: nn.Module) -> None: + """Initializes the weights. + + Args: + module: layer of the model. + """ + + if isinstance(module, nn.Linear): + # initilaization of the weights of a linear module + module.weight.data.normal_(mean=0.0, std=self.config.init_std) + if module.bias is not None: + module.bias.data.zero_() + elif isinstance(module, nn.Embedding): + # initialization of the weights of an embedding module + module.weight.data.normal_(mean=0.0, std=self.config.init_std) + if module.padding_idx is not None: + module.weight.data[module.padding_idx].zero_() + elif isinstance(module, nn.LayerNorm): + # initialization of the weights of a layer normalization module + module.bias.data.zero_() + module.weight.data.fill_(1.0) + + +@dataclass +class VanillaEncoderOutput(ModelOutput): + """VanillaEncoderOutput implementation. Represents the output of the encoder block.""" + + last_hidden_state: torch.FloatTensor + + +class VanillaTransformerEncoder(VanillaTransformerPretrainedModel): + """VanillaTransformerEncoder implementation.""" + + def __init__(self, config: VanillaTransformerConfig) -> None: + """Initializes a VanillaTransformerEncoder. + + Args: + config: configuration of the VanillaTransformer. + """ + super().__init__(config) + + # set how the data is expected by torch encoder and decoder modules + self.batch_first = True + + # instanciating the embedding layer for the input tokens + self.src_token_embedding = TokenEmbedding( + vocabulary_size=config.vocabulary_size, + embedding_dim=config.embedding_dim, + padding_idx=config.pad_token_id, + ) + # instanciating the embedding layer for the position of input tokens + self.positional_encoding = PositionalEncoding( + embedding_dim=config.embedding_dim, + dropout=config.dropout, + max_len=config.max_position_embeddings, + ) + # instanciating an encoder layer + encoder_layer = ExtendedTransformerEncoderLayer( + d_model=config.embedding_dim, + nhead=config.num_attention_heads, + dim_feedforward=config.ffnn_hidden_dim, + dropout=config.dropout, + activation=config.activation, + batch_first=True, + device=config.device, + + ) + # instanciating the normalization layer of the encoder + encoder_norm = ExtendedLayerNorm( + normalized_shape=config.embedding_dim, + device=config.device, + ) + # instanciating the encoder block using `config.num_encoder_layers` encoder layers + self.encoder = TransformerEncoder( + encoder_layer=encoder_layer, + num_layers=config.num_encoder_layers, + norm=encoder_norm, + ) + + # initialization of the weights + self.post_init() + + def _make_encoder_masks( + self, + input: Tensor, + mask_size: int, + padding_mask: Optional[Tensor] = None, + ) -> Tuple[Tensor, Tensor]: + """Create the attention masks for the encoder. + + Args: + input: input of the encoder. + mask_size: shape of the input. + padding_mask: mask for the padding tokens. Defaults to None. + + Returns: + tuple containing: + - mask for the encoder input sequence. + - padding mask for the encoder input sequence. + """ + + mask = torch.zeros( + (mask_size, mask_size), + device=self.config.device, + ).type(torch.bool) + + if padding_mask is None: + # creating the padding mask if is not available + padding_mask = input.eq(self.config.pad_token_id) + # the padding mask is always taken with batch_first + padding_mask = ( + padding_mask if self.batch_first else padding_mask.transpose(1, 0) # type: ignore + ) + else: + # adjusting the dimensions of the padding mask, in case it is available + # the padding mask is always taken with batch_first + padding_mask = padding_mask.view(-1, padding_mask.shape[-1]) + + mask = mask.to(self.config.device) + padding_mask = padding_mask.to(self.config.device) + + return mask, padding_mask + + def forward( + self, + input_ids: torch.LongTensor, + padding_mask: Optional[Tensor] = None, + **kwargs, + ) -> VanillaEncoderOutput: + """Implements a forward pass through the VanillaDecoder. + + Args: + input_ids: tokens given as input to the encoder. + + Raises: + ValueError: in case of missing input_ids. + + Returns: + a VanillaEncoderOutput. + """ + + # adjust input_ids dimension to allow not batched input i.e. of size (seq_length,) + if input_ids is None: + raise ValueError("You have to specify input_ids.") + if input_ids.dim() == 1: + input_ids = input_ids.unsqueeze(0) # Add batch dimension + + input = input_ids # Already batch_first + input_shape = input.shape + + # compute input embeddings + inputs_embeds = self.positional_encoding(self.src_token_embedding(input)) + + # create masks required by the encoder + src_mask, src_padding_mask = self._make_encoder_masks( + input=input, mask_size=input_shape[1], padding_mask=padding_mask + ) + + # compute encoder output + encoder_output = self.encoder( + inputs_embeds, + mask=src_mask, + src_key_padding_mask=src_padding_mask, + ) + + # permute the dimension to batch first in case the output is computed without batch_first + encoder_output = ( + encoder_output if self.batch_first else encoder_output.permute(1, 0, 2) + ) + + return VanillaEncoderOutput(last_hidden_state=encoder_output) + + +@dataclass +class VanillaDecoderOutput(ModelOutput): + """VanillaDecoderOutput implementation. Represents the output of the decoder block.""" + + last_hidden_state: torch.FloatTensor + + +class VanillaTransformerDecoder(VanillaTransformerPretrainedModel): + """VanillaTransformerDecoder implementation.""" + + def __init__(self, config: VanillaTransformerConfig) -> None: + """Initializes a VanillaTransformerDecoder. + + Args: + config: configuration of the VanillaTransformer. + """ + super().__init__(config) + + # set how the data is expected by torch encoder and decoder modules + self.batch_first = True + + # instanciating the embedding layer for the input tokens + self.tgt_token_embedding = TokenEmbedding( + config.vocabulary_size, config.embedding_dim, config.pad_token_id + ) + + # instanciating the embedding layer for the position of input tokens + self.positional_encoding = PositionalEncoding( + config.embedding_dim, + dropout=config.dropout, + max_len=config.max_position_embeddings, + ) + + # instanciating an decoder layer + decoder_layer = ExtendedTransformerDecoderLayer( + d_model=config.embedding_dim, + nhead=config.num_attention_heads, + dim_feedforward=config.ffnn_hidden_dim, + dropout=config.dropout, + activation=config.activation, + batch_first=True, + device=config.device, + ) + + # instanciating the normalization layer of the decoder + decoder_norm = ExtendedLayerNorm( + normalized_shape=config.embedding_dim, + device=config.device, + ) + + # instanciating the decoder block using `config.num_encoder_layers` decoder layers + self.decoder = TransformerDecoder( + decoder_layer=decoder_layer, + num_layers=config.num_decoder_layers, + norm=decoder_norm, + ) + + # initializing the weights + self.post_init() + + def _make_square_subsequent_mask( + self, mask_size: int, device: torch.device + ) -> Tensor: + """Creates the Look-ahead mask, needed to mask future inputs for the decoder. + + Args: + mask_size: size of the mask. + device: device on which the mask tensor will be allocated. + + Returns: + the look-ahead mask. + """ + return torch.triu( + torch.ones(mask_size, mask_size) * self.config.attention_mask, + diagonal=1, + ).to(device) + + def _make_decoder_masks( + self, + input: Tensor, + mask_size: int, + padding_mask: Optional[Tensor] = None, + ) -> Tuple[Tensor, Tensor]: + """Create the masks for the decoder. + + Args: + input: input of the decoder. + mask_size: shape of the input. + padding_mask: mask for the padding tokens. Defaults to None. + + Returns: + tuple containing: + - mask for the decoder input sequence. + - padding mask for the decoder input sequence. + """ + + # creating the look-ahead mask + mask = self._make_square_subsequent_mask(mask_size, self.config.device) + + if padding_mask is None: + # creating the padding mask if is not available + padding_mask = input.eq(self.config.pad_token_id) + # the padding mask is always taken with batch_first + padding_mask = ( + padding_mask if self.batch_first else padding_mask.transpose(1, 0) # type: ignore + ) + else: + # adjusting the dimensions padding mask, in case it is available + # the padding mask is always taken with batch_first + padding_mask = padding_mask.view(-1, padding_mask.shape[-1]) + + mask = mask.bool() + padding_mask = padding_mask.bool() + + mask = mask.to(self.config.device) + padding_mask = padding_mask.to(self.config.device) + + return mask, padding_mask + + def forward( + self, + input_ids: torch.LongTensor, + encoder_output: Optional[torch.FloatTensor] = None, + padding_mask: Optional[Tensor] = None, + **kwargs, + ) -> VanillaDecoderOutput: + """Implements a forward pass through the VanillaDecoder. + + Args: + input_ids: tokens given as input to the decoder. + encoder_output: hidden state created by the encoder. Defaults to None. + + Raises: + ValueError: in case of missing input_ids. + + Returns: + A VanillaDecoderOutput. + """ + + # adjusting input_ids dimension to allow not batched input i.e. of size (seq_length,) + if input_ids is None: + raise ValueError("You have to specify decoder_input_ids.") + + + if input_ids.dim() == 1: + input_ids = input_ids.unsqueeze(0) + + input = input_ids # already batch_first + input_shape = input.shape + + # computing input embeddings + inputs_embeds = self.positional_encoding(self.tgt_token_embedding(input)) + + # creating masks required by the decoder + tgt_mask, tgt_padding_mask = self._make_decoder_masks( + input=input, + mask_size=input_shape[1], + padding_mask=padding_mask, + ) + + # computing decoder's output + decoder_output = self.decoder( + inputs_embeds, + encoder_output, + tgt_mask=tgt_mask, + # memory_mask=memory_mask, + tgt_key_padding_mask=tgt_padding_mask, + # memory_key_padding_mask=memory_key_padding_mask + ) + + # permuting the dimension to batch_first in case the output is computed without batch_first + decoder_output = ( + decoder_output if self.batch_first else decoder_output.permute(1, 0, 2) + ) + + return VanillaDecoderOutput( + last_hidden_state=decoder_output, + ) + + +@dataclass +class VanillaTransformerOutput(ModelOutput): + """VanillaDecoderOutput implementation. + + Args: + loss: language modeling loss. Must be the first argument for huggingface Trainer. + logits: prediction scores of the language modeling head (scores for each vocabulary token before SoftMax). + """ + + loss: Optional[torch.Tensor] = None + logits: torch.Tensor = None # type: ignore + + +class VanillaTransformer(VanillaTransformerPretrainedModel): + """VanillaTransformer implementation.""" + + def __init__(self, config: VanillaTransformerConfig) -> None: + """Initializes a VanillaTransformer. + + Args: + config: configuration of the VanillaTransformer. + """ + super().__init__(config) + + # save vocab size + self.vocab_size = config.vocabulary_size + # instanciating the vanilla transformer encoder + self.encoder = VanillaTransformerEncoder(config) + # instanciating the vanilla transformer decoder + self.decoder = VanillaTransformerDecoder(config) + # instanciating the prediction head + self.lm_head = nn.Linear(config.embedding_dim, config.vocabulary_size) + + # defining the loss function + self.my_loss_function = CrossEntropyLoss(ignore_index=config.pad_token_id) + #self.loss_type = "ForSequenceClassification" + + # initializing the weights (i.e. calling _init_weights of VanillaTransformerPretrainedModel) + self.post_init() + + def forward( + self, + encoder_input_ids: Optional[torch.LongTensor] = None, + decoder_input_ids: Optional[torch.LongTensor] = None, + encoder_padding_mask: Optional[Tensor] = None, + decoder_padding_mask: Optional[Tensor] = None, + encoder_output: Optional[ModelOutput] = None, + **kwargs, + ) -> VanillaTransformerOutput: + """Implements a forward pass through the VanillaTransformer. + + Args: + encoder_input_ids: tokens given as input to the encoder. Defaults to None. + decoder_input_ids: tokens given as input to the decoder. Defaults to None. + encoder_output: hidden state created by the encoder. Defaults to None. + + Raises: + ValueError: in case of missing decoder_input_ids. + ValueError: in case the encoder output is not computed. + + Returns: + a VanillaTransformerOutput. + """ + + # adjusting encoder_input_ids dimension to allow not batched input i.e. of size (seq_length,) + if encoder_input_ids is not None: + + if encoder_input_ids.dim() == 1: + encoder_input_ids = encoder_input_ids.unsqueeze(0) + + input = encoder_input_ids + input_shape = encoder_input_ids.size() + + # adjusting decoder_input_ids dimension to allow not batched input i.e. of size (seq_length,) + if decoder_input_ids is None: + raise ValueError( + "The value of decoder_input_ids must be specified. Is None" + ) + + if decoder_input_ids.dim() == 1: + decoder_input_ids = decoder_input_ids.unsqueeze(0) + + target = decoder_input_ids + target_shape = decoder_input_ids.size() + + # adjusting the dimensions of the decoder padding mask in case it is available + if decoder_padding_mask is not None: + decoder_padding_mask = decoder_padding_mask.view( + -1, decoder_padding_mask.size()[-1] + ) + + # adjusting the dimensions of the encoder padding mask in case it is available + if encoder_padding_mask is not None: + encoder_padding_mask = encoder_padding_mask.view( + -1, encoder_padding_mask.size()[-1] + ) + + if encoder_input_ids is None: + target_input = target + else: + target_input = target[:, :-1] # remove the last element of the sequence + if decoder_padding_mask is not None: + decoder_padding_mask = decoder_padding_mask[:, :-1] + + # encoding the sequence + if encoder_input_ids is not None and encoder_output is None: + encoder_output = self.encoder( + input_ids=input, padding_mask=encoder_padding_mask + ) + + # decoding the sequence + if encoder_output is not None: + decoder_output = self.decoder( + input_ids=target_input, + encoder_output=encoder_output.last_hidden_state, + padding_mask=decoder_padding_mask, + ) + else: + raise ValueError("The value of encoder_output must be computed. Is None.") + + # predicting the next token + logits = self.lm_head(decoder_output.last_hidden_state) + + # shifting the tgt by one so with the we predict the token at pos 1 + target_output = target[:, 1:] + + if encoder_input_ids is not None: + loss = self.my_loss_function( + logits.reshape(-1, logits.shape[-1]), target_output.reshape(-1) + ) + else: + loss = None + + return VanillaTransformerOutput( + loss=loss, + logits=logits, + ) + + def encode( + self, input_ids: Tensor, padding_mask: Optional[Tensor] = None + ) -> VanillaEncoderOutput: + """Encodes the input through a forward pass in the encoder of the model. + + Needed for the generation of new tokens. + + Args: + input_ids: tokens given as input to the encoder. + padding_mask: mask for the padding tokens in the input. Defaults to None. + + Returns: + the encoder output. + """ + return self.encoder(input_ids=input_ids, padding_mask=padding_mask) + + def decode( + self, + input_ids: Tensor, + encoder_output: Tensor, + padding_mask: Optional[Tensor] = None, + ) -> Tensor: + """Decode the target through a forward pass in the decoder of the model. + + Needed for the generation of new tokens. + + Args: + input_ids: tokens given as input to the decoder. + encoder_output: output of the encoder. + padding_mask: mask for the padding tokens in the input. Defaults to None. + + Returns: + the decoder output. + """ + return self.decoder( + input_ids=input_ids, + encoder_output=encoder_output, + padding_mask=padding_mask, + ) + + def prepare_inputs_for_generation( + self, decoder_input_ids: Tensor, **kwargs + ) -> Dict[str, Any]: + """Prepare the input for the foward pass of the model during generation. + + It is called during the generation, passing the input for the decoder and the encoded input sequence. + + Args: + decoder_input_ids: tokens given as input to the decoder. + + Returns: + dictionary with the arguments for a forward pass through the vanilla transformer. + """ + + return { + "encoder_input_ids": None, # encoder_outputs is defined. input_ids not needed + "encoder_output": kwargs.get("encoder_outputs"), + "decoder_input_ids": decoder_input_ids, + } + + def get_encoder(self): + """Returns the encoder of the model. + + Needed for the generation of new tokens. + + Returns: + encoder module of the vanilla transformer. + """ + return self.encoder + + def get_decoder(self): + """Returns the decoder of the model. + + Needed for the generation of new tokens. + + Returns: + decoder module of the vanilla transformer. + """ + return self.decoder diff --git a/worker/transformers_model/vocab.txt b/worker/transformers_model/vocab.txt new file mode 100644 index 0000000000000000000000000000000000000000..def7c5845e502e7a541475efb4eade8cc16d822b --- /dev/null +++ b/worker/transformers_model/vocab.txt @@ -0,0 +1,624 @@ +[PAD] +[unused1] +[unused2] +[unused3] +[unused4] +[unused5] +[unused6] +[unused7] +[unused8] +[unused9] +[unused10] +[UNK] +[CLS] +[SEP] +[MASK] +c +C +( +) +O +1 +2 += +N +. +n +3 +F +Cl +4 +- +[C@H] +[C@@H] +S +# +Br +[O] +5 +[nH] +/ +[N+] +s +[K] +[O-] +[Na] +P +o +[Si] +B +[Cl] +[Pd] +6 +I +[C@] +[P] +[C@@] +[OH] +[Cs] +[H] +[Li] +\ +[CH2] +[C] +[N-] +[Na+] +[P-] +7 +[NaH] +[NH4+] +[Cl-] +[I] +[Fe] +[Cu] +[Mg] +[BH4-] +[CH3] +[Br] +[N] +[n+] +8 +[BH-] +[Zn] +[Sn] +[Li+] +[B-] +[Al] +[c] +[c-] +[BH3-] +[CH] +[Fe+2] +[F-] +[P+] +[Br-] +[OH-] +[O+] +[AlH4-] +[I-] +[S@] +[Pd+2] +[F] +[C-] +[Ni] +[SiH] +[S@@] +9 +[S] +[Ti] +[Ag] +[K+] +[AlH] +[Mn] +[NH3+] +[Pt] +[nH+] +[Cr] +[I+3] +[PH] +[Ru] +[S+] +[Ca] +[NH2+] +[NH+] +[Se] +[Os] +%10 +[NH] +[H+] +[se] +[Cu+2] +[SH] +[SiH2] +[Rh] +[Cl+3] +[Cl+] +[P@] +[SiH3] +[H-] +[P@@] +[PH+] +[S-] +[Co] +[B] +[Hg] +[SnH] +[NH2] +%11 +[Mg+2] +[I+] +[Ce] +[Pb] +[NH-] +[Mo] +[Ge] +[N@@] +[cH-] +[PH2] +[C+] +[In] +[W] +[n] +%12 +[s+] +[n-] +[N@] +[Zr] +[OH2+] +[Au] +[Ar] +[As] +[Ba] +[CH-] +[Ir] +[o+] +[Bi] +[CH2-] +[AlH2-] +%13 +[Ag+] +[Ba+2] +[KH] +[S@@+] +[Sb] +[OH3+] +[Sc+3] +p +[LiH] +[Os-2] +[Hg+2] +[S@+] +[Rh-3] +[Pd+] +[Cu+] +[te] +[Fe-3] +[V] +[OH+] +[Ca+2] +[Te] +[Cs+] +[Re] +[La] +[Ce+3] +[SiH4] +[Yb+3] +[Si@H] +[PH4+] +[Zr+2] +[Sm] +%14 +[Rh+2] +[I+2] +[Si-] +[Pd-] +[Sb-] +[O-2] +[BH] +[CH3-] +[Zn+2] +[SeH] +[Gd] +[Rh+3] +[Yb] +[N@+] +[SH+] +[Br+2] +[S-2] +[Pt-2] +[Ir+3] +[Th] +[CH+] +[Cu-2] +[Cr+3] +[Al+] +[N@@+] +%15 +[PH3+] +[Hf] +[CH2+] +[Al-] +[Pd-2] +[Ga] +%16 +[Si@] +[Cl+2] +[NH2-] +[Co+2] +[Gd+3] +[Bi+3] +[AlH-] +[Y+3] +[cH+] +[AlH3] +[Rb] +[Sn+2] +[Cd] +[Hg+] +[Zr+4] +[Ru-2] +[SH-] +[Ni+2] +[Si@@] +[SnH4] +[SH2] +[c+] +[Ti+3] +[S+2] +[Pb+2] +[Cr+2] +[PH3] +[SnH3] +[Cu-] +%17 +[Ir+] +[Sr] +[PdH2] +[Rh+] +[Sr+2] +[P+3] +[Y] +[Ru+2] +[Tl+] +[Zr+3] +[Nb] +[B+] +[Al+3] +[Au+] +[V+2] +[IH2+] +[Eu] +[B+3] +[IH] +%18 +[Br+] +[Pt+2] +[Sc] +[Cd+2] +%19 +[SnH2] +[Ce+4] +[SH3+] +[Fe+3] +[CaH2] +%20 +[Ta] +[Mg+] +[AlH2] +[Tl] +[Tl+3] +[Hf+2] +[IH+] +[Sb+3] +[Si+] +[Xe] +[H+2] +[PH5] +[Zn+] +[TeH] +[ClH+] +[Er+3] +[Sm+3] +b +[SH2+] +[sH+] +%21 +[Eu+3] +[Ac] +[Lu] +[YH] +%22 +[Ti+2] +[Tb] +[Si-2] +[Be] +[se+] +[N+3] +[Ti+4] +[Ga+3] +[pH] +[N@@H+] +[Pr] +[Au-] +[Ir-] +[U+2] +[N@H+] +[Mn+2] +[InH2] +[U] +%23 +[Ho] +[RuH] +[si] +[PH2+] +[GeH] +[F+] +%24 +[Si@@H] +[SeH2] +[At] +[Nd] +[BH2-] +[Be+2] +[Fm] +[Nd+3] +[RuH2] +[Dy+3] +[SeH-] +[Ru+] +[siH] +[B+2] +[Tc] +[LaH] +[Zr+] +[Si+4] +[P@@H] +[p] +[BH2] +[AlH2+] +[IH-] +[NiH] +[Ir+2] +[As-] +[As+] +[Dy] +%25 +[BrH+] +[Na-] +[B@-] +[Ni-2] +[IH3] +[Fe-4] +[PbH2] +[GeH2] +[SiH-] +[SH4] +[FH+] +[Hf+3] +[Y-] +[UH] +[Er] +[Co+3] +[Hf+4] +[Tb+3] +[AsH2] +[VH] +[PH4] +[Sn-] +[BH+] +[IrH] +[B@@-] +[SbH2] +[Pr+3] +[Se-] +[Se+] +[Ni+] +[Se@+] +[Ge-] +[Bi+] +[PH-] +[C+4] +[Po] +[Nb+5] +[Cm] +[GeH3] +[La+3] +[ClH2+] +[NaH3] +[V+5] +[Mg-] +[Sn+4] +[In+3] +[K-] +[Rh-2] +[Ac-] +[Ce+2] +[SeH+] +[W+] +[p+] +[Au-2] +[Ag+3] +[AsH] +[Fe-] +[p-] +[Tl+2] +[Sn+] +[CuH] +[Mn+4] +[SH5] +[Se@@+] +[H-2] +[Li-2] +[Ga+2] +[Yb+2] +[Ru-] +[CH3+] +[WH] +[Sn-2] +[Pt-] +[Mn+3] +[Bi+2] +[Te+] +[Rh-] +[Sg] +[Y+] +[NaH+] +[SnH2-] +[Au+3] +[Cr+] +[Co+] +[Fr] +[He] +[Mo+2] +[PH4-] +[TeH3] +[RgH] +[P@H] +[P-3] +[W+4] +[Cf] +[Mn+] +[Ga+] +[MgH] +[InH] +[SH3] +[Ta+2] +[IrH2] +[Ni-] +[OH+2] +[Mo+4] +[Re+] +[te+] +[N-2] +[N+2] +[IH2] +[Db] +[Pt+4] +[P@H+] +[ClH+2] +[Al+2] +[Re+4] +[TiH] +[BiH3] +[AsH4+] +[As+3] +[C-4] +[Zn-] +[AuH] +[Xe+] +[Sc+2] +[Fe+] +[CaH] +[PtH2] +[Cl-2] +[PtH] +[oH+] +[Lu+3] +[I-2] +[Sb+2] +[NaH4] +[RuH3] +[F-2] +[VH2] +[Pa] +[BiH2] +[AsH3] +[Br-2] +[Ca+] +[C-2] +[YH4] +[Si+2] +[Se+2] +[TeH2] +[CoH] +[Sb+] +[YH3] +[Cn] +[PH2-] +[Os+2] +[Cr-] +[Ru+3] +[YH2] +[Na-2] +[GeH4] +[V+4] +[Na+2] +[InH3] +[SbH] +[BrH2+] +[Fe-2] +[P+5] +[Ra] +[No] +[Te-] +[Ba+] +[Tm] +[GaH] +[CuH2] +[NH3+2] +[CaH3] +[Ne] +[Se-2] +[AlH3-] +[CeH] +[Mg-2] +[Li-] +[Hs] +[N-3] +[Ta+3] +[NiH2] +[Ru+6] +[YH5] +[LiH+] +[Rf] +[CuH2-] +[Pd-3] +[Cd+] +[OH2+2] +[W-2] +[FH2+] +[BrH+2] +[V+3] +[Ta-] +[C+3] +[MnH] +[BaH] +[Tm+3] +[C+2] +[Ti+] +[P+2] +[Ge+4] +[RaH] +[Sc+] +[PtH+] +[Gd+2] +[CmH5] +[Ce+] +[Pu] +[Am] +[NaH+2] +[AsH+] +[Ac+3] +[Ag-2] +[Ho+3] +[Co-2] +[Bh-] +[TiH2] +[Ta+5] +[Tc+4] +[TaH3] +[Ag+2] +[FH+2] +[BaH2] +[PtH+2] +[Hf+] +[Br+3] +[SbH3] +[AcH] +[NH+3] +[YH7] +[Pt+] +[Eu+2] +[Fe+4] +[Sm+2] +[Es] diff --git a/worker/utils.py b/worker/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..7e41a4dcb56b33a9d5b7b5e1f613c6fdf3137087 --- /dev/null +++ b/worker/utils.py @@ -0,0 +1,157 @@ +"""Utilities used in the RXN forward/retrosynthesis workers.""" +import logging +from typing import Optional + +from rxn.chemutils.conversion import ( + inchi_to_mol, + mol_to_inchi, + mol_to_smiles, + smiles_to_mol, +) +from rxn.chemutils.exceptions import InvalidInchi, InvalidSmiles + +logger = logging.getLogger(__name__) +logger.addHandler(logging.NullHandler()) + +RXN_SMILES_SEPARATOR = ">>" + + +def standardize_smiles( + smiles: str, + canonicalize: bool = True, + sanitize: bool = True, + find_radicals: bool = True, + inchify: bool = False, +) -> str: + """ + Ensure that a SMILES follows a desired standard. + + It allows canonicalization, sanitization and inchification keeping stereochemistry with isomericSmile=True. + It can process multiple molecules separated by ".". + Note that inchify set to True will also canonicalize the molecule. + + Args: + smiles (str): SMILES representation of a molecule. + canonicalize (bool): canonicalize SMILES. Defaults to True. + sanitize (bool): sanitize SMILES. Defaults to True. + inchify (bool): inchify the SMILES. Defaults to False. + + Returns: + a SMILES following the desired standard. + """ + try: + molecule = smiles_to_mol(smiles, sanitize=sanitize, find_radicals=find_radicals) + except InvalidSmiles: + logger.error(f"SMILES parsing failure: {smiles}.") + raise + + if inchify: + try: + inchi_string = mol_to_inchi(molecule) + except InvalidInchi: + logger.error( + f"Inchification failure for SMILES: {smiles}. Returning its canonical version." + ) + return mol_to_smiles(molecule, isomericSmiles=True) + else: + # canonical set to True because we can't guarantee no canonicalization + try: + molecule_from_inchi = inchi_to_mol(inchi_string) + except InvalidInchi: + logger.error( + f"De-inchification failure for InChi: {inchi_string}. Returning its canonical version." + ) + return mol_to_smiles(molecule, isomericSmiles=True) + return mol_to_smiles(molecule_from_inchi, canonical=True) + if canonicalize: + return mol_to_smiles(molecule, isomericSmiles=True) + else: + return smiles + + +def standardize_molecules( + molecules: str, + canonicalize: bool = True, + sanitize: bool = True, + inchify: bool = False, + fragment_bond: str = "~", + ordered_precursors: bool = True, + molecule_token_delimiter: Optional[str] = None, + is_enzymatic: bool = False, + enzyme_separator: str = "|", +) -> str: + """ + Ensure that a set of molecules represented by a string follows a desired standard. + + Args: + molecules (str): molecules SMILES. Molecules can be separated via a ".". + Fragments are supported with a custom `fragment_bond`. + canonicalize (bool): canonicalize SMILES. Defaults to True. + sanitize (bool): sanitize SMILES. Defaults to True. + inchify (bool): inchify the SMILES. Defaults to False. + fragment_bond (str): fragment bond. Defaults to '~'. + ordered_precursors (bool): order precursors. Defaults to True. + molecule_token_delimiter (Optional[str]): delimiter for big molecule tokens. Defaults to None + is_enzymatic (bool): the molecules are representing an enzymatic reaction. Defaults to False. + enzyme_separator (str): separator for molecules and the enzyme. Defaults to '|'. + + Returns: + standardized molecules. + + Examples: + Standardize multiple molecules: + + >>> standardize_molecules('CCO.CC') + 'CC.CCO' + + Standardize multiple molecules including fragment information: + + >>> standardize_molecules('CCO.CC~C') + 'CCO.C~CC' + """ + enzyme = "" + if is_enzymatic: + splitted_molecules = molecules.split(enzyme_separator) + molecules = splitted_molecules[0] + if len(splitted_molecules) > 1: + enzyme = splitted_molecules[1] + enzyme = "{}{}".format(enzyme_separator, enzyme) + if molecule_token_delimiter is not None: + molecules = molecules.replace(molecule_token_delimiter, "") + if fragment_bond in molecules: + standardized_molecules_list = [ + # make sure we remove the fragment to have valid SMILES + standardize_smiles( + molecule.replace(fragment_bond, "."), + canonicalize=canonicalize, + sanitize=sanitize, + inchify=inchify, + ).replace(".", fragment_bond) + for molecule in molecules.split(".") + ] + if ordered_precursors: + standardized_molecules_list = sorted(standardized_molecules_list) + standardized_molecules = ".".join(standardized_molecules_list) + else: + if ordered_precursors: + # RDKit guarantees ordered precursors + standardized_molecules = standardize_smiles( + molecules, + canonicalize=canonicalize, + sanitize=sanitize, + inchify=inchify, + ) + else: + standardized_molecules_list = [ + standardize_smiles( + molecule, + canonicalize=canonicalize, + sanitize=sanitize, + inchify=inchify, + ) + for molecule in molecules.split(".") + ] + standardized_molecules = ".".join(standardized_molecules_list) + # add optional enzyme information + standardized_molecules = "{}{}".format(standardized_molecules, enzyme) + return standardized_molecules