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, 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:
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.
Training
- Base: meta-llama/Llama-3.1-8B-Instruct
- Framework: 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 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. Please cite the paper and the underlying GEOM dataset (Axelrod & Gómez-Bombarelli, Scientific Data, 2022) if you use this model.
@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. By using this model you agree to its terms. Built with Llama.