125 lines
4.5 KiB
Markdown
125 lines
4.5 KiB
Markdown
|
|
---
|
|||
|
|
license: llama3.1
|
|||
|
|
base_model: meta-llama/Llama-3.1-8B-Instruct
|
|||
|
|
library_name: transformers
|
|||
|
|
pipeline_tag: text-generation
|
|||
|
|
language:
|
|||
|
|
- en
|
|||
|
|
tags:
|
|||
|
|
- chemistry
|
|||
|
|
- molecular-geometry
|
|||
|
|
- conformer-generation
|
|||
|
|
- smiles
|
|||
|
|
- z-matrix
|
|||
|
|
- geom
|
|||
|
|
- llama
|
|||
|
|
- axolotl
|
|||
|
|
- arxiv:2607.13350
|
|||
|
|
datasets:
|
|||
|
|
- GEOM
|
|||
|
|
---
|
|||
|
|
|
|||
|
|
# Llama-3.1-8B-GeomLlama-zmatrix
|
|||
|
|
|
|||
|
|
**Built with Llama. Built with Axolotl.**
|
|||
|
|
|
|||
|
|
GeomLlama-zmatrix is a fine-tune of **Llama-3.1-8B-Instruct** that generates
|
|||
|
|
3D molecular conformer geometries directly from a SMILES string, emitting each
|
|||
|
|
structure as a **Fenske–Hall Z-matrix** (internal coordinates: bond length, bond
|
|||
|
|
angle, dihedral, referenced to previously placed atoms). It is one of two models
|
|||
|
|
from our paper; the companion model, [Llama-3.1-8B-GeomLlama-xyz](https://huggingface.co/THGLab/Llama-3.1-8B-GeomLlama-xyz),
|
|||
|
|
emits Cartesian XYZ coordinates instead.
|
|||
|
|
|
|||
|
|
The model was trained jointly ("hybrid") on **GEOM-QM9** and **GEOM-Drugs**, so it
|
|||
|
|
covers both small molecules and larger drug-like molecules with a single set of
|
|||
|
|
weights.
|
|||
|
|
|
|||
|
|
## Quick start
|
|||
|
|
|
|||
|
|
The model was trained in the Alpaca instruction format. Reproduce the exact
|
|||
|
|
inference prompt used for the paper's numbers:
|
|||
|
|
|
|||
|
|
```python
|
|||
|
|
import torch
|
|||
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer
|
|||
|
|
|
|||
|
|
model_id = "THGLab/Llama-3.1-8B-GeomLlama-zmatrix"
|
|||
|
|
tok = AutoTokenizer.from_pretrained(model_id)
|
|||
|
|
model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16, device_map="auto")
|
|||
|
|
|
|||
|
|
smiles = "Cc1cccc(CSc2nnnn2-c2ccccc2)c1"
|
|||
|
|
prompt = (
|
|||
|
|
"### Instruction:\n"
|
|||
|
|
"You can generate accurate molecular coordinates from a prompt "
|
|||
|
|
"containing a SMILES string.\n\n"
|
|||
|
|
"### Input:\n"
|
|||
|
|
"Generate a realistic equilibrium geometry for the molecule with the "
|
|||
|
|
f"following SMILES string in Fenske-Hall Z-matrix format: {smiles}\n\n"
|
|||
|
|
"### Response:\n"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
inputs = tok(prompt, return_tensors="pt").to(model.device)
|
|||
|
|
out = model.generate(**inputs, max_new_tokens=3072, do_sample=True, temperature=1.0, top_p=0.95)
|
|||
|
|
print(tok.decode(out[0][inputs.input_ids.shape[1]:], skip_special_tokens=True))
|
|||
|
|
```
|
|||
|
|
|
|||
|
|
Sample many completions per SMILES (each is one candidate conformer) to build a
|
|||
|
|
conformer ensemble. `T=1.0, top_p=0.95` and `T=1.2, top_p=0.95` are good defaults.
|
|||
|
|
|
|||
|
|
### Output format
|
|||
|
|
|
|||
|
|
Each line places one atom. The first token is the element; the remaining tokens
|
|||
|
|
are the internal coordinates relative to earlier atoms:
|
|||
|
|
|
|||
|
|
```
|
|||
|
|
O 1 # atom 1: reference origin
|
|||
|
|
C 1 1.388 # atom 2: bonded to atom 1 at 1.388 Å
|
|||
|
|
H 2 1.090 1 105.769 # atom 3: bond to 2, angle to 1
|
|||
|
|
C 2 1.547 1 117.304 3 129.0 # atom 4: bond, angle, dihedral
|
|||
|
|
...
|
|||
|
|
```
|
|||
|
|
|
|||
|
|
Atoms are emitted in an arbitrary order with **no atom labels or connectivity
|
|||
|
|
block** — the model learns geometry from the SMILES alone. Parse the Z-matrix back
|
|||
|
|
to Cartesian coordinates with a standard NeRF/internal-to-Cartesian routine.
|
|||
|
|
|
|||
|
|
> View z-matrix or xyz coordinates easily at
|
|||
|
|
> [doublemolview.streamlit.app](https://doublemolview.streamlit.app).
|
|||
|
|
|
|||
|
|
## Training
|
|||
|
|
|
|||
|
|
- **Base:** meta-llama/Llama-3.1-8B-Instruct
|
|||
|
|
- **Framework:** [Axolotl](https://github.com/axolotl-ai-cloud/axolotl)
|
|||
|
|
- **Method:** LoRA (r = 32, α = 16, dropout 0.05, all linear layers), merged into the base weights
|
|||
|
|
- **Epochs:** 4 · **LR:** 3e-4, cosine · **Optimizer:** adamw_bnb_8bit
|
|||
|
|
- **Sequence length:** 4096, sample packing · **Precision:** bf16
|
|||
|
|
- **Data:** GEOM-QM9 + GEOM-Drugs, Fenske–Hall Z-matrix targets (`ori_fh`), plus the
|
|||
|
|
[Alpaca](https://huggingface.co/datasets/tatsu-lab/alpaca) instruction dataset for
|
|||
|
|
general-instruction rehearsal
|
|||
|
|
|
|||
|
|
## Citation
|
|||
|
|
|
|||
|
|
Paper: *How Well Can Frontier Large Language Models Generate Structures? High Quality
|
|||
|
|
Prediction of Molecular Geometries with Help from Fine-Tuning* —
|
|||
|
|
[arXiv:2607.13350](https://arxiv.org/abs/2607.13350). Please cite the paper and the
|
|||
|
|
underlying GEOM dataset (Axelrod & Gómez-Bombarelli, *Scientific Data*, 2022) if you
|
|||
|
|
use this model.
|
|||
|
|
|
|||
|
|
```bibtex
|
|||
|
|
@misc{cavanagh2026geomllama,
|
|||
|
|
title = {How Well Can Frontier Large Language Models Generate Structures?
|
|||
|
|
High Quality Prediction of Molecular Geometries with Help from Fine-Tuning},
|
|||
|
|
author = {Cavanagh, Joseph M. and Arnold, Jonathan B. and Alteri, Giovanni Battista
|
|||
|
|
and Gritsevskiy, Andrew and Head-Gordon, Teresa},
|
|||
|
|
year = {2026},
|
|||
|
|
eprint = {2607.13350},
|
|||
|
|
archivePrefix = {arXiv},
|
|||
|
|
url = {https://arxiv.org/abs/2607.13350}
|
|||
|
|
}
|
|||
|
|
```
|
|||
|
|
|
|||
|
|
## License
|
|||
|
|
|
|||
|
|
Governed by the [Llama 3.1 Community License](https://github.com/meta-llama/llama-models/blob/main/models/llama3_1/LICENSE).
|
|||
|
|
By using this model you agree to its terms. Built with Llama.
|