---
license: apache-2.0
base_model: poolside/Laguna-XS.2
library_name: peft
pipeline_tag: text-generation
datasets:
  - poolside-laguna-hackathon/protein-ligand-design
tags:
  - lora
  - peft
  - reinforcement-learning
  - grpo
  - tool-use
  - chemistry
  - biology
  - drug-discovery
  - prime-intellect
  - verifiers
pretty_name: "Protein-Ligand Design LoRA (Team JAMMY)"
---

![Protein-ligand interaction header](header.png)

# 🧪 Protein-Ligand Design — LoRA adapter for `poolside/Laguna-XS.2`

> **poolside Laguna Hackathon submission — Team JAMMY.** A LoRA adapter trained
> with reinforcement learning (GRPO) to make `poolside/Laguna-XS.2` reason like a
> bench computational chemist / protein engineer: **measure with tools, then
> commit an answer.**

This is the trained adapter that goes with our environment and dataset:

➡️ **Gym / dataset:** [`poolside-laguna-hackathon/protein-ligand-design`](https://huggingface.co/datasets/poolside-laguna-hackathon/protein-ligand-design)

The gym hands the model a molecule or protein plus a scientist's question, and the
model must call CPU-only cheminformatics/proteomics tools (RDKit + Biopython) to
*measure* the answer before committing. The reward is **answer correctness only**,
and every ground-truth answer is computed by those same tools, so scoring is exact.

## What this adapter is

| | |
|---|---|
| Type | PEFT **LoRA** adapter (not a merged model) |
| Base model | `poolside/Laguna-XS.2` |
| Rank `r` | 16 |
| `lora_alpha` | 32 |
| `lora_dropout` | 0.0 |
| Target modules | `q_proj`, `k_proj`, `v_proj`, `o_proj`, `gate_proj`, `up_proj`, `down_proj`, `experts` |
| Dtype | F32 |

LoRA is applied to the attention projections **and the MoE expert MLPs**, which is
why the adapter is large (~4.6 GB) despite being rank-16.

## Training

Trained on [Prime Intellect](https://app.primeintellect.ai) Hosted Training —
**[shared run dashboard](https://app.primeintellect.ai/training/shared/buoc18kj8s9odmkdorp0kto6)**:

- **Algorithm:** GRPO
- **Reward:** binary final-answer correctness (1.0 correct / 0.0 wrong) — using
  tools is the *means*, never the reward
- **Learning rate:** 1e-5
- **Rollouts per example:** 16
- **Batch size:** 128
- **Max tokens:** 4096, thinking enabled
- Stopped early once held-out eval reached 100% (≈30 steps).

## What training changed

RL taught the model to run the **measure-then-commit** loop reliably: gather the
tool evidence a question needs, then commit a final answer. Held-out eval rose to
a perfect score and training stayed stable throughout (no collapse).

| | Base model (step 0) | Trained (step 30) |
|---|---|---|
| Held-out eval `pass@1` (n=20) | 0.95 (19/20) | **1.00 (20/20)** |
| Sampled training rollouts | a Veber question lost to "no answer submitted" | **every** sampled rollout submitted a correct answer (reward mean 0.99) |

**A concrete win.** The one question the base model got wrong in the first batch
(step 0, before any gradient update) was this Veber-filter problem:

> *Among C1–C4, find the single candidate passing Veber (`veber_pass == 1`) with
> `tpsa ≤ 120` and `rotatable_bonds ≤ 6`.*
>
> - C1 `O=C(O)CC(=O)NCC(=O)NCC(=O)O`
> - C2 `CCCCCCCCCCN`
> - **C3 `COc1ccc(CCN(C)C)cc1`  ← correct answer**
> - C4 `NCCCCCCCCCCCCNCC(=O)O`

The base model made 12 tool calls and even measured C3 — seeing it clearly passes
(TPSA 12.5, 4 rotatable bonds) — but **never called `submit_answer`**, ran out of
turns, and scored 0. The trained adapter handles exactly this case: it measures
the candidates and commits the answer the evidence supports.

## Choosing the training recipe

The stable hyperparameters above didn't come for free — they're the output of a
sweep on a precursor environment, `allan/science-gym-bio`. The lesson: **learning
rate is the stability lever** (5e-5 peaks then collapses; 1e-5 holds), and larger
rollout groups (8 → 16) cut GRPO advantage variance. That recipe — **LR 1e-5,
16 rollouts/example, thinking on** — is what we carried into the protein-ligand run.

![science-gym-bio hyperparameter sweep](science_gym_bio_sweep.png)

## Usage

```python
import torch
from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer

base = "poolside/Laguna-XS.2"
tok = AutoTokenizer.from_pretrained(base)
model = AutoModelForCausalLM.from_pretrained(base, torch_dtype=torch.bfloat16, device_map="auto")
model = PeftModel.from_pretrained(model, "poolside-laguna-hackathon/protein-ligand-design")
```

For the full tool-use evaluation loop, install and run the gym:

```bash
prime env install jdthewlis/protein-ligand-design
prime eval run jdthewlis/protein-ligand-design -m <your-deployment> -n 20 -r 3
```

---

*Built by **Team JAMMY** for the poolside Laguna hackathon. Trained with GRPO on
Prime Intellect Hosted Training; environment questions generated with Claude Opus 4.8.*