Mol-MoE-6x1b

ddidacus/Mol-MoE-6x1b

Mol-MoE-6x1b is a text generation model from ddidacus. Use it when you need the model to write or continue text. It is set up for transformers.

Diego Calanzone (1, 2), Pierluca D'Oro (2), Pierre-Luc Bacon (1, 2) <br (1) Universite de Montreal, (2) Mila Quebec AI Institute <br arXiv: https://arxiv.org/abs/2502.05633

Downloads · 30 days

20% of all-time downloads

All-time downloads

344

Public

Parameters

5.3B

21.1 GB on disk

Likes

Public

Hugging Face

Repo makeup

Click a slice to open those files.

.safetensors21.1 GB · 100%

At a glance

Task: Text Generation
Library: transformers
Model type: mixtral
Access: Public
Created: Jan 17, 2025
Updated: Feb 20, 2025
SHA: 2dfe5629

Try a prompt

Model card

From the Hugging Face model README

Mol-MoE: Training Preference-Guided Routers for Molecule Generation

Diego Calanzone (1, 2), Pierluca D'Oro (2), Pierre-Luc Bacon (1, 2) <br> (1) Universite de Montreal, (2) Mila Quebec AI Institute <br> arXiv: https://arxiv.org/abs/2502.05633

Abstract: Recent advances in language models have enabled framing molecule generation as sequence modeling. However, existing approaches often rely on single-objective reinforcement learning, limiting their applicability to real-world drug design, where multiple competing properties must be optimized. Traditional multi-objective reinforcement learning (MORL) methods require costly retraining for each new objective combination, making rapid exploration of trade-offs impractical. To overcome these limitations, we introduce Mol-MoE, a mixture-of-experts (MoE) architecture that enables efficient test-time steering of molecule generation without retraining. Central to our approach is a preference-based router training objective that incentivizes the router to combine experts in a way that aligns with user-specified trade-offs. This provides improved flexibility in exploring the chemical property space at test time, facilitating rapid trade-off exploration. Benchmarking against state-of-the-art methods, we show that Mol-MoE achieves superior sample quality and steerability.

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 <br> [2] https://github.com/arcee-ai/mergekit

Task: Text Generation
Library: transformers
Type: mixtral
Created: Jan 17, 2025
Updated: Feb 20, 2025