Code for Reinforcement Learning-Based Tuning of Transformer Models for Chemical Reaction Prediction

Code base to replicate all the results and conclusions from the paper "Negative Chemical Data Boosts Language Models in Reaction Outcome Prediction". 

This repository contains the code used in our study on extending reinforcement learning-based tuning of language models to the chemistry domain. Specifically, we train a Transformer model for chemical reaction prediction using a rigorously controlled dataset and a high-throughput dataset with extensive reaction screenings across diverse catalyst sets and experimental conditions. Our approach demonstrates state-of-the-art performance, even when successful reactions are significantly underrepresented, by leveraging as few as twenty positive data points alongside a substantially larger negative dataset. This repository includes instructions for dataset extraction, preprocessing, and splitting, covering both controlled and high-throughput reaction datasets. The repository also details the steps for pretraining, fine-tuning with maximum likelihood estimation (MLE), training a reward model, and reinforcement learning optimization. Additionally, it explains how to generate appropriate training splits, create a joint vocabulary, and implement baseline lookup tables for RL training.

The provided framework enables effective model training even with a limited number of successful reactions, highlighting the importance of optimization strategies and negative data inclusion.