RS-mol-llama-1b

ddidacus/RS-mol-llama-1b

RS-mol-llama-1b 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

8% of all-time downloads

All-time downloads

308

Public

Parameters

1.2B

5 GB on disk

Likes

Public

Hugging Face

Repo makeup

Click a slice to open those files.

.safetensors4.9 GB · 100%

At a glance

Task: Text Generation
Library: transformers
Model type: llama
Access: Public
Created: Jan 17, 2025
Updated: Feb 20, 2025
SHA: f2d84dbf

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>. <br> For conditioned generation, please refer to the paper and the official codebase to derive different conditioned models. <br> This model is the merging result of 5 fine-tuned versions (JNK3, DRD2, GSK3B, CYP2D6, CYP2D19) with equal interpolation weight: w_i = 0.2.

An example of the generation pipeline:

from transformers import AutoTokenizer, AutoModelForCausalLM
import re

# Setup
device = "cuda"
tokenizer = AutoTokenizer.from_pretrained("ddidacus/RS-mol-llama-1b")
model = AutoModelForCausalLM.from_pretrained("ddidacus/RS-mol-llama-1b")
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 = "<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]

Developed by: Diego Calanzone ([email protected])
Model type: Decoder-only Transformer
Finetuned from model [optional]: LLaMA 3.2 1B

Read the paper for further details.

Sources

[1] https://tdcommons.ai/single_pred_tasks/overview

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