Files
Llama-3.1-8B-GeomLlama-xyz/README.md

120 lines
4.0 KiB
Markdown
Raw Normal View History

---
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
- xyz
- geom
- llama
- axolotl
- arxiv:2607.13350
datasets:
- GEOM
---
# Llama-3.1-8B-GeomLlama-xyz
**Built with Llama. Built with Axolotl.**
GeomLlama-xyz 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 **Cartesian XYZ coordinates** (one `element x y z` line per atom). It
is one of two models from our paper; the companion model,
[Llama-3.1-8B-GeomLlama-zmatrix](https://huggingface.co/THGLab/Llama-3.1-8B-GeomLlama-zmatrix),
emits Fenske–Hall Z-matrix internal 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-xyz"
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 xyz 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
One atom per line, `element x y z` in Ångström:
```
C -4.344237 -2.044144 -0.978303
C -3.501234 ...
...
```
Parse directly as an XYZ block (no header row is emitted).
> 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, Cartesian XYZ targets (`ori_xyz`), 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.