Building Custom Oracles

Step-by-step guide to creating custom oracle functions for molecule optimization

Overview

Oracles are scoring functions that evaluate molecular properties during optimization. ChemLactica uses oracles to guide the generation of molecules with desired characteristics.

Oracle Structure

A custom oracle is a Python class that implements specific methods and properties. Here’s the basic structure:

from typing import List
from chemlactica.mol_opt.utils import MoleculeEntry

class CustomOracle:
    def __init__(self, max_oracle_calls: int):
        # Maximum number of oracle calls
        self.max_oracle_calls = max_oracle_calls
        
        # Logging frequency
        self.freq_log = 100
        
        # Buffer to track unique molecules
        self.mol_buffer = {}
        
        # Maximum possible score (upper bound)
        self.max_possible_oracle_score = 1.0
        
        # Set to True if __call__ takes MoleculeEntry objects
        self.takes_entry = True
    
    def __call__(self, molecules: List[MoleculeEntry]):
        """Evaluate and return oracle scores for molecules."""
        pass
    
    def __len__(self):
        return len(self.mol_buffer)
    
    @property
    def budget(self):
        return self.max_oracle_calls
    
    @property
    def finish(self):
        return len(self.mol_buffer) >= self.max_oracle_calls

Step-by-Step Implementation

1. Define the Oracle Class

   Create a new class and initialize required attributes in __init__:

   class TPSA_Weight_Oracle:
       def __init__(self, max_oracle_calls: int):
           self.max_oracle_calls = max_oracle_calls
           self.freq_log = 100
           self.mol_buffer = {}
           self.max_possible_oracle_score = 800
           self.takes_entry = True
   

   Set takes_entry = True if your oracle needs access to RDKit molecule objects and fingerprints from MoleculeEntry.

2. Implement the Scoring Logic

   Define the __call__ method to compute scores:

   from rdkit.Chem import rdMolDescriptors
   
   def __call__(self, molecules: List[MoleculeEntry]):
       oracle_scores = []
       for molecule in molecules:
           if self.mol_buffer.get(molecule.smiles):
               # Return cached score
               oracle_scores.append(sum(self.mol_buffer[molecule.smiles][0]))
           else:
               try:
                   # Calculate TPSA
                   tpsa = rdMolDescriptors.CalcTPSA(molecule.mol)
                   oracle_score = tpsa
                   
                   # Apply constraints
                   weight = rdMolDescriptors.CalcExactMolWt(molecule.mol)
                   num_rings = rdMolDescriptors.CalcNumRings(molecule.mol)
                   
                   if weight >= 350:
                       oracle_score = 0
                   if num_rings < 2:
                       oracle_score = 0
                       
               except Exception as e:
                   print(e)
                   oracle_score = 0
               
               # Cache the score
               self.mol_buffer[molecule.smiles] = [oracle_score, len(self.mol_buffer) + 1]
               
               # Log progress
               if len(self.mol_buffer) % 100 == 0:
                   self.log_intermediate()
                   
               oracle_scores.append(oracle_score)
       return oracle_scores
   
3. Add Logging Functionality

   Implement log_intermediate() to track optimization progress:

   import numpy as np
   
   def log_intermediate(self):
       scores = [v[0] for v in self.mol_buffer.values()][-self.max_oracle_calls:]
       scores_sorted = sorted(scores, reverse=True)[:100]
       n_calls = len(self.mol_buffer)
   
       score_avg_top1 = np.max(scores_sorted)
       score_avg_top10 = np.mean(scores_sorted[:10])
       score_avg_top100 = np.mean(scores_sorted)
   
       print(f"{n_calls}/{self.max_oracle_calls} | ",
             f'avg_top1: {score_avg_top1:.3f} | '
             f'avg_top10: {score_avg_top10:.3f} | '
             f'avg_top100: {score_avg_top100:.3f}')
   
4. Implement Required Properties

   Add the required properties and methods:

   def __len__(self):
       return len(self.mol_buffer)
   
   @property
   def budget(self):
       return self.max_oracle_calls
   
   @property
   def finish(self):
       # Stopping condition for optimization
       return len(self.mol_buffer) >= self.max_oracle_calls
   

Using Your Oracle

Once your oracle is defined, use it with the optimization workflow:

import yaml
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
from chemlactica.mol_opt.optimization import optimize
from chemlactica.mol_opt.utils import set_seed

# Load configuration
config = yaml.safe_load(open("config.yaml"))

# Load model and tokenizer
model = AutoModelForCausalLM.from_pretrained(
    config["checkpoint_path"], 
    torch_dtype=torch.bfloat16
).to(config["device"])
tokenizer = AutoTokenizer.from_pretrained(
    config["tokenizer_path"], 
    padding_side="left"
)

# Initialize oracle
set_seed(42)
oracle = TPSA_Weight_Oracle(max_oracle_calls=1000)

# Set up config
config["log_dir"] = "results.log"
config["max_possible_oracle_score"] = oracle.max_possible_oracle_score

# Run optimization
optimize(model, tokenizer, oracle, config)

Oracle Configuration

Key configuration parameters for optimization (defined in your YAML config):

Oracle Config

checkpoint_path: "path/to/model"
tokenizer_path: "path/to/tokenizer"
device: "cuda:0"

# Optimization settings
pool_size: 100
validation_perc: 0.2
num_gens_per_iter: 50
generation_batch_size: 32
num_mols: 3
num_similars: 5

# Generation settings
generation_temperature: [0.8, 1.2]
generation_config:
  max_new_tokens: 512
  do_sample: true
  top_k: 50
  top_p: 0.95

# Strategy
strategy: ["rej-sample-v2"]

Example: QED Oracle

Here’s a simpler oracle that maximizes QED (drug-likeness):

from rdkit.Chem.QED import qed

class QED_Oracle:
    def __init__(self, max_oracle_calls: int):
        self.max_oracle_calls = max_oracle_calls
        self.mol_buffer = {}
        self.max_possible_oracle_score = 1.0
        self.takes_entry = True
        
    def __call__(self, molecules: List[MoleculeEntry]):
        oracle_scores = []
        for molecule in molecules:
            if molecule.smiles in self.mol_buffer:
                oracle_scores.append(self.mol_buffer[molecule.smiles][0])
            else:
                try:
                    score = qed(molecule.mol)
                except:
                    score = 0.0
                self.mol_buffer[molecule.smiles] = [score, len(self.mol_buffer) + 1]
                oracle_scores.append(score)
        return oracle_scores
    
    def __len__(self):
        return len(self.mol_buffer)
    
    @property
    def budget(self):
        return self.max_oracle_calls
    
    @property
    def finish(self):
        return len(self.mol_buffer) >= self.max_oracle_calls

Best Practices

Caching Strategy

Always cache computed scores in mol_buffer to avoid redundant calculations:

if molecule.smiles in self.mol_buffer:
    return cached_score
else:
    compute_and_cache()

Error Handling

Wrap scoring calculations in try-except blocks to handle invalid molecules:

try:
    score = calculate_property(molecule.mol)
except Exception as e:
    print(e)
    score = 0  # Return penalty score

Score Normalization

Set max_possible_oracle_score to help the model understand the score range:

# For QED (0-1 range)
self.max_possible_oracle_score = 1.0

# For custom scores, set an upper bound
self.max_possible_oracle_score = 800

Logging Frequency

Log intermediate results periodically to track progress:

if len(self.mol_buffer) % self.freq_log == 0:
    self.log_intermediate()

The takes_entry attribute determines whether __call__ receives MoleculeEntry objects (with .mol and .fingerprint attributes) or just SMILES strings.

Next Steps

Property Prediction

Learn how to fine-tune models for property prediction

Benchmarking

Evaluate your oracle’s performance