Initial commit

.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

.gitignore

@@ -0,0 +1,15 @@
+# Environment
+env
+.venv
+
+# Jupyter
+.ipynb_checkpoints
+
+# PyCharm
+.idea
+
+# Python bytecode
+__pycache__/
+*.pyc
+*.pyo
+*.pyd

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.

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")<br>Predict products from reactants | ![Retrosynthesis Prediction](img/predict_retrosynthesis.svg "Retrosynthesis Prediction")<br>Predict reactants for one step | ![Retrosynthesis Tree Prediction](img/predict_retrosynthesis_tree.svg "Retrosynthesis Tree Prediction")<br>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
+
+<details id="jupyter-notebook">
+
+<summary>Jupyter Notebook (click to expand)</summary>
+
+### 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")
+
+</details>
+
+<details id="mcp-integration-with-openwebui">
+<summary>LLM (click to expand)</summary>
+
+### 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**
+
+<img src="img/openwebui_integrations.png" alt="OpenWeb UI New Integration" width="450"/>
+
+#### 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")
+
+</details>
+
+<details id="python-scripts">
+
+<summary>Python Scripts (click to expand)</summary>
+
+### 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
+```
+</details>
+
+## 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 |

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:

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:

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"]

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(\"<p><em>No results to display.</em></p>\"))\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",
+    "    <div style=\"\n",
+    "        display: flex;\n",
+    "        flex-direction: row;\n",
+    "        justify-content: flex-start;\n",
+    "        align-items: flex-start;\n",
+    "        gap: 20px;\n",
+    "        overflow-x: auto;\n",
+    "        padding: 12px 20px;\n",
+    "        border: 1px solid #ddd;\n",
+    "        margin-bottom: 16px;\n",
+    "        box-sizing: border-box;\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 = \"<div style='font-weight:700; color:#1f77b4;'>expanded ↓</div>\" if is_expanded else \"\"\n",
+    "\n",
+    "        html += f\"\"\"\n",
+    "        <div style=\"\n",
+    "            display: flex;\n",
+    "            flex-direction: column;\n",
+    "            align-items: center;\n",
+    "            flex: 0 0 auto;\n",
+    "            min-width: 260px;\n",
+    "            padding: 6px;\n",
+    "            border: {border};\n",
+    "            border-radius: 6px;\n",
+    "        \">\n",
+    "            <img src=\"data:image/png;base64,{img}\" style=\"display:block;\"/>\n",
+    "            <div style=\"\n",
+    "                margin-top: 6px;\n",
+    "                font-size: 14px;\n",
+    "                font-weight: 700;\n",
+    "                text-align: center;\n",
+    "                word-break: break-all;\n",
+    "                font-family: monospace;\n",
+    "            \">\n",
+    "                {smi}\n",
+    "            </div>\n",
+    "            {label}\n",
+    "        </div>\n",
+    "        \"\"\"\n",
+    "\n",
+    "    html += \"</div>\"\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": [
+       "<style type=\"text/css\">\n",
+       "#T_cf7d5 caption {\n",
+       "  caption-side: top;\n",
+       "  font-weight: 700;\n",
+       "  font-size: 1.2rem;\n",
+       "  margin-bottom: 10px;\n",
+       "}\n",
+       "#T_cf7d5 th.col_heading {\n",
+       "  font-weight: 700;\n",
+       "  padding: 8px;\n",
+       "  text-align: center;\n",
+       "}\n",
+       "#T_cf7d5 td {\n",
+       "  padding: 8px;\n",
+       "  vertical-align: center;\n",
+       "}\n",
+       "#T_cf7d5 table {\n",
+       "  border-collapse: separate;\n",
+       "  border-spacing: 0px;\n",
+       "  width: 100%;\n",
+       "}\n",
+       "#T_cf7d5_row0_col0, #T_cf7d5_row1_col0, #T_cf7d5_row2_col0 {\n",
+       "  background-color: #fafafa;\n",
+       "  text-align: center;\n",
+       "  font-weight: 600;\n",
+       "  font-family: ui-monospace, SFMono-Regular, Menlo, Monaco, Consolas, 'Liberation Mono', 'Courier New', monospace;\n",
+       "}\n",
+       "#T_cf7d5_row0_col1, #T_cf7d5_row2_col1 {\n",
+       "  background-color: #006837;\n",
+       "  color: #f1f1f1;\n",
+       "  font-weight: 500;\n",
+       "}\n",
+       "#T_cf7d5_row1_col1 {\n",
+       "  background-color: #016a38;\n",
+       "  color: #f1f1f1;\n",
+       "  font-weight: 500;\n",
+       "}\n",
+       "</style>\n",
+       "<table id=\"T_cf7d5\">\n",
+       "  <caption>Product Prediction Results</caption>\n",
+       "  <thead>\n",
+       "    <tr>\n",
+       "      <th id=\"T_cf7d5_level0_col0\" class=\"col_heading level0 col0\" >Predicted Reaction</th>\n",
+       "      <th id=\"T_cf7d5_level0_col1\" class=\"col_heading level0 col1\" >confidence</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <td id=\"T_cf7d5_row0_col0\" class=\"data row0 col0\" >CCI.O=Cc1ccc([N+](=O)[O-])c(O)c1 → CCOc1cc(C=O)ccc1[N+](=O)[O-]</td>\n",
+       "      <td id=\"T_cf7d5_row0_col1\" class=\"data row0 col1\" >0.999</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td id=\"T_cf7d5_row1_col0\" class=\"data row1 col0\" >CCOc1cc(O)c(C=O)cc1OCC.OCCCBr → CCOc1cc(C=O)c(OCCCO)cc1OCC</td>\n",
+       "      <td id=\"T_cf7d5_row1_col1\" class=\"data row1 col1\" >0.993</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td id=\"T_cf7d5_row2_col0\" class=\"data row2 col0\" >C=CCc1cc(OCc2ccccc2)ccc1O.CCBr → C=CCc1cc(OCc2ccccc2)ccc1OCC</td>\n",
+       "      <td id=\"T_cf7d5_row2_col1\" class=\"data row2 col1\" >0.997</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "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": [
+       "<style type=\"text/css\">\n",
+       "#T_dbf15 caption {\n",
+       "  caption-side: top;\n",
+       "  font-weight: 700;\n",
+       "  font-size: 1.2rem;\n",
+       "  margin-bottom: 10px;\n",
+       "}\n",
+       "#T_dbf15 th.col_heading {\n",
+       "  font-weight: 700;\n",
+       "  padding: 8px;\n",
+       "  text-align: center;\n",
+       "}\n",
+       "#T_dbf15 td {\n",
+       "  padding: 8px;\n",
+       "  vertical-align: center;\n",
+       "}\n",
+       "#T_dbf15 table {\n",
+       "  border-collapse: separate;\n",
+       "  border-spacing: 0px;\n",
+       "  width: 100%;\n",
+       "}\n",
+       "#T_dbf15_row0_col0, #T_dbf15_row1_col0, #T_dbf15_row2_col0 {\n",
+       "  background-color: #fafafa;\n",
+       "  text-align: left;\n",
+       "  font-weight: 600;\n",
+       "  font-family: ui-monospace, SFMono-Regular, Menlo, Monaco, Consolas, 'Liberation Mono', 'Courier New', monospace;\n",
+       "}\n",
+       "#T_dbf15_row0_col1 {\n",
+       "  background-color: #006837;\n",
+       "  color: #f1f1f1;\n",
+       "  font-weight: 500;\n",
+       "}\n",
+       "#T_dbf15_row1_col1 {\n",
+       "  background-color: #42ac5a;\n",
+       "  color: #f1f1f1;\n",
+       "  font-weight: 500;\n",
+       "}\n",
+       "#T_dbf15_row2_col1 {\n",
+       "  background-color: #70c164;\n",
+       "  color: #000000;\n",
+       "  font-weight: 500;\n",
+       "}\n",
+       "</style>\n",
+       "<table id=\"T_dbf15\">\n",
+       "  <caption>Product Prediction Results</caption>\n",
+       "  <thead>\n",
+       "    <tr>\n",
+       "      <th id=\"T_dbf15_level0_col0\" class=\"col_heading level0 col0\" >Predicted Reaction</th>\n",
+       "      <th id=\"T_dbf15_level0_col1\" class=\"col_heading level0 col1\" >confidence</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <td id=\"T_dbf15_row0_col0\" class=\"data row0 col0\" >CCI.O=Cc1ccc([N+](=O)[O-])c(O)c1 → CCOc1cc(C=O)ccc1[N+](=O)[O-]</td>\n",
+       "      <td id=\"T_dbf15_row0_col1\" class=\"data row0 col1\" >0.999</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td id=\"T_dbf15_row1_col0\" class=\"data row1 col0\" >CCI.O=Cc1ccc([N+](=O)[O-])c(O)c1 → CCOc1cc(C=O)ccc1[N+](=O)[O-].CCOc1cc(C=O)ccc1[N+](=O)[O-]</td>\n",
+       "      <td id=\"T_dbf15_row1_col1\" class=\"data row1 col1\" >0.847</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td id=\"T_dbf15_row2_col0\" class=\"data row2 col0\" >CCI.O=Cc1ccc([N+](=O)[O-])c(O)c1 → CCOc1cc(C=O)ccc1[N+](=O)[O-].CCOc1cc(CO)ccc1[N+](=O)[O-]</td>\n",
+       "      <td id=\"T_dbf15_row2_col1\" class=\"data row2 col1\" >0.782</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "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": [
+       "<style type=\"text/css\">\n",
+       "#T_90827 caption {\n",
+       "  caption-side: top;\n",
+       "  font-weight: 700;\n",
+       "  font-size: 1.2rem;\n",
+       "  margin-bottom: 10px;\n",
+       "}\n",
+       "#T_90827 th.col_heading {\n",
+       "  font-weight: 700;\n",
+       "  padding: 8px;\n",
+       "  text-align: center;\n",
+       "}\n",
+       "#T_90827 td {\n",
+       "  padding: 8px;\n",
+       "  vertical-align: center;\n",
+       "}\n",
+       "#T_90827 table {\n",
+       "  border-collapse: separate;\n",
+       "  border-spacing: 0px;\n",
+       "  width: 100%;\n",
+       "}\n",
+       "#T_90827_row0_col0, #T_90827_row1_col0, #T_90827_row2_col0, #T_90827_row3_col0, #T_90827_row4_col0, #T_90827_row5_col0, #T_90827_row6_col0, #T_90827_row7_col0 {\n",
+       "  background-color: #fafafa;\n",
+       "  text-align: right;\n",
+       "  font-weight: 600;\n",
+       "  font-family: ui-monospace, SFMono-Regular, Menlo, Monaco, Consolas, 'Liberation Mono', 'Courier New', monospace;\n",
+       "}\n",
+       "#T_90827_row0_col1 {\n",
+       "  background-color: #33a456;\n",
+       "  color: #f1f1f1;\n",
+       "  font-weight: 500;\n",
+       "}\n",
+       "#T_90827_row1_col1 {\n",
+       "  background-color: #3ca959;\n",
+       "  color: #f1f1f1;\n",
+       "  font-weight: 500;\n",
+       "}\n",
+       "#T_90827_row2_col1 {\n",
+       "  background-color: #42ac5a;\n",
+       "  color: #f1f1f1;\n",
+       "  font-weight: 500;\n",
+       "}\n",
+       "#T_90827_row3_col1, #T_90827_row4_col1, #T_90827_row5_col1 {\n",
+       "  background-color: #4bb05c;\n",
+       "  color: #f1f1f1;\n",
+       "  font-weight: 500;\n",
+       "}\n",
+       "#T_90827_row6_col1, #T_90827_row7_col1 {\n",
+       "  background-color: #4eb15d;\n",
+       "  color: #f1f1f1;\n",
+       "  font-weight: 500;\n",
+       "}\n",
+       "</style>\n",
+       "<table id=\"T_90827\">\n",
+       "  <caption>Retrosynthesis Prediction Results</caption>\n",
+       "  <thead>\n",
+       "    <tr>\n",
+       "      <th id=\"T_90827_level0_col0\" class=\"col_heading level0 col0\" >Predicted Reaction</th>\n",
+       "      <th id=\"T_90827_level0_col1\" class=\"col_heading level0 col1\" >confidence</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <td id=\"T_90827_row0_col0\" class=\"data row0 col0\" >C1CCOC1.C=CC(=C)Cl.C[Si](C)(C)[CH2][Mg][Cl] → C=CC(=C)C[Si](C)(C)C</td>\n",
+       "      <td id=\"T_90827_row0_col1\" class=\"data row0 col1\" >0.863</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td id=\"T_90827_row1_col0\" class=\"data row1 col0\" >C1CCOC1.C=CC(=C)Cl.C[Si](C)(C)[CH2][Mg][Cl].N#N → C=CC(=C)C[Si](C)(C)C</td>\n",
+       "      <td id=\"T_90827_row1_col1\" class=\"data row1 col1\" >0.854</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td id=\"T_90827_row2_col0\" class=\"data row2 col0\" >C1CCOC1.C=CC(=C)Cl.CCOCC.C[Si](C)(C)[CH2][Mg][Cl] → C=CC(=C)C[Si](C)(C)C</td>\n",
+       "      <td id=\"T_90827_row2_col1\" class=\"data row2 col1\" >0.844</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td id=\"T_90827_row3_col0\" class=\"data row3 col0\" >C1CCOC1.C=CC(=C)Cl.C[Si](C)(C)[CH2][Mg][Cl].O.[Cl-].[NH4+] → C=CC(=C)C[Si](C)(C)C</td>\n",
+       "      <td id=\"T_90827_row3_col1\" class=\"data row3 col1\" >0.834</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td id=\"T_90827_row4_col0\" class=\"data row4 col0\" >C1CCOC1.C=CC(=C)Cl.CCOCC.C[Si](C)(C)[CH2][Mg][Cl].N#N → C=CC(=C)C[Si](C)(C)C</td>\n",
+       "      <td id=\"T_90827_row4_col1\" class=\"data row4 col1\" >0.834</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td id=\"T_90827_row5_col0\" class=\"data row5 col0\" >C1CCOC1.C=CC(=C)Cl.CCOCC.C[Si](C)(C)[CH2][Mg][Cl].O.[Cl-].[NH4+] → C=CC(=C)C[Si](C)(C)C</td>\n",
+       "      <td id=\"T_90827_row5_col1\" class=\"data row5 col1\" >0.833</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td id=\"T_90827_row6_col0\" class=\"data row6 col0\" >C1CCOC1.C=CC(=C)Cl.CCOCC.C[Si](C)(C)[CH2][Mg][Cl].Cl.O → C=CC(=C)C[Si](C)(C)C</td>\n",
+       "      <td id=\"T_90827_row6_col1\" class=\"data row6 col1\" >0.830</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td id=\"T_90827_row7_col0\" class=\"data row7 col0\" >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</td>\n",
+       "      <td id=\"T_90827_row7_col1\" class=\"data row7 col1\" >0.830</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "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": [
+       "<IPython.core.display.Markdown object>"
+      ]
+     },
+     "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": [
+       "<IPython.core.display.Markdown object>"
+      ]
+     },
+     "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": [
+       "<IPython.core.display.Markdown object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/markdown": [
+       "## Route 2\n",
+       "- **Score:** `0.371`\n",
+       "- **Steps:** `1`"
+      ],
+      "text/plain": [
+       "<IPython.core.display.Markdown object>"
+      ]
+     },
+     "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": [
+       "<IPython.core.display.Markdown object>"
+      ]
+     },
+     "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": [
+       "<IPython.core.display.Markdown object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/markdown": [
+       "## Route 3\n",
+       "- **Score:** `0.339`\n",
+       "- **Steps:** `1`"
+      ],
+      "text/plain": [
+       "<IPython.core.display.Markdown object>"
+      ]
+     },
+     "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": [
+       "<IPython.core.display.Markdown object>"
+      ]
+     },
+     "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": [
+       "<IPython.core.display.Markdown object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/markdown": [
+       "## Route 4\n",
+       "- **Score:** `0.249`\n",
+       "- **Steps:** `3`"
+      ],
+      "text/plain": [
+       "<IPython.core.display.Markdown object>"
+      ]
+     },
+     "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": [
+       "<IPython.core.display.Markdown object>"
+      ]
+     },
+     "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": [
+       "<IPython.core.display.Markdown object>"
+      ]
+     },
+     "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": [
+       "<IPython.core.display.Markdown object>"
+      ]
+     },
+     "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": [
+       "<IPython.core.display.Markdown object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/markdown": [
+       "## Route 5\n",
+       "- **Score:** `0.242`\n",
+       "- **Steps:** `3`"
+      ],
+      "text/plain": [
+       "<IPython.core.display.Markdown object>"
+      ]
+     },
+     "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": [
+       "<IPython.core.display.Markdown object>"
+      ]
+     },
+     "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": [
+       "<IPython.core.display.Markdown object>"
+      ]
+     },
+     "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": [
+       "<IPython.core.display.Markdown object>"
+      ]
+     },
+     "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": [
+       "<IPython.core.display.Markdown object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/markdown": [
+       "## Route 6\n",
+       "- **Score:** `0.233`\n",
+       "- **Steps:** `3`"
+      ],
+      "text/plain": [
+       "<IPython.core.display.Markdown object>"
+      ]
+     },
+     "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": [
+       "<IPython.core.display.Markdown object>"
+      ]
+     },
+     "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": [
+       "<IPython.core.display.Markdown object>"
+      ]
+     },
+     "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": [
+       "<IPython.core.display.Markdown object>"
+      ]
+     },
+     "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": [
+       "<IPython.core.display.Markdown object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/markdown": [
+       "## Route 7\n",
+       "- **Score:** `0.218`\n",
+       "- **Steps:** `3`"
+      ],
+      "text/plain": [
+       "<IPython.core.display.Markdown object>"
+      ]
+     },
+     "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": [
+       "<IPython.core.display.Markdown object>"
+      ]
+     },
+     "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": [
+       "<IPython.core.display.Markdown object>"
+      ]
+     },
+     "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": [
+       "<IPython.core.display.Markdown object>"
+      ]
+     },
+     "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": [
+       "<IPython.core.display.Markdown object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/markdown": [
+       "## Route 8\n",
+       "- **Score:** `0.180`\n",
+       "- **Steps:** `3`"
+      ],
+      "text/plain": [
+       "<IPython.core.display.Markdown object>"
+      ]
+     },
+     "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": [
+       "<IPython.core.display.Markdown object>"
+      ]
+     },
+     "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": [
+       "<IPython.core.display.Markdown object>"
+      ]
+     },
+     "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": [
+       "<IPython.core.display.Markdown object>"
+      ]
+     },
+     "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": [
+       "<IPython.core.display.Markdown object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/markdown": [
+       "## Route 9\n",
+       "- **Score:** `0.180`\n",
+       "- **Steps:** `1`"
+      ],
+      "text/plain": [
+       "<IPython.core.display.Markdown object>"
+      ]
+     },
+     "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": [
+       "<IPython.core.display.Markdown object>"
+      ]
+     },
+     "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": [
+       "<IPython.core.display.Markdown object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/markdown": [
+       "## Route 10\n",
+       "- **Score:** `0.176`\n",
+       "- **Steps:** `3`"
+      ],
+      "text/plain": [
+       "<IPython.core.display.Markdown object>"
+      ]
+     },
+     "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": [
+       "<IPython.core.display.Markdown object>"
+      ]
+     },
+     "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": [
+       "<IPython.core.display.Markdown object>"
+      ]
+     },
+     "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": [
+       "<IPython.core.display.Markdown object>"
+      ]
+     },
+     "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": [
+       "<IPython.core.display.Markdown object>"
+      ]
+     },
+     "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
+}

jupyter/requirements.txt

@@ -0,0 +1,7 @@
+celery[redis]
+ipykernel
+ipywidgets
+matplotlib
+notebook
+pandas
+rdkit

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"]

mcp/requirements.txt

@@ -0,0 +1,4 @@
+celery
+mcp
+mcpo
+redis

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())

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

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/<version>/` 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.

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 `<USER>` and `<PASSWORD>` 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.

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 <USER>:<PASSWORD> $BASE_URL/$RELEASE/data/pistachio.tar.gz -o pistachio.tar.gz

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)

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}'")

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

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

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

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.")
+

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}'")

model-training/requirements.txt

@@ -0,0 +1,7 @@
+lightning
+pandas
+python-Levenshtein
+rdkit
+rxn-chem-utils
+torch
+transformers

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 <input_file>
+"""
+
+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 <input_file>")
+        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

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)

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}")

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}")

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.")

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)

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.")

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.")

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()

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

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