metadata

library_name: transformers
tags: []

Mol-MoE: Performant & Steerable Multi-Objective RLHF in Drug Design

Diego Calanzone (1, 2), Pierluca D'Oro (2), Pierre-Luc Bacon (1, 2)
(1) Universite de Montreal, (2) Mila Quebec AI Institute

How to use this model

This LM is fine-tuned to generate molecules in the SMILES format wrt. desired properties. For unconditioned SMILES generation, use the BOS token <s>. For conditioned generation, you can target the following properties: JNK3, DRD2, GSK3B, CYP2D6, CYP2C19.

prompt: <JNK3=0.3><DRD2=0.7><GSK3B=0.2><CYP2D6=0.8><CYP2C19=0.8><s>

An example of the generation pipeline:

from transformers import AutoTokenizer, AutoModelForCausalLM
import re

# Setup
device = "cuda"
tokenizer = AutoTokenizer.from_pretrained("ddidacus/Mol-MoE-6x1b")
model = AutoModelForCausalLM.from_pretrained("ddidacus/Mol-MoE-6x1b")
generation_kwargs = {
    "max_new_tokens": 128,
    "min_length": -1,
    "top_k": 0.0,
    "top_p": 0.9,
    "do_sample": True,
    "pad_token_id": tokenizer.eos_token_id,
    "temperature": 1.0
}

# Inference
query = "<JNK3=0.3><DRD2=0.7><GSK3B=0.2><CYP2D6=0.8><CYP2C19=0.8><s>"
toks = tokenizer([query], return_tensors="pt")["input_ids"].to(device)
output = model.generate(toks, **generation_kwargs)
output = tokenizer.batch_decode(output)

# Parsing
filter = r'<s>(.*?)</s>'
molecule = re.findall(filter, output[0], re.DOTALL)

Model Description

This model is a fine-tuned version of LLaMa 3.2 1B through two stages:

Fine-tuning on ~3.5M molecules extracted from: ZINC 250K, MOSES, CHEMBL
RLHF-tuning using RLOO on 5 distinct reward functions from PyTDC [1]

The six LLaMa models trained in (1) and (2) are merged into mixtral blocks using MergeKit [2].

Developed by: Diego Calanzone ([email protected])
Model type: Mixtral Mixture of Experts
Finetuned from model [optional]: LLaMA 3.2 1B

Read the paper for further details.

Sources

[1] https://tdcommons.ai/single_pred_tasks/overview
[2] https://github.com/arcee-ai/mergekit