---
base_model: meta-llama/Meta-Llama-3-8B
tags:
- molecular-optimization
- chemistry
- llama-3
- grpo
- rlhf
license: apache-2.0
language:
- en
pipeline_tag: text-generation
---

# MEGA-GRPO

Fine-tuned molecular optimization model using Tanimoto-aware GRPO (Group Relative Policy Optimization) on 500K molecular transformations. Based on **Llama 3 8B**.

**Paper**: [MEGA: A Large-Scale Molecular Editing Dataset for Guided-Action Optimization](https://openreview.net/pdf?id=wzou4rm3Tt)
**Official Repository**: [https://github.com/nfsrules/MEGA-moledit](https://github.com/nfsrules/MEGA-moledit)

## Installation

```bash
pip install unsloth torch
```

## Usage

```python
from unsloth import FastLanguageModel
from unsloth.chat_templates import get_chat_template

# Configuration
max_seq_length = 1024
lora_rank = 32

# Load model
model, tokenizer = FastLanguageModel.from_pretrained(
    model_name = "nfsrulesFR/mega-grpo",
    max_seq_length = max_seq_length,
    load_in_4bit = True,
    fast_inference = True,
    max_lora_rank = lora_rank,
    gpu_memory_utilization = 0.6,
)

# Configure tokenizer
tokenizer.padding_side = 'left'
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

tokenizer = get_chat_template(
    tokenizer,
    chat_template="llama-3",
    mapping={"role": "from", "content": "value", "user": "human", "assistant": "gpt"},
)

# Generate
input_smiles = "CCO"
task = "Can you make molecule CCO more soluble in water? The output molecule should be similar to the input molecule."

messages = [{"from": "human", "value": task}]

encoded = tokenizer.apply_chat_template(
    messages,
    tokenize=True,
    add_generation_prompt=True,
    return_tensors="pt",
    padding=True,
)

outputs = model.generate(
    input_ids=encoded["input_ids"].cuda(),
    attention_mask=encoded["attention_mask"].cuda(),
    max_new_tokens=64,
    use_cache=True,
    pad_token_id=tokenizer.pad_token_id,
)

response = tokenizer.decode(outputs[0][encoded["input_ids"].shape[1]:], skip_special_tokens=True)
print(response)
```

## Supported Tasks

| Task ID | Description |
|---------|-------------|
| 101 | Increase water solubility |
| 102 | Decrease water solubility |
| 103 | Increase drug-likeness |
| 104 | Decrease drug-likeness |
| 105 | Increase permeability |
| 106 | Decrease permeability |
| 107 | Increase hydrogen bond acceptors |
| 108 | Increase hydrogen bond donors |
| 201 | Increase solubility + HBA |
| 202 | Decrease solubility + increase HBA |
| 203 | Increase solubility + HBD |
| 204 | Decrease solubility + increase HBD |
| 205 | Increase solubility + permeability |
| 206 | Increase solubility + decrease permeability |

## Model Details

- **Base Model**: Meta-Llama-3-8B
- **Training**: Tanimoto-aware GRPO on 500K molecular transformations
- **Input**: SMILES string + task description
- **Output**: Modified SMILES string

## Citation

```bibtex
@article{mega2025,
  title={MEGA: A Large-Scale Molecular Editing Dataset for Guided-Action Optimization},
  author={Fernandez, Nelson and Illouz, Maxime and Pinto, Luis and Yang, Entao and Amadou Boubacar, Habiboulaye},
  journal={Under review at International Conference on Learning Representations},
  year={2025},
  url={https://openreview.net/forum?id=wzou4rm3Tt}
}
```