Views
No views yet
transformers library:1from transformers import AutoModelForCausalLM, AutoConfig, AutoTokenizer
2import torch
3
4model_id = "CrystalReasoner/Qwen2.5-3B-CrysReas-ElasticProperties"
5
6tokenizer = AutoTokenizer.from_pretrained(model_id)
7config = AutoConfig.from_pretrained(model_id, trust_remote_code=True)
8model = AutoModelForCausalLM.from_pretrained(
9 model_id,
10 config=config,
11 torch_dtype=torch.bfloat16,
12 trust_remote_code=True
13)
14
15messages = [
16 {"role": "user", "content": "Below is a description of a bulk material. The chemical formula is NaCl. The bulk_modulus is about 100 GPa. Generate a description of the lengths and angles of the lattice vectors and then the element type and coordinates for each atom within the lattice:"},
17]
18
19text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
20model_inputs = tokenizer(text, return_tensors="pt").to(model.device)
21
22generated_ids = model.generate(
23 model_inputs.input_ids,
24 max_new_tokens=2048,
25 pad_token_id=tokenizer.pad_token_id,
26 eos_token_id=tokenizer.eos_token_id,
27 use_cache=True,
28)
29generated_text = tokenizer.batch_decode(generated_ids, skip_special_tokens=False)[0]
30print(generated_text)1def get_structure(generated_text: str):
2 import re
3 from pymatgen.core import Lattice, Structure
4
5 cif_match = re.search(r'<CIF>(.*?)</CIF>', generated_text, re.DOTALL)
6 if cif_match:
7 generated_text = cif_match.group(1)
8
9 lines = [line.strip() for line in generated_text.strip().split('\n') if line.strip()]
10 if lines and not re.match(r'^[-+0-9.eE\s]+$', lines[0]):
11 lines = lines[1:]
12
13 lengths = list(map(float, lines[0].split()))
14 angles = list(map(float, lines[1].split()))
15 lattice = Lattice.from_parameters(*lengths, *angles)
16
17 species = []
18 coords = []
19 for line in lines[2:]:
20 parts = line.split()
21 species.append(parts[0])
22 coords.append([float(parts[2]), float(parts[3]), float(parts[4])])
23
24 structure = Structure(lattice, species, coords)
25 return structure
26
27structure = get_structure(generated_text)
28print(structure)