ChemSage

A hybrid, chemically aware LLM for drug discovery: RAG for facts, a QLoRA-tuned model for behaviour, and live RDKit/PyMOL tool calls for chemical truth. Fine-tuned and served locally on Apple Silicon with MLX-LM.

ChemSage is a locally hosted, open-source assistant that reasons about small molecules (SMILES, scaffolds, properties), drives the cheminformatics and structural tools a medicinal chemist actually uses (RDKit, PyMOL, docking, PLIP), and grounds every factual claim in your own corpus.

Why it matters: off-the-shelf models invent chemistry (wrong logP, malformed SMILES) and cannot see your private SAR data. ChemSage fixes both by separating concerns: volatile knowledge lives in retrieval, stable behaviour lives in the weights, and deterministic chemical facts are computed live by trusted tools rather than guessed. It is useful for: interrogating your own assay and SAR data, triaging compounds, generating correct PyMOL/RDKit scripts, and interpreting docking and interaction output, all on local hardware. A live, no-install terminal session against the real fine-tuned model runs at chemsage.mdeller.com (see Web deployment below).

Initial idea: Fed up with token, model and upload document limits? Then why not run your own AI models on Mac hardware? Thanks to @ai.christianson on TikTok for highlighting the following method to implement MLX-LM optimised fine-tuning of a local LLM on Mac hardware: https://www.tiktok.com/@ai.christianson/video/7538150476467014943?lang=en

Disclaimer: I am not a professional coder! This is an exercise in Vibe Coding as a Structural Biologist to get the job done. Use at your own risk, and double check the science. I am still working on validation of the Chemistry and Biology. This is a proof of concept exercise to demonstrate what can be achieved using moderate hardware and limited coding experience.

Project links:

• GitHub
• My blog on the making of ChemSage
• Model data and evaluation stats
• Models on Hugging Face

🤗 Models on HuggingFace Hub

⚙️ The hybrid in one line

RAG keeps it truthful today; the fine-tune makes it consistent tomorrow; RDKit makes it correct always.

🚀 Quick start

See PROJECT_PLAN.md for the full step-by-step build (phases 0 to 8). Short version:

▶️ How to run

Terminal 1 — start the model server:

cd /Users/dellboy/Documents/Vibe_Coding/chem_sage
source .venv/bin/activate
mlx_lm.server --model models/chem_sage_32b_v5 --port 8080

Terminal 2 — run chat:

cd /Users/dellboy/Documents/Vibe_Coding/chem_sage
source .venv/bin/activate
python scripts/chat.py --model models/chem_sage_32b_v5

Wait for “HTTP server listening” in Terminal 1 before starting Terminal 2.

🎓 Training

🎓 Training data

Behaviour classes by round:

Round 3 additions (targeting R2 eval gaps):

Round 4 additions (structural biology + fidelity):

🎓 Training parameters

Notes: Round 3 was configured for 1,500 iters but stopped early at actual iter 625 after val loss plateaued at ~0.060. Training was resumed from iter 100 checkpoint after a battery shutdown. 64 layers was attempted first but caused swap thrashing at 31.5 GB peak; reverted to 32 layers. batch_size=8 was also attempted but ran 4.5x slower due to memory pressure; reverted to 4. Round 4 uses 32 layers (same as R3) — 48 and 64 layers were found to be too slow (~2–2.5 min/iter) on M1 Max; 32 layers is the throughput sweet spot at ~40 sec/iter. Round 5 doubles the LoRA rank to 64 with RSLoRA scale 90.0 (effective alpha ≈ 11.25, same order as R4’s 11.31), triples the iters to 3,000, extends max_seq_length to 3,072, and doubles val_batches to 50 for a more stable eval signal. Trainable parameters doubled to 268M (0.82%).

🎓 Training loss curves

Validation loss measured every 25 steps over 25 batches. Round 3 was resumed from iter 100 checkpoint after a battery shutdown; actual iter = logged iter + 100 for the resumed portion.

Round 3 val loss (0.054) is 6.4x better than Round 2 (0.347) and 7.2x better than Round 1 (0.389). The improvement is driven by dataset quality (19 behaviour classes with code-heavy examples) and 3.3x more training data (4,000 vs 1,200 examples). Round 3 was stopped early at iter 625 of 1,500 after val loss plateaued.

Round 4 completed 2026-06-27 (1,500 iters, ~43 h). Best val loss: 0.041 at iter 950, beating R3’s best of 0.054 by 24%. The best was also tied at iter 1,425. Fused from iter 950 checkpoint into models/chem_sage_32b_v4/ (~17.2 GB). Peak memory 25.5 GB, 32 adapted layers, 134M trainable parameters (0.41%).

Round 5 early-stopped 2026-06-29 at iter 2000 of 3,000. Val loss plateaued at 0.055 from iter 1600 (true best 0.054 at iters 1650 and 1950, but fell between checkpoint saves). Fused from iter 1600 checkpoint into models/chem_sage_32b_v5/ (~17.2 GB). Higher floor than R4 (0.055 vs 0.041) reflects the harder task: 78 behaviour classes vs 59, 20k examples vs 6.4k.

🧪 Evaluation results

Original scorecard eval (per-round, eval/eval_chem.py)

Scorecard files: eval/scorecards/scorecard_r1_7b.html, scorecard_r2_32b.html, scorecard_r3_32b_v2.html.

Comparative eval — all 5 rounds (2026-06-29, eval/compare/eval_compare.py)

100 shared R5 test examples (seed=42). Full report: eval/compare/results/compare_20260629_1928.html.

Scores: per-example pass/fail (examples where all instances correct / 100). Metrics degrade gracefully to N/A when a metric’s trigger type never appears in a model’s responses.

R5 wins on all exec/fidelity targets: exec +24%, PyExec +30%, fidelity +42%, Code→Quote +42%, SMARTS +45%, Degeneration-free 100%, Refusal 100%.

The R4 exec regression vs R3 (71% vs 79%) was caused by dataset distribution shift towards structural biology content that reduced the density of pure-RDKit drills. R5 corrected this with pyexec_drill (×6), code_then_quote_v2 (×5), and fidelity_multistep (×5) generators.

📚 RAG corpus

The data/corpus/ directory contains the full retrieval knowledge base: tool reference docs plus two databases of experimental chemistry data:

📚 ChEMBL (data/corpus/chembl/, ~70 MB)

📚 PubChem (data/corpus/pubchem/, ~6.6 MB)

📚 PDB (data/corpus/pdb/, ~111 MB)

📝 Recent changes

Live web deployment (2026-07-24)

• chemsage.mdeller.com is live. The real scripts/chat.py CLI now runs in the browser via an xterm.js terminal, bridged over WebSocket to a server-side PTY — same Rich output, corpus tables, and slash commands as the local CLI, not a rebuilt chat-bubble UI.
• No local GPU on the droplet: mlx_lm‘s stream_generate/generate are patched at import time (chat_remote.py) to call a HuggingFace Space running the real fused 32B model on ZeroGPU, then the unmodified chat.py is run unaware of the substitution.
• Serving stack: gunicorn (gevent worker, 1 worker / 100 connections — a PTY session is I/O-bound, not CPU-bound) behind nginx with WebSocket-upgrade support on /ws, systemd service, Let’s Encrypt TLS. Full detail in Web deployment below.

Chat CLI and dependency fixes (2026-06-30)

• mlx-lm 0.31.3 API fix (scripts/chat.py): stream_generate() and generate() no longer accept verbose=False as a keyword argument — removed from both call sites.
• Corpus fast-path (rag/corpus_lookup.py): Added "get" to the lookup trigger words so queries like "get all EGFR structures" route directly to the corpus table view instead of falling through to the LLM.
• System prompt (scripts/chat.py): Added explicit instruction to render retrieved structured data (PDB structures, bioactivity measurements) as markdown tables rather than generating external API code.
• Dependencies (requirements.txt): Added torch explicitly (required by sentence-transformers for local embeddings).
• How to run: Added quick-start run instructions to the main README — single terminal, venv python, --no-rag flag for model-only mode.
• All 5 models published to HuggingFace Hub (Dellboy/chem_sage_*).

🧪 Eval harness fixes and 5-round comparative eval (2026-06-29)

Tokenizer fix (all models):
– tokenizer_config.json in chem_sage_7b_v1, chem_sage_32b_v2, chem_sage_32b_v3, chem_sage_32b_v4 had extra_special_tokens as a JSON list of strings. Newer transformers calls .keys() on this field — causing an AttributeError on load. The server silently swallowed the error, loaded a broken tokenizer, and then hung indefinitely on every inference request (all queries timed out). Fixed by converting to an empty dict {} in all four files. chem_sage_32b_v5 was already in the correct format.

Eval harness model-path fix (eval/compare/eval_compare.py):
– query_model() previously sent "model": "chemsage" in every HTTP POST to the local mlx_lm.server. Since mlx_lm 0.31.3, an unknown model name triggers a 120-second HuggingFace lazy-load timeout, causing all R1–R4 evals to time out. Fixed by passing the full local model directory path (model_path) as the "model" field.

Per-example pass/fail scoring:
– Changed the eval metric methodology throughout eval_compare.py from per-instance accuracy (ok / total_attempts) to per-example pass/fail (passed_examples / 100). An example “passes” a metric if every instance in that response is correct (ok == tot, tot > 0). This makes round-over-round comparisons meaningful: R1 calling 2 SMILES correctly is no longer the same score as R3 calling 100 SMILES correctly. The _per_example_agg() helper re-derives this from results[].scores in raw_results.json.

5-round comparative eval (complete, 2026-06-29):
– All 5 rounds evaluated on 100 shared R5 test examples (seed=42). Full report: eval/compare/results/compare_20260629_1928.html. R5 overall 95% (vs R4 80%, R3 87%). R5 wins on exec (+24%), PyExec (+30%), fidelity (+42%), SMARTS (+45%). All README eval tables updated.

🎓 Round 5: Higher rank + 2.5× corpus (2026-06-28, training in progress)

• Expanded SFT dataset from 8,000 to 20,000 examples (16,000 train / 2,000 val / 2,000 test) across 78 behaviour classes (19 new vs R4)
• Raised LoRA rank to 64 (from 32); scale set to 90.0 (RSLoRA effective alpha ≈ 11.25, same order as R4); trainable parameters doubled to 268M (0.82%)
• Extended max_seq_length to 3,072 (from 2,048); all corpus examples fit (p99=488, max=558 tokens)
• Tripled training iters to 3,000; doubled val_batches to 50 for a more stable eval signal
• New generators targeting R4 exec/PyExec gap and numerical fidelity: pyexec_drill (×6), code_then_quote_v2 (×5), rounding_explicit (×5), fidelity_multistep (×5)
• New content generators: herg_liability, selectivity_profile, prodrug_bcs, sar_delta, mdanalysis, dssp, ppi_interface, uniprot_api, conformer_3d, mcs_search, reaction_smarts, recap_fragmentation, electron_density, biological_assembly, drug_target_family (×3-4 each)
• 707 examples rejected during validation (96.5% pass rate); token lengths all under max_seq_length (p99=488)
• Training launched 2026-06-28 with caffeinate -i under single-process MLX launch (/tmp/r5_launch.py) — memory limits set in-process before importing mlx_lm; cache 16.7 GB / memory limit 50.1 GB
• Iter 1 val loss: 2.210 (lower than R4’s 2.397 — rank=64 adapter has more capacity from the start)
• Config: config/train_config_r5.yaml; adapter path: adapters/chem_sage_32b_v5_lora/
• Early stop at iter 2000 — val loss plateaued at 0.055 from iter 1600; fused from iter 1600 checkpoint into models/chem_sage_32b_v5/ (~17.2 GB)

🎓 Round 4: Expanded dataset training (2026-06-25 to 2026-06-27, complete)

• Expanded SFT dataset from 5,000 to 8,000 examples (6,400 train / 800 val / 800 test) across 59 behaviour classes
• New behaviour classes: drug mechanisms/indications/ATC, multi-step reasoning, refusals, personality, structural biology, docking, visualisation (Plotly, R/ggplot2), numerical fidelity drills (code_then_quote, rounding_precision)
• steps_per_report aligned with steps_per_eval (both 25) so train/val loss appear at the same iters
• Training ran to completion: 1,500 iters, ~43 h (2026-06-25 15:39 to 2026-06-27 10:59 BST)
• Best val loss: 0.041 at iter 950 — beats R3’s all-time best of 0.054 by 24% (tied at iter 1425)
• Peak memory 25.5 GB, 32 adapted layers, 134M trainable parameters (0.41%)
• Fused from iter 950 checkpoint into models/chem_sage_32b_v4/ (~17.2 GB, 4 safetensors shards)

🧪 Comparative eval harness (2026-06-26)

New script eval/compare/eval_compare.py runs all 4 model rounds side-by-side on the same 100 R4 test examples and produces a rich HTML + Markdown comparison report.

Features:
– Option A auto-server: spawns and kills mlx_lm.server per model automatically (no manual terminal juggling); health-checks the OpenAI-compatible endpoint before starting eval; sends SIGTERM to the process group on completion
– Shared examples: 100 examples sampled from the R4 test set with a fixed seed (reproducible)
– All 13 metrics: 3 core + 7 new metrics, all reporting n/a gracefully on older models
– R3 as baseline: metric table includes a Δ column (±%) vs Round 3; colour-coded green/red
– Model profile panel: full spec comparison across architecture, training config, training outcomes, corpus, and eval runtime
– Runtime resource stats: server process CPU mean/peak %, RSS mean/peak GB (via psutil), mean latency (s), throughput (tok/s)
– Validation loss curves: SVG comparison chart of all 4 models’ val loss curves (clipped at 0.60; dashed line for estimated curves)
– Per-example table: 100 collapsible rows; click “responses” to see all model outputs side-by-side
– Resume flag: --resume loads raw_results.json cache and skips already-evaluated models
– Branding: ChemSage ASCII art (pyfiglet), clickable marcdeller.com links throughout

Usage (run after R4 training completes and fuse is done):

cd chem_sage/
# Full run — all 4 models
python eval/compare/eval_compare.py

# Subset and resume
python eval/compare/eval_compare.py --models round3 round4 --resume

# Quick smoke-test (10 examples)
python eval/compare/eval_compare.py --limit 10

Config: eval/compare/models.yaml — all model metadata, loss curves, and training stats in one place; edit to update R4 final values after training.

🔗 ChromaDB 1.x compatibility fix (2026-06-25)

• ChromaDB 1.0+ switched to a Rust connection backend; PersistentClient timed out on databases created with the old Python backend
• Fix: pass Settings(allow_reset=True, anonymized_telemetry=False) to PersistentClient in both rag/retrieve.py and scripts/ingest_rag.py
• No data loss, no re-ingest — all 55,742 indexed chunks preserved; confirmed connected and serving queries
• Downgrade to chromadb<1.0.0 is not possible on Python 3.14 (no pre-built chroma-hnswlib wheel)

📊 Chat CLI: Streaming output and UX polish (2026-06-25)

Streaming generation:
– Response tokens now stream live to the terminal in real-time using mlx_lm.stream_generate + Rich Live; raw text streams token-by-token during generation, then the final response re-renders as formatted Markdown — no more waiting behind a “Thinking…” spinner
– Falls back to blocking generate + spinner if stream_generate is unavailable

New slash commands:
– /save — exports the full conversation to a timestamped Markdown file on the Desktop (includes marcdeller.com attribution header)
– /info — shows a panel with loaded model name/size, training config (iters, RSLoRA rank, lr, ctx), corpus stats (chunks, CSV files, MB), and embed model; clickable hyperlinks to marcdeller.com and marc@marcdeller.com
– /retry — re-generates the previous response using the same RAG context; replaces the last assistant turn in history

UX refinements:
– Tab completion: type / and press Tab for a dropdown of all slash commands; completion menu styled to match the ChemSage colour scheme
– RAG retrieval indicator: 🔍 Retrieving context… spinner during retrieval, then 📚 N chunks retrieved confirmation line before generation begins
– Context window indicator in toolbar: bottom toolbar now shows ctx N/3 alongside the turn counter so you can see how much of the 3-turn context window is filled
– /reset confirmation: asks Clear history? [y/N] before wiping — no accidental data loss
– Generation stats: now includes elapsed seconds — ⚡ N tok · X.X tok/s · X.Xs
– Styled goodbye: exit via quit/q/Ctrl-C shows a full-screen Panel with the ChemSage ASCII banner, “Thank you for using ChemSage”, and clickable links to marcdeller.com and marc@marcdeller.com
– Hyperlinks throughout: banner, /help, /info, /save output, and goodbye Panel all include clickable [link=...] Rich markup pointing to marcdeller.com and marc@marcdeller.com

Earlier polish (same date):
– prompt_toolkit input with arrow-key history and styled You: prompt
– Persistent bottom toolbar: ⬡ ChemSage · model · QLoRA · RAG · turn N · ctx N/3
– Corpus tables: Rich Table with rounded cyan borders, single-line rows (ellipsis overflow), abbreviated headers, ✓/✗ booleans, method abbreviations
– RDKit output in yellow-bordered Panel; retrieval errors in red-bordered Panel
– Terminal title set to ChemSage on launch and /clear
– Added rich, prompt_toolkit, pyfiglet to requirements.txt

🧪 Eval harness: R4 extended metrics (2026-06-25)

Extended eval/eval_chem.py with new metrics that degrade gracefully to n/a when run against R1-R3 models/datasets, enabling a clean side-by-side comparison across all rounds:

Backwards compatibility: R1-R3 test sets have no refusal examples (refusal shows n/a), no QED mentions (QED shows n/a), code blocks may not print (code-then-quote shows n/a). All five new metrics return (0, 0) when inapplicable. The core 3-metric overall score formula is unchanged for direct cross-round comparison.

New scorecard features: per-class breakdown table (if test examples carry a class field), overall score card, collapsible response viewer in per-example table, extended exec breakdown by library category, mean latency in header, cleaner terminal summary with n/a annotation.

🎓 Round 3: Code-heavy dataset and extended training (2026-06-25)

• Expanded SFT dataset from 1,500 to 5,000 examples (4,000 train) across 19 behaviour classes
• 8 new generator classes targeting eval gaps: single_property, full_profile, qed_score, molecular_formula, smiles_validation, property_comparison, corpus_drug, murcko_sar
• 51 seed SMILES + ~1,000 corpus drug examples from approved_drugs.csv
• Configured for 1,500 iters; stopped early at iter 625 after val loss plateaued at ~0.060
• Best val loss 0.054 at iter 625 (vs 0.347 for Round 2): 6.4x improvement
• Attempted 64 layers (swap thrashing at 31.5 GB) and batch_size=8 (4.5x slower); reverted to 32 layers / batch 4
• Resumed from iter 100 checkpoint after battery shutdown; no data loss
• Eval: SMILES 99%, Exec 99%, Fidelity 68% (vs R2: 100%, 41%, 10%)
• Fused model saved to models/chem_sage_32b_v3/ (17 GB)

Round 2: Qwen 2.5 32B upgrade (2026-06-24)

• Upgraded base model from Qwen 2.5 7B to 32B (4-bit, MLX); trained with 32 adapted layers (vs 16 for 7B)
• Best val loss 0.347 at step 675 (vs 0.389 for 7B): 10.8% improvement
• Fused model saved to models/chem_sage_32b_v2/ (17 GB); serves on 64 GB Apple Silicon at 24.9 GB peak memory
• Eval shows 100% SMILES validity, 41% tool executability, 10% numerical fidelity (7B scored 0% across the board)
• Added tqdm progress bar and live metrics to eval/eval_chem.py
• System prompt injection in eval harness (test examples lack system messages)
• Fixed merge_export.py: uses sys.executable -m mlx_lm fuse (corrected deprecated syntax), resolves adapter path relative to config directory

📚 Round 1: Corpus fast-path and chat CLI

• All 17 corpus CSVs searchable directly from the chat CLI via keyword fast-path (rag/corpus_lookup.py) that bypasses the LLM for enumeration queries
• Up to 500 matching rows loaded per source; pageable in 50-row pages with CSV export
• PDB binding affinity cross-reference: resolves target names to PDB IDs via pdb_all_entries.csv
• Deduplication by comp_id sorted by comp_mw descending (highest-MW representative per compound)
• Intent filtering: SMILES/affinity/sequence queries drop irrelevant sources from the menu
• _CAPS_STOPWORDS prevents format words (“SMILES”) from winning keyword extraction over target names (“EGFR”)
• RAG hybrid loop: top-5 corpus chunks per query via ChromaDB (BAAI/bge-base-en-v1.5, 55k chunks)
• RDKit auto-execution of code blocks in model responses
• ASCII art banner (ChemSage, small_slant font), Rich Markdown rendering, Rich Table corpus tables with rounded cyan borders
• Column display config (_COL_CFG) maps all corpus column names to abbreviated headers and per-column max widths; ellipsis overflow keeps all rows single-line

🌐 Web deployment

Live at chemsage.mdeller.com. Not a rebuilt chat-bubble interface: the browser gets a real terminal (xterm.js) connected over WebSocket to a server-side PTY running the actual, unmodified scripts/chat.py — the same Rich-rendered markdown, corpus tables, and slash commands (/save, /info, /retry, tab completion) as running it locally.

No GPU on the droplet. chat_remote.py is a drop-in launcher: at import time it patches mlx_lm.generate/mlx_lm.stream_generate to call a HuggingFace Space running the real fused 32B model on ZeroGPU (a serverless, on-demand GPU allocation), then runs chat.py completely unaware anything is different. chat.py itself is never modified for the hosted path.

Serving stack: gunicorn with a single gevent worker (100 concurrent connections) rather than sync workers — a terminal session is I/O-bound (waiting on the remote model and on keystrokes), not CPU-bound, so one greenlet-based worker serves many concurrent sessions cheaply. nginx reverse-proxies both the WebSocket endpoint (/ws, with Upgrade/Connection headers and a long read/send timeout for idle-but-open sessions) and the static app shell, behind a Let’s Encrypt certificate. Runs as a dedicated systemd service on the droplet alongside AlphaFraud and BoltzMaker’s own web apps.

Source isolation: the web app (web/flask_app/) lives on a separate web git branch/worktree, so it never touches the tested training/eval code on main.

🧱 Stack

Everything runs locally on Apple Silicon via MLX-LM (mlx_lm.lora to train, mlx_lm.fuse to merge, mlx_lm.server for an OpenAI-compatible endpoint the GUI talks to). Unified memory means no VRAM wall: 16 GB handles ~8B, 32 GB ~14B, 64 GB ~32B. QLoRA is automatic because the base is a 4-bit MLX model.

The chat CLI (scripts/chat.py) uses Rich for streaming Markdown output (Live), corpus tables (rounded borders, ellipsis overflow, abbreviated headers), panels, progress spinners, and a themed console; prompt_toolkit for the interactive input session (arrow-key history, slash-command tab completion, styled prompt, persistent bottom toolbar with context window indicator); ChromaDB + sentence-transformers for the RAG retrieval layer; and RDKit for live chemical property computation. Responses stream token-by-token via mlx_lm.stream_generate and re-render as Markdown on completion. New commands: /save (export conversation to Markdown), /info (session/model details), /retry (regenerate last response).

✅ To Do

Completed (Rounds 1-5)

• ✅ Increase SFT dataset from 1,500 to 5,000 examples with code-heavy generators (Round 3)
• ✅ Add examples that teach “emit RDKit code, don’t compute manually” (Round 3: single_property, full_profile)
• ✅ Include SMILES validation examples (Round 3: smiles_validation)
• ✅ Increase iters from 750 to 1,500 (Round 3; stopped early at 625 — val loss plateaued)
• ✅ Try num_layers: 64 — caused swap thrashing at 31.5 GB; reverted to 32
• ✅ Try batch_size: 8 — 4.5x slower due to memory pressure; reverted to 4
• ✅ Add structural biology generators: PyMOL selections/viz, PLIP, PDB format, pdb-tools, Biopython, ChimeraX, gemmi
• ✅ Add Plotly and R/ggplot2/bio3d recipe generators
• ✅ Add protein family generators: kinases, GPCRs, antibodies, GTPases, cytokines, receptors (real PDB data)
• ✅ Add numerical fidelity drills: code_then_quote (3x), rounding_precision (2x)
• ✅ Add refusal/out-of-scope examples
• ✅ Expand dataset to 8,000 examples across 51 behaviour classes

Round 5 dataset (complete)

• ✅ 20,000 examples across 78 behaviour classes (build_dataset.py --n 20000 --seed 51)
• ✅ 19 new generator classes targeting R4 exec/PyExec gap: pyexec_drill (×6), code_then_quote_v2 (×5), rounding_explicit (×5), fidelity_multistep (×5), herg_liability, selectivity_profile, prodrug_bcs, sar_delta, mdanalysis, dssp, ppi_interface, uniprot_api, conformer_3d, mcs_search, reaction_smarts, recap_fragmentation, electron_density, biological_assembly, drug_target_family
• ✅ Token lengths p50=322, p99=488, max=558 — all under max_seq_length=3072
• ✅ 707 examples rejected (96.5% pass rate); 78 classes, class balance max:median < 3:1
• ✅ data/sft_r5/ — 16,000 train / 2,000 val / 2,000 test

🎓 Round 5 training (complete — early stop)

• ✅ Config: config/train_config_r5.yaml — rank=64, scale=90, iters=3000, max_seq=3072, val_batches=50
• ✅ Launch: single-process caffeinate -i .venv/bin/python /tmp/r5_launch.py (cache 16.7 GB, memory limit 50.1 GB)
• ✅ Training started 2026-06-28 — Iter 1 val loss 2.210
• ✅ No rank=64 instability — converged normally
• ✅ Early stop at iter 2000 — val loss plateaued at 0.055 from iter 1600 (best: 0.054 at iters 1650 and 1950, unsaved)
• ✅ Best available checkpoint: iter 1600, val loss 0.055 → copied to adapters.safetensors
• ✅ Fused into models/chem_sage_32b_v5/ (~17.2 GB)
• ✅ Run eval/compare/eval_compare.py after fuse

Round 4 dataset (complete)

• ✅ 8,000 examples across 59 behaviour classes (build_dataset.py --n 8000 --seed 47)
• ✅ Protein family coverage: kinases (7,687 PDB), antibodies (3,322), GPCRs (103), GTPases (359), cytokines (769), receptors (6,891) — real PDB IDs and ligand SMILES from corpus
• ✅ Numerical fidelity drills: code_then_quote enforces “compute then quote” pattern; rounding_precision teaches MW 1dp, logP 2dp, TPSA 1dp
• ✅ Structural biology: PyMOL selection algebra, visualization recipes, PLIP interaction analysis, PDB format, pdb-tools, Biopython Bio.PDB, ChimeraX, gemmi, structure preparation, docking interpretation
• ✅ Visualization: Plotly Python (scatter, radar), R/ggplot2 (SAR plots), bio3d (RMSD, RMSF)
• ✅ Set steps_per_report: 25 to align train/val reporting at same iters

🎓 Round 4 training (complete)

• ✅ num_layers: 32 — 48 and 64 layers tested and found too slow (~2 min/iter); 32 is the M1 Max sweet spot
• ✅ Adapter path: adapters/chem_sage_32b_v4_lora (checkpoints saved every 200 iters)
• ✅ Training completed 2026-06-27 10:59 BST — 1,500 iters, ~43 h
• ✅ Best checkpoint: iter 950, val loss 0.041 (also tied at iter 1425)
• ✅ Copied 0000950_adapters.safetensors → adapters.safetensors
• ✅ Fused into models/chem_sage_32b_v4/ (~17.2 GB)
• ✅ Updated eval/compare/models.yaml with final stats and complete val loss curve
• ✅ Raised rank to 64 in Round 5 for more adapter expressivity

🧪 Evaluation and model quality

• ✅ Complete Round 3 eval: SMILES 99%, Exec 99%, Fidelity 68%
• ✅ Comparative eval harness: eval/compare/eval_compare.py with 13 metrics, loss curves, resource stats
• ✅ Comparative eval run after R4 fuse: full 4-round report eval/compare/results/compare_20260627_1753.html
• ✅ Run R5 comparative eval after fuse: add Round 5 to eval/compare/models.yaml and run eval_compare.py
• ✅ Improve exec/PyExec from R4’s 74%/75% to >85% (R5 result: exec 95%, PyExec 99%)
• ⬜ Add perplexity benchmark against base (unfine-tuned) 32B
• ⬜ Add hallucination rate metric: count invented PDB IDs, ChEMBL IDs, SMILES not in corpus

Infrastructure

• ✅ ChromaDB 1.x fix: allow_reset=True in rag/retrieve.py and scripts/ingest_rag.py — 55,742 chunks preserved, no re-ingest required
• ✅ psutil added to requirements.txt for server process monitoring in comparative eval
• ✅ Set up HuggingFace Hub: all 5 models published at huggingface.co/Dellboy
• ⬜ Create scripts/run_eval.sh (start server, wait, run eval, stop server)

Publishing and deployment

• ✅ Live web deployment at chemsage.mdeller.com — shipped as an xterm.js terminal in the browser, not a chat-bubble Gradio UI: the real, unmodified scripts/chat.py runs server-side in a PTY, bridged to the browser over a WebSocket, so the hosted session is identical to running the CLI locally (same slash commands, same Rich rendering). No local GPU on the droplet, so mlx_lm‘s generate/stream_generate are patched at import time to call a HuggingFace Space running the real fused model on ZeroGPU instead — see Web deployment below.
• ⬜ Packaged CLI installer — publish chat.py as a pipx-installable tool (pyproject.toml entry point); users with a local Qwen 32B model can pipx install chemsage and point it at their own mlx_lm.server endpoint.

Round 6 (planned)

Ideas to try next, based on learnings from Rounds 1–5:

• ⬜ Weighted class-balanced sampling across all 78 classes. Distribution shift was the proven cause of the R4 regression (exec 79→71%, SMARTS 100→55%) and SMARTS has swung 100→55→100 across three rounds purely on dataset composition. Lock the fragile metrics by controlling sampling density rather than relying on raw example counts.
• ⬜ Triple the code_then_quote_v2 drill. Code→Quote fidelity moved from 19% (R4) to 61% (R5) but is still the weakest metric by 28 points. The failure mode is copy-fidelity, not reasoning — add examples where the quoted value must byte-match the printed stdout token, including trailing zeros and units, so the model learns to transcribe rather than re-derive.
• ⬜ Add a dedicated pymol_syntax_drill generator. PyMOL syntax has sat at 89–90% for two rounds and is the lowest structural biology metric. Failures cluster in selection algebra and command chaining — target those specifically rather than adding more generic structural biology classes.
• ⬜ Two-phase curriculum: full run then a short fidelity-focused continuation. Train the full 78-class dataset to plateau (~1,600 iters, matching R5 pattern), then run a short ~400-iter second phase at lower LR on a weak-metrics-only subset (code→quote, PyMOL, fidelity). Phase 2 is cheap (~7 h) and concentrates remaining capacity on the actual gaps rather than fighting the plateau.
• ⬜ Configure for ~2,000 iters, not 3,000. R3 stopped at 42% of configured iters, R5 at 67% — both plateaued well before the ceiling. Stop burning 15–20 hours on a flat val curve; redirect that budget into the phase-2 continuation or more dataset quality work.
• ⬜ Try rank 96 (not 128). The rank 32→64 jump in R5 tracked with the 95% overall score but the val floor rose to 0.055 with 78 classes, suggesting the harder task is capacity-limited. Rank 96 (~400M trainable params) should sit comfortably under the R5 peak of 30.3 GB; skip rank 128 to avoid swap risk on 64 GB unified memory.
• ⬜ Extend fidelity_multistep with longer chains and distractor numbers. Fidelity reached 89%, but the residual 11% is multi-step cases where an intermediate value is transcribed incorrectly. Adding 4–5 step chains with adversarial distractor numbers forces the model to track which value to quote rather than defaulting to the largest or last number.
• ⬜ Add per-class eval breakdown and run it on a mid-training checkpoint (~iter 1,000). R4’s exec/SMARTS regression only surfaced after 43 hours of training. A per-class pass/fail on the 100-example shared test set at iter 1,000 would expose class-level dilution early — while there is still time to rebalance the dataset and restart.

👤 Author

Marc C. Deller, D.Phil.
Structural biologist & drug discovery scientist