---
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/j9g0b2cy8zgusorqktdv0u3l)**:

- **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
- **Steps:** 80 (full run); the final checkpoint is the (tied-)best on eval.

## What training changed

On a **held-out set of 60 questions**, eval accuracy (`avg@1`) rose from the base
model's **91.7%** to **96.7%** — a real **+5-point** gain, with the final
checkpoint tied for best. Training was stable throughout: reward held around 0.9
(normal GRPO variance, dipping to ~0.65 and recovering), completions stayed ~900
tokens (no ballooning), and there were no truncated or failed rollouts.

![held-out eval curve](eval_curve.png)

**A concrete example.** One question the base model got wrong before training:

> *For a kinase backbone-carbonyl halogen bond, is 4-bromo-7-azaindole
> (`Brc1ccc2[nH]ccc2n1`) a candidate — it must carry a heavy halogen (Cl/Br/I)
> for the halogen bond, an aromatic ring for π-stacking, and LogP between 1 and 3?*
> **(correct answer: yes)**

The base model ran the tools (substructure matches for the halogen and ring,
plus descriptors) but **committed the wrong verdict**. This is the kind of
tool-grounded judgement the adapter sharpens.

Honest framing: the base model is already strong on this set (~92%), so the
available headroom was modest — the adapter captured most of it.

## 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.*